mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 16:08:43 +00:00
Compare commits
63 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3fd322a93a | |||
| d8f0982ee8 | |||
| 7276c34c51 | |||
| 1918d1a3b6 | |||
| 3b626efa47 | |||
| 137eac9b50 | |||
| 06b53c97d6 | |||
| 711e05619b | |||
| afc0e5f497 | |||
| c12a6dce9f | |||
| 8ea78e3fbc | |||
| 40238d240a | |||
| 60428e1a32 | |||
| 5b982f2f05 | |||
| cde48fad95 | |||
| 1f2068b9fe | |||
| 7527890607 | |||
| a548e59d49 | |||
| 104fc5a08e | |||
| 715be580d0 | |||
| 0d9c87a079 | |||
| 8e364e6812 | |||
| 32a2776446 | |||
| 285add40dd | |||
| 22bf091de1 | |||
| ff81428a9c | |||
| 75c5c83f12 | |||
| 291e9e37be | |||
| 6c066530e5 | |||
| f428c6a76c | |||
| df89c133ca | |||
| ec763521d4 | |||
| f8dc2f78ee | |||
| 3bcff0165e | |||
| c6db80dd0b | |||
| f84190fe12 | |||
| 122dcd0f66 | |||
| e6661a7285 | |||
| 37466a0390 | |||
| bfce8a510d | |||
| a1261e6299 | |||
| 17c499177f | |||
| d889321b5e | |||
| 8a37f2ad77 | |||
| f94673ae5e | |||
| 3b70fc4c9d | |||
| 3a7b02119b | |||
| bcbc0da090 | |||
| 9bead9f53d | |||
| 0351b77984 | |||
| f6c9d31f98 | |||
| a8f1c5a69f | |||
| 10fecdf051 | |||
| c9ae93a7fa | |||
| 05756f0bbf | |||
| 2a0945443e | |||
| 39e819b6a7 | |||
| 70126943ff | |||
| e01777070d | |||
| 3878adc6dc | |||
| 3df3043563 | |||
| 8a5cd74e48 | |||
| 448d5ec20f |
@@ -0,0 +1,145 @@
|
||||
---
|
||||
name: lancedb-branch-ops
|
||||
description: >-
|
||||
Manage LanceDB table branches through the REST API: list, create, and delete
|
||||
branches; target schema reads, field-metadata updates, and index creation to a
|
||||
named branch; and verify that branch changes remain isolated from main. Use
|
||||
when a task involves branch lifecycle, an experimental or isolated table
|
||||
version, directing an operation to a non-main branch, or confirming that a
|
||||
mutation did not affect main. This skill also explains that LanceDB has no
|
||||
checkout operation; each request selects its target branch in the request
|
||||
body.
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
Manage branches on a LanceDB table: list what exists, create new ones, delete stale ones, and direct read/write operations at a specific branch without touching main.
|
||||
|
||||
## Step 0: Establish the connection
|
||||
|
||||
Use the `lancedb-connect` skill to resolve the base URL and auth headers (`x-api-key`, `x-lancedb-database`). Skip this only if the connection is already known from the current conversation.
|
||||
|
||||
All examples below use `{base_url}` — substitute the resolved endpoint and include the auth headers on every request.
|
||||
|
||||
## The branch model (important)
|
||||
|
||||
LanceDB branches are named snapshots that diverge from the table's current state at creation time. There is **no checkout command** — you never switch the whole table to a branch. Instead, you **pass `"branch": "<name>"` in the request body** of any operation to target that branch. Omitting the key (or sending an empty body) always targets main.
|
||||
|
||||
`branches/list` returns only non-main branches. Main always exists and is not listed.
|
||||
|
||||
## List branches
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/list
|
||||
Content-Type: application/json
|
||||
|
||||
{}
|
||||
```
|
||||
|
||||
Response:
|
||||
```json
|
||||
{
|
||||
"branches": {
|
||||
"experiment-reindex": {"parentVersion": 1, "createAt": 1782506085, "manifestSize": 1029}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
If `branches` is `{}`, the table has no branches besides main.
|
||||
|
||||
## Create a branch
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/create
|
||||
Content-Type: application/json
|
||||
|
||||
{"name": "experiment-reindex"}
|
||||
```
|
||||
|
||||
HTTP 200 with `{}` body = success. The branch is created off the table's current state on main.
|
||||
|
||||
Verify by calling `branches/list` and confirming the new name appears.
|
||||
|
||||
## Delete a branch
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/delete
|
||||
Content-Type: application/json
|
||||
|
||||
{"name": "stale-2024"}
|
||||
```
|
||||
|
||||
HTTP 200 with `{}` body = success. Only the branch pointer is removed — main and all row data remain intact.
|
||||
|
||||
Verify by calling `branches/list` (name gone) and `describe` with no branch param (main still responds).
|
||||
|
||||
## Operate on a specific branch
|
||||
|
||||
Pass `"branch": "<name>"` in the body of any operation to scope it to that branch:
|
||||
|
||||
**Read schema on a branch:**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/describe
|
||||
Content-Type: application/json
|
||||
|
||||
{"branch": "wip-branch"}
|
||||
```
|
||||
|
||||
**Write metadata to a branch (not main):**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/update_field_metadata
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"branch": "wip-branch",
|
||||
"updates": [
|
||||
{
|
||||
"path": "category",
|
||||
"metadata": {"lancedb:description": "Product category label."},
|
||||
"replace": false
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Build an index on a branch:**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/create_index
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"branch": "wip-branch",
|
||||
"column": "category",
|
||||
"index_type": "BTREE"
|
||||
}
|
||||
```
|
||||
|
||||
## Verifying isolation
|
||||
|
||||
After writing to a branch, always confirm the change did NOT land on main:
|
||||
|
||||
```bash
|
||||
# Should show the new metadata
|
||||
curl -s -X POST {base_url}/v1/table/{table_id}/describe \
|
||||
-H "x-api-key: <key>" -H "x-lancedb-database: <db>" \
|
||||
-H "content-type: application/json" \
|
||||
-d '{"branch": "wip-branch"}'
|
||||
|
||||
# Should NOT show the new metadata
|
||||
curl -s -X POST {base_url}/v1/table/{table_id}/describe \
|
||||
-H "x-api-key: <key>" -H "x-lancedb-database: <db>" \
|
||||
-H "content-type: application/json" \
|
||||
-d '{}'
|
||||
```
|
||||
|
||||
## Quick reference
|
||||
|
||||
| Goal | Endpoint | Body |
|
||||
|------|----------|------|
|
||||
| List all branches | `branches/list` | `{}` |
|
||||
| Create a branch | `branches/create` | `{"name": "..."}` |
|
||||
| Delete a branch | `branches/delete` | `{"name": "..."}` |
|
||||
| Read schema on branch | `describe` | `{"branch": "..."}` |
|
||||
| Write metadata on branch | `update_field_metadata` | `{"branch": "...", "updates": [...]}` |
|
||||
| Build index on branch | `create_index` | `{"branch": "...", "column": ..., "index_type": ...}` |
|
||||
| Target main (default) | any endpoint | omit `"branch"` key |
|
||||
@@ -0,0 +1,79 @@
|
||||
---
|
||||
name: lancedb
|
||||
description: Use when writing, reviewing, debugging, or documenting LanceDB pipelines in Python or TypeScript, especially code that should work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables. Helps avoid non-portable full-table materialization, choose idiomatic query/search patterns, and apply LanceDB performance defaults for ingestion, indexing, filtering, and diagnostics.
|
||||
---
|
||||
|
||||
# Building LanceDB Pipelines
|
||||
|
||||
Use this skill to produce LanceDB pipelines that are portable between local and remote tables (for LanceDB Enterprise/Cloud) and idiomatic for the selected SDK.
|
||||
|
||||
## LanceDB Table Modes
|
||||
|
||||
LanceDB has two common execution modes:
|
||||
|
||||
- **Local table**: embedded, open source, in-process LanceDB. The client opens data from a local path or object storage URI and executes queries in the application process.
|
||||
- **Remote table**: LanceDB Enterprise/Cloud table opened through a `db://...` URI. The data may be very large, commonly backed by object storage, and queried through a remote service.
|
||||
|
||||
Do NOT assume local-only table helpers exist on remote tables. If the user asks for LanceDB Enterprise, Cloud, `db://...`, production remote access, or a remote table, focus on the remote table path: use `search()` / `query()`, keep reads bounded with `select()` and `limit()`, and avoid table-level full materialization APIs.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Identify the SDK: Python, TypeScript, or both.
|
||||
2. Identify the table mode: local/embedded OSS, remote Enterprise/Cloud, or portable across both. If the user says "LanceDB Enterprise", choose the remote table path.
|
||||
3. Read the matching language branch before writing or changing code:
|
||||
- Python patterns: `references/python/patterns.md`
|
||||
- Python API quick reference: `references/python/api_reference.md`
|
||||
- Python performance guidance: `references/python/performance.md`
|
||||
- TypeScript patterns: `references/typescript/patterns.md`
|
||||
- TypeScript API quick reference: `references/typescript/api_reference.md`
|
||||
- TypeScript performance guidance: `references/typescript/performance.md`
|
||||
4. Start with `patterns.md` for the selected SDK. Read `api_reference.md` when choosing method names or return collectors. Read `performance.md` when the task involves ingestion, indexing, filtering, query tuning, diagnostics, or large datasets.
|
||||
5. For Python schemas, favor Pydantic models and validate records before writing. Use PyArrow schemas when Arrow-native, streaming, or highly dynamic data makes them materially better suited.
|
||||
6. Prefer `search()` or `query()` builders with explicit `select()` and `limit()` for reads.
|
||||
7. Avoid table-level full materialization in remote or portable code. This is the main local-vs-remote read pitfall.
|
||||
8. After a successful embedded OSS ingestion, call `table.optimize()`. Do not call it for Enterprise/Cloud; remote maintenance is automatic.
|
||||
9. For remote Enterprise/Cloud writes, never drop-then-reuse or `mode="overwrite"` the same table name — see "Enterprise: never drop-then-reuse the same table name" below. This is the main local-vs-remote write pitfall.
|
||||
10. If reviewing an existing file or repo, run `scripts/check_materialization.py` on the relevant paths and inspect each finding before editing.
|
||||
11. Cross-check unfamiliar or non-trivial API claims against the source tree instead of relying on memory.
|
||||
|
||||
## Core Portability Rule
|
||||
|
||||
Do not write code that assumes a local table API will exist on a remote table. Remote tables can be very large, so whole-table materialization helpers are intentionally unavailable or unsafe.
|
||||
|
||||
This does **not** mean result conversion is forbidden. Bounded query/search result collection is normal:
|
||||
|
||||
- Python: `table.search(...).select([...]).limit(10).to_pandas()`
|
||||
- TypeScript: `await table.search(...).select([...]).limit(10).toArray()`
|
||||
|
||||
The unsafe pattern is table-level or unbounded collection, plus local-only dataset escape hatches in remote code:
|
||||
|
||||
- Python: `table.to_pandas()`, `table.to_arrow()`, `table.to_polars()`; `table.to_lance()` is local/OSS-only dataset access, not materialization
|
||||
- TypeScript: `await table.toArrow()`, `await table.query().toArray()` without `limit()`
|
||||
|
||||
## Enterprise: never drop-then-reuse the same table name
|
||||
|
||||
LanceDB Enterprise/Cloud splits a **control plane** (DDL: create/drop/rename) from a **data plane** (query nodes that serve reads). Query nodes cache the resolved dataset for a table name for up to `table_cache_ttl` — **default 300 seconds (5 minutes)**. After you drop or overwrite a table, the control plane updates immediately but the data plane keeps serving the *old* dataset until that cache entry expires. During the window the two planes disagree.
|
||||
|
||||
The failure this causes: you `drop_table("t")` then immediately `create_table("t", ...)` (or `create_table("t", ..., mode="overwrite")`). The DDL returns success, but every query against `t` returns **`500 Internal Server Error`** (the query node resolves the stale/deleted dataset), and a fresh `describe` may still show the *old* schema/version. It looks like your write silently failed; it didn't — the name is cached.
|
||||
|
||||
**`mode="overwrite"` has the same problem** — it is a drop+create of the same name under the hood.
|
||||
|
||||
Rules for portable Enterprise ingestion:
|
||||
|
||||
1. **Never reuse a table name you just dropped/overwrote within the cache TTL.** Do not use `mode="overwrite"` to replace an existing Enterprise table in place.
|
||||
2. To (re)load data, **write to a fresh table name** (e.g. `<table>_v2`, or a run-stamped suffix). A brand-new name has no cached data-plane entry, so writes and reads work immediately.
|
||||
3. Before creating, `list_tables()` and **fail loudly if the name already exists** rather than overwriting — prompt for a new name.
|
||||
4. To land on a specific final name that is currently occupied by an old table: drop the old table, **wait out the TTL (~5 min), then `rename_table(fresh_name, final_name)`**. Renaming onto a name whose old dataset is still cached hits the same race, so the wait is mandatory. `rename_table` is a supported control-plane op.
|
||||
5. When you hand a table name back to a human, tell them which step still needs the propagation wait (usually: "the old `t` was dropped; run the rename in ~5 minutes").
|
||||
|
||||
This is Enterprise/Cloud-specific. Local/OSS tables have no separate data plane, so `mode="overwrite"` and immediate same-name reuse are fine there.
|
||||
|
||||
## Script
|
||||
|
||||
Run the scanner when reviewing or modifying an existing codebase:
|
||||
|
||||
```bash
|
||||
python skills/lancedb/scripts/check_materialization.py path/to/file_or_dir
|
||||
```
|
||||
|
||||
The script reports likely unsafe full-table materialization in Python and TypeScript. Treat results as review prompts, not automatic proof of a bug.
|
||||
@@ -0,0 +1,105 @@
|
||||
# Python API Reference
|
||||
|
||||
Quick method reference for Python LanceDB code. Cross-check source for non-trivial claims.
|
||||
|
||||
## Connect
|
||||
|
||||
```python
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect("./camelot-db") # local/OSS
|
||||
db = lancedb.connect("db://my-db", api_key=api_key, region=region) # remote
|
||||
```
|
||||
|
||||
**Place the local database directory next to the script/entrypoint that opens it** (i.e. resolve the path relative to the script, `Path(__file__).parent / "camelot-db"`), not buried under a shared `data/` folder. The Lance dataset is the database, not a data file — keeping it beside its code makes ownership obvious and paths stable regardless of the working directory the script is launched from.
|
||||
|
||||
**Do not name the directory `lancedb`** (e.g. `./lancedb`, `./data/lancedb`). It collides with the imported `lancedb` package name, which is confusing to read and easy to shadow in scripts. Give it a name derived from the repo or dataset with a clear prefix/suffix — for example `./<dataset>-db`, `./<repo>_lancedb`, or `./vectordb`.
|
||||
|
||||
Async:
|
||||
|
||||
```python
|
||||
db = await lancedb.connect_async("./camelot-db")
|
||||
```
|
||||
|
||||
## Table Reads
|
||||
|
||||
| Task | Preferred API |
|
||||
| --- | --- |
|
||||
| Vector search | `table.search(query_vector).limit(k)` |
|
||||
| Full scan with filters/projection (sync) | `table.search().where(...).select(...).limit(...)` |
|
||||
| Full scan with filters/projection (async) | `table.query().where(...).select(...).limit(...)` |
|
||||
| Filter | `.where("col > 10")` |
|
||||
| Projection | `.select(["id", "text"])` |
|
||||
| Bound result count | `.limit(20)` |
|
||||
| Collect bounded result as Python objects (default, no extra deps) | `.to_list()` on query/search result |
|
||||
| Collect bounded result as Arrow (default, `pyarrow` always available) | `.to_arrow()` on query/search result |
|
||||
| Collect bounded result as pandas (only if project uses pandas) | `.to_pandas()` on query/search result |
|
||||
| Collect bounded result as Polars (only if project uses polars) | `.to_polars()` on query/search result |
|
||||
|
||||
## Sync vs Async Scan API
|
||||
|
||||
The plain-scan entry point differs between the sync and async clients. **Verified against `lancedb` 0.34.0** — re-check if the pinned version changes:
|
||||
|
||||
- **Sync** (`lancedb.connect(...)`): the table has **no `.query()` method**. Use `.search()` with no argument for a plain scan; it returns a query builder that supports `.where()`, `.select()`, `.limit()`, and the `.to_list()` / `.to_arrow()` / `.to_pandas()` / `.to_polars()` collectors.
|
||||
```python
|
||||
rows = table.search().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
- **Async** (`lancedb.connect_async(...)`): the table has **both** `.query()` and `.search()`. Use `.query()` for a plain scan.
|
||||
```python
|
||||
rows = await async_table.query().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
|
||||
Do not call `table.query()` on a sync table — it raises `AttributeError`.
|
||||
|
||||
## Local vs Remote Table Methods
|
||||
|
||||
| API | Local table | Remote table | Agent guidance |
|
||||
| --- | --- | --- | --- |
|
||||
| `table.search(...)` | Yes | Yes | Preferred read path (sync + async) |
|
||||
| `table.query()` | Async only | Async only | Sync scan path is `table.search()`; `.query()` is the async scan builder |
|
||||
| `table.to_pandas()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_arrow()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_polars()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_lance()` | Yes | No | Local/OSS escape hatch only |
|
||||
|
||||
## Indexes
|
||||
|
||||
Use `create_index(...)` for vector indexes and modern index configs. Use scalar indexes for filtered or merge keys.
|
||||
|
||||
Common calls:
|
||||
|
||||
```python
|
||||
table.create_index("vector")
|
||||
table.create_scalar_index("status")
|
||||
table.create_fts_index("text")
|
||||
```
|
||||
|
||||
Check source docs before specifying advanced index config names or parameters.
|
||||
|
||||
## Filtering And Recall Knobs
|
||||
|
||||
```python
|
||||
table.search(query_vector).where("status = 'ready'") # pre-filter by default
|
||||
table.search(query_vector).where("status = 'ready'", prefilter=False)
|
||||
table.search(query_vector).limit(10).refine_factor(20)
|
||||
table.search(query_vector).limit(10).nprobes(50)
|
||||
```
|
||||
|
||||
Use post-filtering only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
```python
|
||||
print(table.search(query_vector).where("year > 2000").limit(10).analyze_plan())
|
||||
print(table.index_stats("vector_idx"))
|
||||
```
|
||||
|
||||
Use these before changing indexes or search tuning.
|
||||
|
||||
## Maintenance
|
||||
|
||||
```python
|
||||
table.optimize()
|
||||
```
|
||||
|
||||
Call this after every successful local/OSS ingestion. It handles compaction, cleanup of old versions according to retention, and index optimization. Do not add this for LanceDB Enterprise/Cloud remote tables; Enterprise handles compaction and cleanup automatically from cluster configuration.
|
||||
@@ -0,0 +1,173 @@
|
||||
# Python Patterns
|
||||
|
||||
Use these patterns when writing Python code with `lancedb`.
|
||||
|
||||
## Before Writing Code
|
||||
|
||||
Choose the output type from what the project actually depends on. **Do not assume `pandas` or `polars` is installed** — they are heavy dependencies that many LanceDB projects do not use. `pyarrow`, by contrast, ships as a LanceDB dependency and is always available, so it is a safe default to lean on.
|
||||
|
||||
Default output (after applying `select()` and `limit()`):
|
||||
|
||||
- **Python objects**: `.to_list()` — a list of dicts, no extra dependencies. Prefer this for scripts, examples, and agent-generated code unless there is a reason to do otherwise.
|
||||
- **PyArrow**: `.to_arrow()` — a `pyarrow.Table`, when the surrounding code is Arrow-native or you need columnar/zero-copy handoff.
|
||||
|
||||
Only reach for a DataFrame when the project *already* declares that dependency:
|
||||
|
||||
- Pandas projects (pandas in `pyproject.toml`/requirements): `.to_pandas()`.
|
||||
- Polars projects (polars declared): `.to_polars()`.
|
||||
|
||||
If unsure, check the dependency manifest or the imports in surrounding files. When in doubt, use `.to_list()` or `.to_arrow()`.
|
||||
|
||||
## Schema Design and Validation
|
||||
|
||||
Favor `LanceModel` and Pydantic validation for Python schemas. They keep field
|
||||
types readable, validate source records before a write, and map directly to a
|
||||
LanceDB schema. Use `Vector(dimension)` for fixed-size vectors:
|
||||
|
||||
```python
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
class Document(LanceModel):
|
||||
id: int
|
||||
text: str
|
||||
vector: Vector(384, nullable=False)
|
||||
|
||||
rows = [Document.model_validate(row) for row in source_rows]
|
||||
table = db.create_table("documents", schema=Document)
|
||||
table.add(rows)
|
||||
```
|
||||
|
||||
Use PyArrow schemas instead when the pipeline is already Arrow-native, needs
|
||||
record-batch streaming, or has runtime schema requirements that would make a
|
||||
Pydantic model harder to understand. Declare Pydantic as a direct project
|
||||
dependency when application code imports it, even if LanceDB also depends on it.
|
||||
|
||||
## Recommended Patterns
|
||||
|
||||
### Bounded search or query
|
||||
|
||||
Use this for application reads, examples, notebooks, and agent-generated scripts:
|
||||
|
||||
```python
|
||||
results = (
|
||||
table.search(query_vector)
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.to_list() # or .to_arrow(); .to_pandas()/.to_polars() only if the project uses them
|
||||
)
|
||||
```
|
||||
|
||||
Why: `search()` works across local and remote tables and on both the sync and async clients. `select()` avoids fetching unused columns. `limit()` prevents accidental full-table reads. `.to_list()` and `.to_arrow()` avoid assuming pandas/polars is installed (see "Before Writing Code").
|
||||
|
||||
For a **plain scan** (no query vector), the entry point differs by client:
|
||||
|
||||
```python
|
||||
# Sync client: no .query() method — use .search() with no argument.
|
||||
rows = table.search().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
|
||||
# Async client: use .query().
|
||||
rows = await async_table.query().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
|
||||
`table.query()` on a sync table raises `AttributeError` (verified on `lancedb` 0.34.0). See the "Sync vs Async Scan API" section in `api_reference.md`.
|
||||
|
||||
### Bounded query result conversion
|
||||
|
||||
It is fine to collect bounded query/search results:
|
||||
|
||||
```python
|
||||
arrow_table = table.search().select(["id"]).limit(100).to_arrow() # sync plain scan
|
||||
rows = table.search(query_vector).limit(10).to_list()
|
||||
df = table.search(query_vector).limit(10).to_pandas() # only if pandas is a project dep
|
||||
```
|
||||
|
||||
### Local-only Lance dataset API
|
||||
|
||||
`table.to_lance()` does not itself materialize the full dataset. It returns the underlying `lance.LanceDataset`, making the table accessible through the PyLance dataset API. Use it when the task is explicitly local/OSS and needs Lance dataset methods not exposed by LanceDB:
|
||||
|
||||
```python
|
||||
# Local/OSS only: RemoteTable does not expose table.to_lance().
|
||||
ds = table.to_lance()
|
||||
for batch in ds.to_batches(columns=["id", "text"], batch_size=10_000):
|
||||
process(batch)
|
||||
```
|
||||
|
||||
### Async Python
|
||||
|
||||
Keep the same shape and bound the result before collecting:
|
||||
|
||||
```python
|
||||
results = await (
|
||||
async_table.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.to_list() # or .to_arrow()
|
||||
)
|
||||
```
|
||||
|
||||
## Anti-Patterns
|
||||
|
||||
**Avoid the following anti-patterns in your code.**
|
||||
|
||||
### Table-level full materialization
|
||||
|
||||
Avoid whole-table collectors in portable or large-table code:
|
||||
|
||||
```python
|
||||
df = table.to_pandas()
|
||||
arrow_table = table.to_arrow()
|
||||
polars_df = table.to_polars()
|
||||
```
|
||||
|
||||
Why: local tables expose these whole-table collectors, but remote tables intentionally do not — a remote production table can be far larger than a local development table, so it is easy to accidentally pull the entire table into memory.
|
||||
|
||||
`table.to_lance()` is different: it is not a full materialization call, but it is still local/OSS-only and should not appear in code meant to run against remote Enterprise tables.
|
||||
|
||||
### Unbounded result collection
|
||||
|
||||
Avoid query/search collection without a meaningful limit:
|
||||
|
||||
```python
|
||||
rows = table.search().to_list() # unbounded plain scan
|
||||
rows = table.search(query_vector).to_list() # unbounded vector search
|
||||
```
|
||||
|
||||
Prefer `select(...).limit(...)` before collecting; for large reads, stream in batches instead.
|
||||
|
||||
### Per-row writes
|
||||
|
||||
Avoid loops that write one row per call:
|
||||
|
||||
```python
|
||||
for row in rows:
|
||||
table.add([row]) # one commit + fragment per row
|
||||
```
|
||||
|
||||
Each `add()` creates a new version and fragment. Pass the whole batch in a single call, or chunk very large inputs:
|
||||
|
||||
```python
|
||||
table.add(rows) # single commit
|
||||
# for very large inputs, add batches of several thousand rows
|
||||
```
|
||||
|
||||
After the final successful write to an embedded OSS table, call
|
||||
`table.optimize()`. Skip this for Enterprise/Cloud tables because their
|
||||
maintenance is automatic.
|
||||
|
||||
### Drop-then-reuse the same table name (Enterprise/Cloud)
|
||||
|
||||
Avoid dropping or overwriting a remote table and then reusing that name right away:
|
||||
|
||||
```python
|
||||
db.drop_table("my_table")
|
||||
table = db.create_table("my_table", data=rows) # reads 500 for ~5 min
|
||||
table = db.create_table("my_table", data=rows, mode="overwrite") # same problem
|
||||
```
|
||||
|
||||
Why: Enterprise/Cloud splits DDL (control plane) from query serving (data plane). The data plane caches the dataset behind a table name for up to `table_cache_ttl` (default 300s / 5 min), so after a drop/overwrite the DDL succeeds but queries against the reused name return `500 Internal Server Error` until the cache expires — and a fresh `describe` may still show the old schema. Instead, write to a **fresh name**, use `list_tables()` and fail if it already exists, then `rename_table(fresh, final)` onto the final name only after the old table's drop has propagated (~5 min). See the "Enterprise: never drop-then-reuse the same table name" section in `SKILL.md`. Local/OSS tables have no separate data plane — overwrite freely there.
|
||||
|
||||
### Guessing performance fixes
|
||||
|
||||
Avoid changing `nprobes`, `refine_factor`, or index types before checking the query plan and index stats. Diagnose first, then tune one knob at a time.
|
||||
@@ -0,0 +1,131 @@
|
||||
# Python Performance Guidance
|
||||
|
||||
Use this when writing Python code that ingests data, queries large tables, builds indexes, or investigates latency.
|
||||
|
||||
## Ingestion
|
||||
|
||||
### Recommended: validate schemas and records with Pydantic
|
||||
|
||||
Favor `LanceModel` for readable Python schema definitions and validate source
|
||||
records before writing. Use PyArrow directly for Arrow-native or streaming
|
||||
pipelines where it is the clearer representation.
|
||||
|
||||
```python
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
class Document(LanceModel):
|
||||
id: int
|
||||
text: str
|
||||
vector: Vector(384, nullable=False)
|
||||
|
||||
rows = [Document.model_validate(row) for row in source_rows]
|
||||
table = db.create_table("documents", schema=Document)
|
||||
table.add(rows)
|
||||
```
|
||||
|
||||
### Recommended: bulk ingestion for materialized data
|
||||
|
||||
```python
|
||||
table.add(arrow_table)
|
||||
table.add(df)
|
||||
table.add(pa.dataset("data/", format="parquet"))
|
||||
```
|
||||
|
||||
For very large initial loads, create the table empty first, then call `add(...)`. Passing data directly to `create_table(name, data)` can skip the auto-parallel write path.
|
||||
|
||||
### Recommended: iterator ingestion for generated or streamed data
|
||||
|
||||
```python
|
||||
def batches():
|
||||
for raw in source:
|
||||
vectors = model.encode(raw["text"])
|
||||
yield pa.RecordBatch.from_pydict({**raw, "vector": vectors})
|
||||
|
||||
table.add(batches())
|
||||
```
|
||||
|
||||
Use chunks of several thousand rows or more when practical. Tiny batches and per-row writes create many small fragments.
|
||||
|
||||
### Anti-pattern: per-row `add()`
|
||||
|
||||
```python
|
||||
for row in rows:
|
||||
table.add([row])
|
||||
```
|
||||
|
||||
Each call creates a version and fragment. This slows ingestion and later queries.
|
||||
|
||||
## Indexing
|
||||
|
||||
- Build a vector index once brute-force vector search becomes too slow. As a rule of thumb, local brute force is fine below roughly 100K vectors; beyond that, build an index.
|
||||
- Use `IVF_PQ` as the general-purpose default. Enterprise builds this automatically.
|
||||
- Use scalar indexes for filtered columns and merge/upsert keys.
|
||||
- Use `BTREE` for mostly distinct numeric/string/temporal columns, `BITMAP` for booleans and low-cardinality columns, and `LABEL_LIST` for list membership queries.
|
||||
- Keep full-text defaults unless phrase queries require position data.
|
||||
|
||||
## Querying
|
||||
|
||||
Always be explicit:
|
||||
|
||||
```python
|
||||
table.search(query_vector).select(["id", "title"]).limit(20)
|
||||
```
|
||||
|
||||
- `select()` reduces bytes read and transferred.
|
||||
- `limit()` prevents accidental full-table materialization.
|
||||
- Pre-filtering is the default and guarantees returned rows satisfy the predicate.
|
||||
- Use post-filtering only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Recall Tuning
|
||||
|
||||
Tune one knob at a time:
|
||||
|
||||
- Quantized indexes: raise `refine_factor` to rescore more candidates on full vectors.
|
||||
- HNSW-backed indexes: raise `ef`; start around `1.5 * k`, increase toward `10 * k` if recall is short.
|
||||
- IVF candidate breadth: `nprobes` is auto-tuned; override only when a selective pre-filter leaves too few neighbors.
|
||||
|
||||
## Maintenance
|
||||
|
||||
After every successful embedded OSS/local ingestion, call `table.optimize()`.
|
||||
Do not add this to LanceDB Enterprise/Cloud remote table code; remote compaction
|
||||
and cleanup are handled automatically based on the Enterprise cluster
|
||||
configuration.
|
||||
|
||||
Why local maintenance is needed:
|
||||
|
||||
- Frequent writes can create many small fragments. Queries then need to scan across more files, which can increase latency.
|
||||
- Updates, deletes, and appends create new table versions. Old versions are retained for time travel and rollback, which can grow disk usage.
|
||||
- Indexes may have newly added rows that are not yet fully optimized into the index structure.
|
||||
|
||||
For local/OSS tables, run `optimize()` after the final successful ingestion
|
||||
write. Also run it after later batches of update/delete operations or on a
|
||||
regular maintenance schedule:
|
||||
|
||||
```python
|
||||
table.optimize()
|
||||
```
|
||||
|
||||
If the user wants more aggressive local disk cleanup, pass a shorter cleanup retention window:
|
||||
|
||||
```python
|
||||
from datetime import timedelta
|
||||
|
||||
table.optimize(cleanup_older_than=timedelta(days=1))
|
||||
```
|
||||
|
||||
Do not use very short cleanup windows when the application depends on time travel, rollback, or old versions.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
Before changing code or indexes, inspect:
|
||||
|
||||
```python
|
||||
print(table.search(query_vector).where("year > 2000").limit(10).analyze_plan())
|
||||
print(table.index_stats("vector_idx"))
|
||||
```
|
||||
|
||||
Look for high scan bytes, missing indexes, fragmented data, and unindexed rows.
|
||||
|
||||
## Python Multiprocessing
|
||||
|
||||
When using multiprocessing, use `spawn` rather than `fork`. LanceDB is multi-threaded internally, and `fork` plus a multi-threaded process is unsafe.
|
||||
@@ -0,0 +1,78 @@
|
||||
# TypeScript API Reference
|
||||
|
||||
Quick method reference for TypeScript LanceDB code. Cross-check source for non-trivial claims.
|
||||
|
||||
## Connect
|
||||
|
||||
```typescript
|
||||
import * as lancedb from "@lancedb/lancedb";
|
||||
|
||||
const db = await lancedb.connect("./camelot-db");
|
||||
```
|
||||
|
||||
**Place the local database directory next to the script/entrypoint that opens it** (resolve the path relative to the module, e.g. via `import.meta.dirname` / `__dirname`), not buried under a shared `data/` folder. The Lance dataset is the database, not a data file — keeping it beside its code makes ownership obvious and paths stable regardless of the working directory the script is launched from.
|
||||
|
||||
**Do not name the directory `lancedb`** (e.g. `./lancedb`, `./data/lancedb`). It collides with the imported `lancedb` package/namespace, which is confusing to read. Give it a name derived from the repo or dataset with a clear prefix/suffix — for example `./<dataset>-db`, `./<repo>_lancedb`, or `./vectordb`.
|
||||
|
||||
Remote connections use `db://...` plus Enterprise/Cloud credentials and deployment settings. Check current source/docs for exact connection options.
|
||||
|
||||
## Table Reads
|
||||
|
||||
| Task | Preferred API |
|
||||
| --- | --- |
|
||||
| Vector search | `table.search(queryVector).limit(k)` |
|
||||
| Full scan with filters/projection | `table.query().where(...).select(...).limit(...)` |
|
||||
| Filter | `.where("col > 10")` |
|
||||
| Projection | `.select(["id", "text"])` |
|
||||
| Bound result count | `.limit(20)` |
|
||||
| Collect bounded result as objects | `.toArray()` on query/search result |
|
||||
| Collect bounded result as Arrow | `.toArrow()` on query/search result |
|
||||
| Stream result batches | `for await (const batch of table.query()...)` |
|
||||
|
||||
## Local vs Remote Safety
|
||||
|
||||
| API | Agent guidance |
|
||||
| --- | --- |
|
||||
| `table.search(...)` | Preferred read path |
|
||||
| `table.query()` | Preferred scan/filter path |
|
||||
| `await table.toArrow()` | Avoid in portable or large-table code |
|
||||
| `await table.query().toArray()` with no `limit()` | Avoid; unbounded collection |
|
||||
| `await table.query().toArrow()` with no `limit()` | Avoid; unbounded collection |
|
||||
|
||||
## Indexes
|
||||
|
||||
```typescript
|
||||
await table.createIndex("vector");
|
||||
await table.createIndex("status");
|
||||
```
|
||||
|
||||
Use vector indexes for large vector search workloads and scalar indexes for filtered columns or merge/upsert keys. Check source/docs before specifying advanced index options.
|
||||
|
||||
## Filtering And Recall Knobs
|
||||
|
||||
```typescript
|
||||
await table.search(queryVector).where("status = 'ready'").limit(10).toArray();
|
||||
await table.search(queryVector).limit(10).refineFactor(20).toArray();
|
||||
await table.search(queryVector).limit(10).nprobes(50).toArray();
|
||||
await table.search(queryVector).limit(10).ef(100).toArray();
|
||||
await table.search(queryVector).where("status = 'ready'").postfilter().limit(10).toArray();
|
||||
```
|
||||
|
||||
Use `postfilter()` only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
```typescript
|
||||
console.log(await table.search(queryVector).where("year > 2000").limit(10).analyzePlan());
|
||||
console.log(await table.indexStats("vector_idx"));
|
||||
```
|
||||
|
||||
Use these before changing indexes or search tuning.
|
||||
|
||||
## Maintenance
|
||||
|
||||
```typescript
|
||||
await table.optimize();
|
||||
```
|
||||
|
||||
Call this after every successful local/OSS ingestion. It handles compaction, cleanup of old versions according to retention, and index optimization. Do not add this for LanceDB Enterprise/Cloud remote tables; Enterprise handles compaction and cleanup automatically from cluster configuration.
|
||||
@@ -0,0 +1,100 @@
|
||||
# TypeScript Patterns
|
||||
|
||||
Use these patterns when writing TypeScript code with `@lancedb/lancedb`.
|
||||
|
||||
## Recommended Patterns
|
||||
|
||||
### Bounded query
|
||||
|
||||
Use this for application reads, scripts, and examples:
|
||||
|
||||
```typescript
|
||||
const rows = await table
|
||||
.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.toArray();
|
||||
```
|
||||
|
||||
### Bounded vector search
|
||||
|
||||
```typescript
|
||||
const rows = await table
|
||||
.search(queryVector)
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.toArray();
|
||||
```
|
||||
|
||||
### Batch streaming for larger reads
|
||||
|
||||
When the task needs many rows, avoid collecting everything at once:
|
||||
|
||||
```typescript
|
||||
for await (const batch of table
|
||||
.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(10_000)) {
|
||||
process(batch);
|
||||
}
|
||||
```
|
||||
|
||||
## Anti-Patterns
|
||||
|
||||
**Avoid the following anti-patterns in your code.**
|
||||
|
||||
### Table-level full materialization
|
||||
|
||||
Avoid whole-table collectors in portable or large-table code:
|
||||
|
||||
```typescript
|
||||
const tableArrow = await table.toArrow();
|
||||
```
|
||||
|
||||
Why: local tables expose these whole-table collectors, but remote tables intentionally do not — a remote production table can be far larger than a local development table, so it is easy to accidentally pull the entire table into memory.
|
||||
|
||||
### Unbounded result collection
|
||||
|
||||
Avoid query/search collection without a meaningful limit:
|
||||
|
||||
```typescript
|
||||
const rows = await table.query().toArray(); // unbounded plain scan
|
||||
const rows = await table.search(queryVector).toArray(); // unbounded vector search
|
||||
```
|
||||
|
||||
Prefer `select(...).limit(...)` before collecting; for large reads, stream in batches instead.
|
||||
|
||||
### Per-row writes
|
||||
|
||||
Avoid loops that write one row per call:
|
||||
|
||||
```typescript
|
||||
for (const row of rows) {
|
||||
await table.add([row]); // one commit + fragment per row
|
||||
}
|
||||
```
|
||||
|
||||
Each `add()` creates a new version and fragment. Pass the whole batch in a single call, or chunk very large inputs:
|
||||
|
||||
```typescript
|
||||
await table.add(rows); // single commit
|
||||
// for very large inputs, add in chunks of several thousand rows
|
||||
```
|
||||
|
||||
### Drop-then-reuse the same table name (Enterprise/Cloud)
|
||||
|
||||
Avoid dropping or overwriting a remote table and then reusing that name right away:
|
||||
|
||||
```typescript
|
||||
await db.dropTable("my_table");
|
||||
const table = await db.createTable("my_table", rows); // reads 500 for ~5 min
|
||||
const table = await db.createTable("my_table", rows, { mode: "overwrite" }); // same problem
|
||||
```
|
||||
|
||||
Why: Enterprise/Cloud splits DDL (control plane) from query serving (data plane). The data plane caches the dataset behind a table name for up to `table_cache_ttl` (default 300s / 5 min), so after a drop/overwrite the DDL succeeds but queries against the reused name return `500 Internal Server Error` until the cache expires — and a fresh `describe` may still show the old schema. Instead, write to a **fresh name**, use `tableNames()` and fail if it already exists, then `renameTable(fresh, final)` onto the final name only after the old table's drop has propagated (~5 min). See the "Enterprise: never drop-then-reuse the same table name" section in `SKILL.md`. Local/OSS tables have no separate data plane — overwrite freely there.
|
||||
|
||||
### Guessing performance fixes
|
||||
|
||||
Avoid changing `nprobes`, `refineFactor`, `ef`, or index settings before checking `analyzePlan()` and `indexStats(...)`. Diagnose first, then tune one knob at a time.
|
||||
@@ -0,0 +1,78 @@
|
||||
# TypeScript Performance Guidance
|
||||
|
||||
Use this when writing TypeScript code that ingests data, queries large tables, builds indexes, or investigates latency.
|
||||
|
||||
## Ingestion
|
||||
|
||||
- Prefer bulk or batched writes.
|
||||
- Avoid per-row write loops; they create many small commits/fragments.
|
||||
- For generated data, accumulate reasonable batches before adding.
|
||||
- For file-backed data, prefer APIs that stream from Arrow/Parquet-style inputs when available.
|
||||
|
||||
## Indexing
|
||||
|
||||
- Build a vector index once brute-force vector search becomes too slow. As a rule of thumb, local brute force is fine below roughly 100K vectors; beyond that, build an index.
|
||||
- Use the general-purpose vector index defaults unless the task has explicit recall/latency requirements.
|
||||
- Build scalar indexes for filtered columns and merge/upsert keys.
|
||||
- Use full-text index phrase options only when phrase queries require them.
|
||||
|
||||
## Querying
|
||||
|
||||
Always be explicit:
|
||||
|
||||
```typescript
|
||||
await table.search(queryVector).select(["id", "title"]).limit(20).toArray();
|
||||
```
|
||||
|
||||
- `select()` reduces bytes read and transferred.
|
||||
- `limit()` prevents accidental full-table collection.
|
||||
- Pre-filtering is the default behavior. Use `postfilter()` only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Recall Tuning
|
||||
|
||||
Tune one knob at a time:
|
||||
|
||||
- Quantized indexes: raise `refineFactor(...)` to rescore more candidates on full vectors.
|
||||
- HNSW-backed indexes: raise `ef(...)`; start around `1.5 * k`, increase toward `10 * k` if recall is short.
|
||||
- IVF candidate breadth: `nprobes(...)` is usually auto-tuned; override only when a selective pre-filter leaves too few neighbors.
|
||||
|
||||
## Maintenance
|
||||
|
||||
After every successful embedded OSS/local ingestion, call `table.optimize()`.
|
||||
Do not add this to LanceDB Enterprise/Cloud remote table code; remote compaction
|
||||
and cleanup are handled automatically based on the Enterprise cluster
|
||||
configuration.
|
||||
|
||||
Why local maintenance is needed:
|
||||
|
||||
- Frequent writes can create many small fragments. Queries then need to scan across more files, which can increase latency.
|
||||
- Updates, deletes, and appends create new table versions. Old versions are retained for time travel and rollback, which can grow disk usage.
|
||||
- Indexes may have newly added rows that are not yet fully optimized into the index structure.
|
||||
|
||||
For local/OSS tables, run `optimize()` after the final successful ingestion
|
||||
write. Also run it after later batches of update/delete operations or on a
|
||||
regular maintenance schedule:
|
||||
|
||||
```typescript
|
||||
await table.optimize();
|
||||
```
|
||||
|
||||
If the user wants more aggressive local disk cleanup, pass a shorter cleanup retention window:
|
||||
|
||||
```typescript
|
||||
const olderThan = new Date(Date.now() - 24 * 60 * 60 * 1000);
|
||||
await table.optimize({ cleanupOlderThan: olderThan });
|
||||
```
|
||||
|
||||
Do not use very short cleanup windows when the application depends on time travel, rollback, or old versions.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
Before changing code or indexes, inspect:
|
||||
|
||||
```typescript
|
||||
console.log(await table.search(queryVector).where("year > 2000").limit(10).analyzePlan());
|
||||
console.log(await table.indexStats("vector_idx"));
|
||||
```
|
||||
|
||||
Look for high scan cost, missing indexes, fragmented data, and unindexed rows.
|
||||
@@ -0,0 +1,135 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Scan Python and TypeScript for likely unsafe LanceDB materialization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PY_FULL_TABLE = re.compile(r"\b\w+\.(to_pandas|to_arrow|to_polars)\s*\(")
|
||||
TS_TABLE_TO_ARROW = re.compile(r"\b\w+\.toArrow\s*\(")
|
||||
TS_QUERY_COLLECTOR = re.compile(r"\.query\s*\(\s*\)[\s\S]*?\.to(Array|Arrow)\s*\(")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Finding:
|
||||
path: Path
|
||||
line: int
|
||||
message: str
|
||||
text: str
|
||||
|
||||
|
||||
def iter_files(paths: list[Path]) -> list[Path]:
|
||||
files: list[Path] = []
|
||||
for path in paths:
|
||||
if path.is_dir():
|
||||
files.extend(
|
||||
p
|
||||
for p in path.rglob("*")
|
||||
if p.suffix in {".py", ".ts", ".tsx"} and "node_modules" not in p.parts
|
||||
)
|
||||
elif path.suffix in {".py", ".ts", ".tsx"}:
|
||||
files.append(path)
|
||||
return sorted(set(files))
|
||||
|
||||
|
||||
def line_number(text: str, offset: int) -> int:
|
||||
return text.count("\n", 0, offset) + 1
|
||||
|
||||
|
||||
def scan_python(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in PY_FULL_TABLE.finditer(text):
|
||||
line_start = text.rfind("\n", 0, match.start()) + 1
|
||||
line_end = text.find("\n", match.start())
|
||||
if line_end == -1:
|
||||
line_end = len(text)
|
||||
line = text[line_start:line_end].strip()
|
||||
if ".search(" in line or ".query(" in line:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
f"Review Python `{match.group(1)}()` call; table-level materialization is not portable to remote tables.",
|
||||
line,
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def statement_around(text: str, start: int, end: int) -> str:
|
||||
before = max(text.rfind(";", 0, start), text.rfind("\n\n", 0, start))
|
||||
after_candidates = [pos for pos in (text.find(";", end), text.find("\n\n", end)) if pos != -1]
|
||||
after = min(after_candidates) if after_candidates else len(text)
|
||||
return text[before + 1 : after].strip()
|
||||
|
||||
|
||||
def scan_typescript(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in TS_TABLE_TO_ARROW.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".query(" in stmt or ".search(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review TypeScript `table.toArrow()`-style call; table-level materialization is not portable for large/remote tables.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in TS_QUERY_COLLECTOR.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".limit(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review unbounded TypeScript query collection; add `limit()` or stream batches.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def scan_file(path: Path) -> list[Finding]:
|
||||
text = path.read_text(encoding="utf-8", errors="replace")
|
||||
if path.suffix == ".py":
|
||||
return scan_python(path, text)
|
||||
if path.suffix in {".ts", ".tsx"}:
|
||||
return scan_typescript(path, text)
|
||||
return []
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("paths", nargs="+", type=Path)
|
||||
parser.add_argument(
|
||||
"--no-fail", action="store_true", help="Always exit 0 after reporting findings."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
findings: list[Finding] = []
|
||||
for path in iter_files(args.paths):
|
||||
findings.extend(scan_file(path))
|
||||
|
||||
for finding in findings:
|
||||
print(f"{finding.path}:{finding.line}: {finding.message}")
|
||||
print(f" {finding.text}")
|
||||
|
||||
if findings:
|
||||
print(
|
||||
f"\n{len(findings)} finding(s). Review manually; bounded query result conversion may be OK."
|
||||
)
|
||||
return 0 if args.no_fail or not findings else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.31.0-beta.2"
|
||||
current_version = "0.32.0-beta.1"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -34,15 +34,16 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
target: x86_64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
curl -L https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-$(uname -m).zip > /tmp/protoc.zip \
|
||||
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||
&& rm /tmp/protoc.zip
|
||||
curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-x86_64.zip -o /tmp/protoc.zip
|
||||
unzip /tmp/protoc.zip -d /usr/local
|
||||
rm /tmp/protoc.zip
|
||||
/usr/local/bin/protoc --version
|
||||
- name: Build Arm Manylinux Wheel
|
||||
if: ${{ inputs.arm-build == 'true' }}
|
||||
uses: PyO3/maturin-action@v1
|
||||
@@ -50,13 +51,14 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
target: aarch64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
yum install -y clang \
|
||||
&& curl -L https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-aarch_64.zip > /tmp/protoc.zip \
|
||||
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||
&& rm /tmp/protoc.zip
|
||||
yum install -y clang
|
||||
curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-aarch_64.zip -o /tmp/protoc.zip
|
||||
unzip /tmp/protoc.zip -d /usr/local
|
||||
rm /tmp/protoc.zip
|
||||
/usr/local/bin/protoc --version
|
||||
|
||||
@@ -25,7 +25,7 @@ jobs:
|
||||
# Only runs on tags that matches the make-release action
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
workspaces: rust
|
||||
@@ -47,7 +47,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,14 +36,14 @@ jobs:
|
||||
echo "guidelines = ${{ inputs.guidelines }}"
|
||||
|
||||
- name: Checkout Repo
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.branch }}
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 (used by the nodejs install step below) requires
|
||||
# Node >= 22.13; use 24 since 22 hits EOL in October.
|
||||
@@ -82,7 +82,7 @@ jobs:
|
||||
cache: maven
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Install Node.js dependencies for TypeScript bindings
|
||||
|
||||
@@ -30,13 +30,13 @@ jobs:
|
||||
echo "tag = ${{ inputs.tag || 'latest' }}"
|
||||
|
||||
- name: Checkout Repo LanceDB
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
name: Verify PR title / description conforms to semantic-release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: "18"
|
||||
# These rules are disabled because Github will always ensure there
|
||||
|
||||
@@ -35,7 +35,7 @@ jobs:
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install dependencies needed for ubuntu
|
||||
run: |
|
||||
sudo apt install -y protobuf-compiler libssl-dev
|
||||
@@ -53,7 +53,7 @@ jobs:
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -e .
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -r ../docs/requirements.txt
|
||||
- name: Set up node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
cache: 'npm'
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Set up Java 8
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
@@ -73,7 +73,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,7 +36,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Set up Java 17
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
|
||||
@@ -19,7 +19,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install license-header-checker
|
||||
working-directory: /tmp
|
||||
run: |
|
||||
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
steps:
|
||||
- name: Output Inputs
|
||||
run: echo "${{ toJSON(github.event.inputs) }}"
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
@@ -38,14 +38,14 @@ jobs:
|
||||
CC: gcc-12
|
||||
CXX: g++-12
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October. The library itself still supports Node >= 18
|
||||
@@ -86,14 +86,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
name: Setup Node.js 24 for build
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
@@ -130,7 +130,7 @@ jobs:
|
||||
echo "Run 'pnpm run docs', fix any warnings, and commit the changes."
|
||||
exit 1
|
||||
fi
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
name: Setup Node.js ${{ matrix.node-version }} for test
|
||||
with:
|
||||
node-version: ${{ matrix.node-version }}
|
||||
@@ -166,14 +166,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -103,7 +103,7 @@ jobs:
|
||||
features: fp16kernels
|
||||
pre_build: brew install protobuf
|
||||
- target: x86_64-pc-windows-msvc
|
||||
host: windows-latest
|
||||
host: windows-2025-8x-x64
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
@@ -111,12 +111,21 @@ jobs:
|
||||
# There is an issue where choco doesn't add nasm to the path
|
||||
export PATH="$PATH:/c/Program Files/NASM"
|
||||
nasm -v
|
||||
# Fat LTO of the cdylib is single-threaded and the peak-memory
|
||||
# step of the build, and had started hitting rustc-LLVM OOM on the
|
||||
# Windows runners. ThinLTO parallelizes it across the runner's
|
||||
# cores and keeps peak memory well under the limit.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: aarch64-pc-windows-msvc
|
||||
host: windows-latest
|
||||
host: windows-2025-8x-x64
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc
|
||||
rustup target add aarch64-pc-windows-msvc
|
||||
# See ThinLTO note on the x86_64-pc-windows-msvc target above.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-unknown-linux-gnu
|
||||
host: ubuntu-latest
|
||||
features: fp16kernels
|
||||
@@ -170,13 +179,13 @@ jobs:
|
||||
run:
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -190,7 +199,7 @@ jobs:
|
||||
toolchain: stable
|
||||
targets: ${{ matrix.settings.target }}
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v4
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry/index/
|
||||
@@ -244,7 +253,7 @@ jobs:
|
||||
if: ${{ !matrix.settings.docker }}
|
||||
shell: bash
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/*.node
|
||||
@@ -256,7 +265,7 @@ jobs:
|
||||
run: pnpm tsc
|
||||
- name: Upload Generic Artifacts
|
||||
if: ${{ matrix.settings.target == 'aarch64-apple-darwin' }}
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: |
|
||||
@@ -287,13 +296,13 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup Node.js 24 for install
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -303,18 +312,18 @@ jobs:
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Setup Node.js ${{ matrix.node }} for test
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: ${{ matrix.node }}
|
||||
- name: Download artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
@@ -339,13 +348,13 @@ jobs:
|
||||
needs:
|
||||
- test-lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
@@ -353,14 +362,14 @@ jobs:
|
||||
registry-url: "https://registry.npmjs.org"
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
name: Download arch-specific binaries
|
||||
with:
|
||||
pattern: lancedb-*
|
||||
@@ -398,7 +407,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -41,7 +41,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -66,7 +66,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -95,7 +95,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -160,7 +160,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -189,7 +189,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -212,7 +212,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
+27
-11
@@ -40,7 +40,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -65,7 +65,7 @@ jobs:
|
||||
timeout-minutes: 10
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: EmbarkStudios/cargo-deny-action@v2
|
||||
with:
|
||||
command: check advisories bans licenses sources
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
CC: clang
|
||||
CXX: clang++
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
# Building without a lock file often requires the latest Rust version since downstream
|
||||
# dependencies may have updated their minimum Rust version.
|
||||
- uses: actions-rust-lang/setup-rust-toolchain@v1
|
||||
@@ -113,7 +113,7 @@ jobs:
|
||||
CXX: clang++-18
|
||||
GH_TOKEN: ${{ secrets.SOPHON_READ_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -125,10 +125,26 @@ jobs:
|
||||
- uses: rui314/setup-mold@v1
|
||||
- name: Make Swap
|
||||
run: |
|
||||
sudo fallocate -l 16G /swapfile
|
||||
sudo chmod 600 /swapfile
|
||||
sudo mkswap /swapfile
|
||||
sudo swapon /swapfile
|
||||
swapfile=/swapfile
|
||||
min_swap_bytes=$((15 * 1024 * 1024 * 1024))
|
||||
active_swap_bytes="$(sudo swapon --show=NAME,SIZE --bytes --noheadings | awk '$1 == "/swapfile" { print $2 }')"
|
||||
if [ -n "$active_swap_bytes" ]; then
|
||||
if [ "$active_swap_bytes" -ge "$min_swap_bytes" ]; then
|
||||
echo "/swapfile is already active with enough space; skipping swap creation"
|
||||
exit 0
|
||||
fi
|
||||
echo "/swapfile is already active but smaller than 16G; using /mnt/lancedb-swapfile"
|
||||
swapfile=/mnt/lancedb-swapfile
|
||||
fi
|
||||
if sudo swapon --show=NAME --noheadings | grep -Fxq "$swapfile"; then
|
||||
echo "$swapfile is already active; skipping swap creation"
|
||||
exit 0
|
||||
fi
|
||||
sudo rm -f "$swapfile"
|
||||
sudo fallocate -l 16G "$swapfile"
|
||||
sudo chmod 600 "$swapfile"
|
||||
sudo mkswap "$swapfile"
|
||||
sudo swapon "$swapfile"
|
||||
- name: Build
|
||||
run: cargo build --profile ci --all-features --tests --locked --examples
|
||||
- name: Run feature tests
|
||||
@@ -152,7 +168,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: rust
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -181,7 +197,7 @@ jobs:
|
||||
run:
|
||||
working-directory: rust/lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Set target
|
||||
run: rustup target add ${{ matrix.target }}
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -210,7 +226,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Install dependencies
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
Generated
+394
-141
File diff suppressed because it is too large
Load Diff
+17
-14
@@ -13,24 +13,25 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.0.0-beta.8", default-features = false, "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.8", default-features = false, "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.8", default-features = false, "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.8", "tag" = "v9.0.0-beta.8", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
arrow-array = "58.0.0"
|
||||
arrow-buffer = "58.0.0"
|
||||
arrow-data = "58.0.0"
|
||||
arrow-ipc = "58.0.0"
|
||||
arrow-ord = "58.0.0"
|
||||
@@ -53,6 +54,8 @@ half = { "version" = "2.7.1", default-features = false, features = [
|
||||
] }
|
||||
futures = "0"
|
||||
log = "0.4"
|
||||
metrics = "0.24"
|
||||
metrics-util = "0.19"
|
||||
moka = { version = "0.12", features = ["future"] }
|
||||
object_store = "0.13.2"
|
||||
pin-project = "1.0.7"
|
||||
|
||||
@@ -51,18 +51,6 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0436
|
||||
{ id = "RUSTSEC-2024-0436", reason = "transitive via datafusion; awaiting ecosystem migration" },
|
||||
|
||||
# encoding: unmaintained. Reached through lindera-dictionary, which is
|
||||
# required by the native Lindera tokenizer path. Lindera has not migrated
|
||||
# off this crate yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2021-0153
|
||||
{ id = "RUSTSEC-2021-0153", reason = "transitive via lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# fast-float: unsound and unmaintained. Reached only through polars-arrow
|
||||
# from the optional Polars integration; replacement requires a Polars
|
||||
# dependency upgrade.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0379
|
||||
{ id = "RUSTSEC-2024-0379", reason = "transitive via polars-arrow; waiting on Polars migration" },
|
||||
|
||||
# tantivy: segfault on malformed input due to missing bounds check.
|
||||
# Pulled in via lance for full-text search. We only feed tantivy
|
||||
# documents we construct ourselves, not attacker-controlled bytes.
|
||||
@@ -80,18 +68,6 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0119
|
||||
{ id = "RUSTSEC-2025-0119", reason = "transitive via hf-hub/indicatif; cosmetic formatting crate" },
|
||||
|
||||
# bincode: unmaintained. Reached through lindera and lindera-dictionary,
|
||||
# which are required by the native Lindera tokenizer path. Lindera has not
|
||||
# migrated to another serialization format yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0141
|
||||
{ id = "RUSTSEC-2025-0141", reason = "transitive via lindera/lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# lru: soundness issue in IterMut. Reached only through aws-sdk-s3 in
|
||||
# LanceDB's dev-dependency graph; LanceDB does not use that iterator
|
||||
# directly. Clearing this requires the AWS SDK chain to update lru.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0002
|
||||
{ id = "RUSTSEC-2026-0002", reason = "transitive via aws-sdk-s3 dev-dependency; waiting on AWS SDK lru upgrade" },
|
||||
|
||||
# rustls-webpki 0.101.7 (old major line): name-constraint checks for
|
||||
# URI / wildcard names. Pulled in only via the legacy rustls 0.21 chain
|
||||
# from aws-smithy-http-client. The 0.103 line we actively use is patched.
|
||||
@@ -108,17 +84,23 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0104
|
||||
{ id = "RUSTSEC-2026-0104", reason = "only affects rustls-webpki 0.101 from legacy aws-smithy/rustls 0.21 chain" },
|
||||
|
||||
# rand 0.8.5: soundness issue only when ThreadRng reseeds inside a custom
|
||||
# logger. Reached through several transitive chains. LanceDB does not use
|
||||
# rand from a custom logger; upgrade once all pinned chains accept 0.8.6+.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0097
|
||||
{ id = "RUSTSEC-2026-0097", reason = "transitive rand 0.8.5; LanceDB does not call ThreadRng from custom logging" },
|
||||
|
||||
# pyo3 advisories in the Python bindings; tracked pending a patched pyo3 release.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0176
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0177
|
||||
{ id = "RUSTSEC-2026-0176", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
{ id = "RUSTSEC-2026-0177", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
|
||||
# quick-xml < 0.41.0: quadratic runtime on duplicate attribute names (DoS).
|
||||
# quick-xml < 0.41.0: unbounded namespace-declaration allocation in NsReader (DoS).
|
||||
# Pulled in transitively by inferno (dev-only flame-graph dep), lance-namespace-impls
|
||||
# (git dep from lance), and opendal/reqsign (cloud storage XML parsing). The XML
|
||||
# parsed by opendal/reqsign comes from trusted cloud-storage endpoints (S3, GCS,
|
||||
# Azure), not attacker-controlled input. Clearing requires upstream crates to migrate
|
||||
# to quick-xml >= 0.41.0.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0194
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0195
|
||||
{ id = "RUSTSEC-2026-0194", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
{ id = "RUSTSEC-2026-0195", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.31.0-beta.2</version>
|
||||
<version>0.32.0-beta.1</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -518,6 +518,9 @@ 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
|
||||
|
||||
@@ -398,6 +398,26 @@ Drop an index from the table.
|
||||
|
||||
***
|
||||
|
||||
### getLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
abstract getLsmWriteSpec(): Promise<undefined | LsmWriteSpec>
|
||||
```
|
||||
|
||||
Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) currently installed on this table.
|
||||
|
||||
Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
spec has been set, or it was removed with [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)).
|
||||
The returned spec — including its `maintainedIndexes` and
|
||||
`writerConfigDefaults` — mirrors what was passed to
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec).
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`undefined` \| [`LsmWriteSpec`](../interfaces/LsmWriteSpec.md)>
|
||||
|
||||
***
|
||||
|
||||
### indexStats()
|
||||
|
||||
```ts
|
||||
@@ -914,6 +934,32 @@ Return the table as an arrow table
|
||||
|
||||
***
|
||||
|
||||
### tokenize()
|
||||
|
||||
```ts
|
||||
abstract tokenize(query, options): Promise<FtsToken[]>
|
||||
```
|
||||
|
||||
Tokenize a full-text search query using the tokenizer configured on an FTS index.
|
||||
|
||||
Specify exactly one of `column` or `indexName`.
|
||||
|
||||
Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
the client process from index metadata. For remote tables, this means the
|
||||
same tokenizer model files must also exist locally.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **query**: `string`
|
||||
|
||||
* **options**: [`TokenizeTableOptions`](../type-aliases/TokenizeTableOptions.md)
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]>
|
||||
|
||||
***
|
||||
|
||||
### unsetLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -767,6 +767,9 @@ 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
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / OAuthFlowType
|
||||
|
||||
# Enumeration: OAuthFlowType
|
||||
|
||||
OAuth authentication flow types.
|
||||
|
||||
## Enumeration Members
|
||||
|
||||
### AzureManagedIdentity
|
||||
|
||||
```ts
|
||||
AzureManagedIdentity: "azure_managed_identity";
|
||||
```
|
||||
|
||||
Azure Managed Identity via IMDS.
|
||||
|
||||
***
|
||||
|
||||
### ClientCredentials
|
||||
|
||||
```ts
|
||||
ClientCredentials: "client_credentials";
|
||||
```
|
||||
|
||||
Client Credentials grant (service-to-service / M2M).
|
||||
@@ -0,0 +1,42 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / instrumentLanceDbMetrics
|
||||
|
||||
# Function: instrumentLanceDbMetrics()
|
||||
|
||||
```ts
|
||||
function instrumentLanceDbMetrics(meterProvider?): boolean
|
||||
```
|
||||
|
||||
Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
|
||||
Installs a process-global metrics recorder and creates one observable
|
||||
instrument per LanceDB metric (currently object store request counts, bytes,
|
||||
latency, errors, and throttles) on the given (or global) `MeterProvider`. The
|
||||
configured `MetricReader` then collects them on its own schedule.
|
||||
|
||||
Counters and gauges map directly to observable counters/gauges. Because
|
||||
OpenTelemetry has no asynchronous histogram instrument, each histogram is
|
||||
exported Prometheus-style as cumulative `le` bucket counts (`<name>_bucket`,
|
||||
with an `le` attribute) plus `<name>_count` and `<name>_sum`.
|
||||
|
||||
Requires `@opentelemetry/api` (a dependency) and, to actually export, an
|
||||
OpenTelemetry SDK such as `@opentelemetry/sdk-metrics`.
|
||||
|
||||
## Parameters
|
||||
|
||||
* **meterProvider?**: `MeterProvider`
|
||||
The provider to register instruments on. Defaults to the
|
||||
global provider from `@opentelemetry/api`.
|
||||
|
||||
## Returns
|
||||
|
||||
`boolean`
|
||||
|
||||
`true` if the recorder is installed and instruments are registered.
|
||||
`false` if a different `metrics` recorder is already installed in this
|
||||
process (only one global recorder is permitted), in which case a warning is
|
||||
emitted and no instruments are created. Calling this more than once is safe;
|
||||
instruments are created only on the first successful call.
|
||||
@@ -0,0 +1,26 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / tokenize
|
||||
|
||||
# Function: tokenize()
|
||||
|
||||
```ts
|
||||
function tokenize(query, options?): Promise<FtsToken[]>
|
||||
```
|
||||
|
||||
Tokenize a full-text search query using an explicit tokenizer.
|
||||
|
||||
This does not require a table or FTS index. The tokenizer options match
|
||||
[Index.fts](../classes/Index.md#fts).
|
||||
|
||||
## Parameters
|
||||
|
||||
* **query**: `string`
|
||||
|
||||
* **options?**: `Partial`<[`TokenizeOptions`](../interfaces/TokenizeOptions.md)>
|
||||
|
||||
## Returns
|
||||
|
||||
`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]>
|
||||
@@ -12,6 +12,7 @@
|
||||
## Enumerations
|
||||
|
||||
- [FullTextQueryType](enumerations/FullTextQueryType.md)
|
||||
- [OAuthFlowType](enumerations/OAuthFlowType.md)
|
||||
- [Occur](enumerations/Occur.md)
|
||||
- [Operator](enumerations/Operator.md)
|
||||
|
||||
@@ -71,6 +72,7 @@
|
||||
- [FragmentStatistics](interfaces/FragmentStatistics.md)
|
||||
- [FragmentSummaryStats](interfaces/FragmentSummaryStats.md)
|
||||
- [FtsOptions](interfaces/FtsOptions.md)
|
||||
- [FtsToken](interfaces/FtsToken.md)
|
||||
- [FullTextQuery](interfaces/FullTextQuery.md)
|
||||
- [FullTextSearchOptions](interfaces/FullTextSearchOptions.md)
|
||||
- [HnswPqOptions](interfaces/HnswPqOptions.md)
|
||||
@@ -85,6 +87,8 @@
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
||||
- [MergeResult](interfaces/MergeResult.md)
|
||||
- [NativeOAuthConfig](interfaces/NativeOAuthConfig.md)
|
||||
- [OAuthConfig](interfaces/OAuthConfig.md)
|
||||
- [OpenTableOptions](interfaces/OpenTableOptions.md)
|
||||
- [OptimizeOptions](interfaces/OptimizeOptions.md)
|
||||
- [OptimizeStats](interfaces/OptimizeStats.md)
|
||||
@@ -104,6 +108,7 @@
|
||||
- [TimeoutConfig](interfaces/TimeoutConfig.md)
|
||||
- [TlsConfig](interfaces/TlsConfig.md)
|
||||
- [TokenResponse](interfaces/TokenResponse.md)
|
||||
- [TokenizeOptions](interfaces/TokenizeOptions.md)
|
||||
- [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md)
|
||||
- [UpdateOptions](interfaces/UpdateOptions.md)
|
||||
- [UpdateResult](interfaces/UpdateResult.md)
|
||||
@@ -113,6 +118,7 @@
|
||||
|
||||
## Type Aliases
|
||||
|
||||
- [BaseTokenizer](type-aliases/BaseTokenizer.md)
|
||||
- [Data](type-aliases/Data.md)
|
||||
- [DataLike](type-aliases/DataLike.md)
|
||||
- [FieldLike](type-aliases/FieldLike.md)
|
||||
@@ -122,12 +128,15 @@
|
||||
- [RecordBatchLike](type-aliases/RecordBatchLike.md)
|
||||
- [SchemaLike](type-aliases/SchemaLike.md)
|
||||
- [TableLike](type-aliases/TableLike.md)
|
||||
- [TokenizeTableOptions](type-aliases/TokenizeTableOptions.md)
|
||||
|
||||
## Functions
|
||||
|
||||
- [RecordBatchIterator](functions/RecordBatchIterator.md)
|
||||
- [connect](functions/connect.md)
|
||||
- [connectNamespace](functions/connectNamespace.md)
|
||||
- [instrumentLanceDbMetrics](functions/instrumentLanceDbMetrics.md)
|
||||
- [makeArrowTable](functions/makeArrowTable.md)
|
||||
- [packBits](functions/packBits.md)
|
||||
- [permutationBuilder](functions/permutationBuilder.md)
|
||||
- [tokenize](functions/tokenize.md)
|
||||
|
||||
@@ -64,6 +64,19 @@ client used by manifest-enabled native connections.
|
||||
|
||||
***
|
||||
|
||||
### oauthConfig?
|
||||
|
||||
```ts
|
||||
optional oauthConfig: NativeOAuthConfig;
|
||||
```
|
||||
|
||||
(For LanceDB cloud only): OAuth configuration for IdP-based
|
||||
authentication (e.g., Azure Entra ID). When set, token acquisition
|
||||
and refresh are handled entirely in Rust. TypeScript users should pass
|
||||
the public `OAuthConfig` type exported from `@lancedb/lancedb`.
|
||||
|
||||
***
|
||||
|
||||
### readConsistencyInterval?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -23,7 +23,7 @@ whether to remove punctuation
|
||||
### baseTokenizer?
|
||||
|
||||
```ts
|
||||
optional baseTokenizer: "raw" | "simple" | "whitespace" | "ngram";
|
||||
optional baseTokenizer: BaseTokenizer;
|
||||
```
|
||||
|
||||
The tokenizer to use when building the index.
|
||||
@@ -37,6 +37,10 @@ The following tokenizers are available:
|
||||
|
||||
"raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token.
|
||||
|
||||
"icu" - ICU dictionary-based word segmentation.
|
||||
|
||||
"icu/split" - ICU segmentation with simple-style delimiter splitting.
|
||||
|
||||
***
|
||||
|
||||
### language?
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / FtsToken
|
||||
|
||||
# Interface: FtsToken
|
||||
|
||||
Token produced by the tokenizer configured on a full-text search index.
|
||||
|
||||
## Properties
|
||||
|
||||
### position
|
||||
|
||||
```ts
|
||||
position: number;
|
||||
```
|
||||
|
||||
Token position used by full-text query matching.
|
||||
|
||||
***
|
||||
|
||||
### text
|
||||
|
||||
```ts
|
||||
text: string;
|
||||
```
|
||||
|
||||
Token text after tokenizer filters have been applied.
|
||||
@@ -0,0 +1,88 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / NativeOAuthConfig
|
||||
|
||||
# Interface: NativeOAuthConfig
|
||||
|
||||
OAuth configuration for LanceDB authentication.
|
||||
|
||||
This is the generated napi-rs binding shape. TypeScript users should prefer
|
||||
the public `OAuthConfig` type exported from `@lancedb/lancedb`.
|
||||
|
||||
All token acquisition and refresh is handled in the Rust layer.
|
||||
|
||||
## Properties
|
||||
|
||||
### clientId
|
||||
|
||||
```ts
|
||||
clientId: string;
|
||||
```
|
||||
|
||||
Application / Client ID.
|
||||
|
||||
***
|
||||
|
||||
### clientSecret?
|
||||
|
||||
```ts
|
||||
optional clientSecret: string;
|
||||
```
|
||||
|
||||
Client secret (required for client_credentials).
|
||||
|
||||
***
|
||||
|
||||
### flow?
|
||||
|
||||
```ts
|
||||
optional flow: string;
|
||||
```
|
||||
|
||||
Authentication flow: "client_credentials" or "azure_managed_identity"
|
||||
|
||||
***
|
||||
|
||||
### issuerUrl
|
||||
|
||||
```ts
|
||||
issuerUrl: string;
|
||||
```
|
||||
|
||||
OIDC issuer URL or OAuth authority URL.
|
||||
For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
|
||||
|
||||
***
|
||||
|
||||
### managedIdentityClientId?
|
||||
|
||||
```ts
|
||||
optional managedIdentityClientId: string;
|
||||
```
|
||||
|
||||
Client ID for user-assigned managed identity (azure_managed_identity).
|
||||
|
||||
***
|
||||
|
||||
### refreshBufferSecs?
|
||||
|
||||
```ts
|
||||
optional refreshBufferSecs: number;
|
||||
```
|
||||
|
||||
Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
Keep this well below the token TTL; if it is greater than or equal to
|
||||
the TTL, each request refreshes the token.
|
||||
|
||||
***
|
||||
|
||||
### scopes
|
||||
|
||||
```ts
|
||||
scopes: string[];
|
||||
```
|
||||
|
||||
OAuth scopes to request. For Azure managed identity, exactly one scope
|
||||
or resource is required. For example: `["api://{app_id}/.default"]`
|
||||
@@ -0,0 +1,111 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / OAuthConfig
|
||||
|
||||
# Interface: OAuthConfig
|
||||
|
||||
OAuth configuration for LanceDB authentication.
|
||||
|
||||
This is the public TypeScript OAuth configuration type. The generated
|
||||
`NativeOAuthConfig` type has the same runtime shape but is an implementation
|
||||
detail of the napi-rs binding.
|
||||
|
||||
All token acquisition and refresh is handled in the Rust layer.
|
||||
This config is passed through to Rust via napi-rs.
|
||||
|
||||
## Examples
|
||||
|
||||
```typescript
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
clientId: "app-id",
|
||||
clientSecret: "secret",
|
||||
scopes: ["api://lancedb-api/.default"],
|
||||
};
|
||||
```
|
||||
|
||||
```typescript
|
||||
const config: OAuthConfig = {
|
||||
issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
clientId: "app-id",
|
||||
scopes: ["api://lancedb-api/.default"],
|
||||
flow: OAuthFlowType.AzureManagedIdentity,
|
||||
};
|
||||
```
|
||||
|
||||
## Properties
|
||||
|
||||
### clientId
|
||||
|
||||
```ts
|
||||
clientId: string;
|
||||
```
|
||||
|
||||
Application / Client ID.
|
||||
|
||||
***
|
||||
|
||||
### clientSecret?
|
||||
|
||||
```ts
|
||||
optional clientSecret: string;
|
||||
```
|
||||
|
||||
Client secret (required for ClientCredentials).
|
||||
|
||||
***
|
||||
|
||||
### flow?
|
||||
|
||||
```ts
|
||||
optional flow: OAuthFlowType;
|
||||
```
|
||||
|
||||
Authentication flow (default: ClientCredentials).
|
||||
|
||||
***
|
||||
|
||||
### issuerUrl
|
||||
|
||||
```ts
|
||||
issuerUrl: string;
|
||||
```
|
||||
|
||||
OIDC issuer URL or OAuth authority URL.
|
||||
For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
|
||||
|
||||
***
|
||||
|
||||
### managedIdentityClientId?
|
||||
|
||||
```ts
|
||||
optional managedIdentityClientId: string;
|
||||
```
|
||||
|
||||
Client ID for user-assigned managed identity (AzureManagedIdentity).
|
||||
|
||||
***
|
||||
|
||||
### refreshBufferSecs?
|
||||
|
||||
```ts
|
||||
optional refreshBufferSecs: number;
|
||||
```
|
||||
|
||||
Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
Keep this well below the token TTL; if it is greater than or equal to
|
||||
the TTL, each request refreshes the token.
|
||||
|
||||
***
|
||||
|
||||
### scopes
|
||||
|
||||
```ts
|
||||
scopes: string[];
|
||||
```
|
||||
|
||||
OAuth scopes to request.
|
||||
For Azure managed identity, exactly one scope or resource is required.
|
||||
For example: `["api://{app_id}/.default"]`
|
||||
@@ -8,6 +8,14 @@
|
||||
|
||||
## Properties
|
||||
|
||||
### clumpSize?
|
||||
|
||||
```ts
|
||||
optional clumpSize: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### counts?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenizeOptions
|
||||
|
||||
# Interface: TokenizeOptions
|
||||
|
||||
Options for tokenizing a full-text search query without a table index.
|
||||
|
||||
## Properties
|
||||
|
||||
### asciiFolding?
|
||||
|
||||
```ts
|
||||
optional asciiFolding: boolean;
|
||||
```
|
||||
|
||||
Whether to fold ASCII characters.
|
||||
|
||||
***
|
||||
|
||||
### baseTokenizer?
|
||||
|
||||
```ts
|
||||
optional baseTokenizer: BaseTokenizer;
|
||||
```
|
||||
|
||||
The tokenizer to use. The default is "simple".
|
||||
|
||||
***
|
||||
|
||||
### language?
|
||||
|
||||
```ts
|
||||
optional language: string;
|
||||
```
|
||||
|
||||
Language for stemming and stop words.
|
||||
|
||||
***
|
||||
|
||||
### lowercase?
|
||||
|
||||
```ts
|
||||
optional lowercase: boolean;
|
||||
```
|
||||
|
||||
Whether to lowercase tokens.
|
||||
|
||||
***
|
||||
|
||||
### maxTokenLength?
|
||||
|
||||
```ts
|
||||
optional maxTokenLength: number;
|
||||
```
|
||||
|
||||
Maximum token length; tokens longer than this are ignored.
|
||||
|
||||
***
|
||||
|
||||
### ngramMaxLength?
|
||||
|
||||
```ts
|
||||
optional ngramMaxLength: number;
|
||||
```
|
||||
|
||||
N-gram maximum length.
|
||||
|
||||
***
|
||||
|
||||
### ngramMinLength?
|
||||
|
||||
```ts
|
||||
optional ngramMinLength: number;
|
||||
```
|
||||
|
||||
N-gram minimum length.
|
||||
|
||||
***
|
||||
|
||||
### prefixOnly?
|
||||
|
||||
```ts
|
||||
optional prefixOnly: boolean;
|
||||
```
|
||||
|
||||
Whether to only emit token prefixes for the n-gram tokenizer.
|
||||
|
||||
***
|
||||
|
||||
### removeStopWords?
|
||||
|
||||
```ts
|
||||
optional removeStopWords: boolean;
|
||||
```
|
||||
|
||||
Whether to remove stop words.
|
||||
|
||||
***
|
||||
|
||||
### stem?
|
||||
|
||||
```ts
|
||||
optional stem: boolean;
|
||||
```
|
||||
|
||||
Whether to stem tokens.
|
||||
@@ -0,0 +1,19 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BaseTokenizer
|
||||
|
||||
# Type Alias: BaseTokenizer
|
||||
|
||||
```ts
|
||||
type BaseTokenizer:
|
||||
| "simple"
|
||||
| "whitespace"
|
||||
| "raw"
|
||||
| "ngram"
|
||||
| "icu"
|
||||
| "icu/split"
|
||||
| `jieba/${string}`
|
||||
| `lindera/${string}`;
|
||||
```
|
||||
@@ -0,0 +1,11 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenizeTableOptions
|
||||
|
||||
# Type Alias: TokenizeTableOptions
|
||||
|
||||
```ts
|
||||
type TokenizeTableOptions: object | object;
|
||||
```
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.31.0-beta.2</version>
|
||||
<version>0.32.0-beta.1</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.31.0-beta.2</version>
|
||||
<version>0.32.0-beta.1</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>9.0.0-beta.8</lance-core.version>
|
||||
<lance-core.version>9.0.0-beta.23</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>
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.31.0-beta.2"
|
||||
version = "0.32.0-beta.1"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
@@ -44,6 +44,6 @@ aws-lc-rs = "=1.16.3"
|
||||
napi-build = "2.3.1"
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
MeterProvider,
|
||||
type MetricData,
|
||||
MetricReader,
|
||||
} from "@opentelemetry/sdk-metrics";
|
||||
import * as tmp from "tmp";
|
||||
import { connect, instrumentLanceDbMetrics } from "../lancedb";
|
||||
// snapshotLancedbMetrics is internal plumbing (not part of the public API), so
|
||||
// it is imported from the native module rather than the package entry point.
|
||||
import { snapshotLancedbMetrics } from "../lancedb/native";
|
||||
|
||||
// The metrics recorder is process-global and installed once, so the whole
|
||||
// bridge is exercised in a single test to avoid cross-test global-state coupling.
|
||||
|
||||
// A minimal pull-based reader whose `collect()` we drive directly, invoking the
|
||||
// observable-instrument callbacks. `@opentelemetry/sdk-metrics` ships no
|
||||
// in-memory reader, so we subclass the abstract base.
|
||||
class TestMetricReader extends MetricReader {
|
||||
protected async onForceFlush(): Promise<void> {
|
||||
// no-op: collection is driven directly via collect()
|
||||
}
|
||||
protected async onShutdown(): Promise<void> {
|
||||
// no-op: nothing to release
|
||||
}
|
||||
}
|
||||
|
||||
async function metricsByName(
|
||||
reader: TestMetricReader,
|
||||
): Promise<Map<string, MetricData>> {
|
||||
const collected = await reader.collect();
|
||||
const result = new Map<string, MetricData>();
|
||||
for (const scope of collected.resourceMetrics.scopeMetrics) {
|
||||
for (const metric of scope.metrics) {
|
||||
result.set(metric.descriptor.name, metric);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
describe("OpenTelemetry metrics bridge", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("snapshot is safe to call regardless of install state", () => {
|
||||
expect(Array.isArray(snapshotLancedbMetrics())).toBe(true);
|
||||
});
|
||||
|
||||
it("exports object store metrics via observable instruments", async () => {
|
||||
const reader = new TestMetricReader();
|
||||
const provider = new MeterProvider({ readers: [reader] });
|
||||
expect(instrumentLanceDbMetrics(provider)).toBe(true);
|
||||
|
||||
// Generate object store activity on the local filesystem (scheme "file").
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = Array.from({ length: 256 }, (_, i) => ({ id: i }));
|
||||
const table = await db.createTable("t", data);
|
||||
expect(await table.countRows()).toBe(256);
|
||||
|
||||
const metrics = await metricsByName(reader);
|
||||
|
||||
const requests = metrics.get("lance_object_store_requests_total");
|
||||
expect(requests).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const requestPoints = (requests!.dataPoints as any[]) ?? [];
|
||||
expect(requestPoints.length).toBeGreaterThan(0);
|
||||
for (const p of requestPoints) {
|
||||
// Labelled by `operation` and `base` (the store scheme by default).
|
||||
expect(p.attributes).toHaveProperty("base");
|
||||
expect(p.attributes).toHaveProperty("operation");
|
||||
}
|
||||
const totalRequests = requestPoints.reduce((acc, p) => acc + p.value, 0);
|
||||
expect(totalRequests).toBeGreaterThan(0);
|
||||
|
||||
// Histograms are decomposed into bucket / count / sum observable counters.
|
||||
const bucket = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_bucket",
|
||||
);
|
||||
expect(bucket).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const bucketPoints = (bucket!.dataPoints as any[]) ?? [];
|
||||
expect(bucketPoints.length).toBeGreaterThan(0);
|
||||
expect(bucketPoints.every((p) => "le" in p.attributes)).toBe(true);
|
||||
// The implicit +Inf bucket must be present.
|
||||
expect(bucketPoints.some((p) => p.attributes.le === "+Inf")).toBe(true);
|
||||
|
||||
const count = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_count",
|
||||
);
|
||||
expect(count).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const countPoints = (count!.dataPoints as any[]) ?? [];
|
||||
expect(countPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
const sum = metrics.get("lance_object_store_request_duration_seconds_sum");
|
||||
expect(sum).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const sumPoints = (sum!.dataPoints as any[]) ?? [];
|
||||
expect(sumPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
// Unit handling: only `_sum` keeps the histogram's unit (seconds); `_bucket`
|
||||
// and `_count` observe cumulative counts and are unitless.
|
||||
expect(sum!.descriptor.unit).toBe("s");
|
||||
expect(bucket!.descriptor.unit).toBe("");
|
||||
expect(count!.descriptor.unit).toBe("");
|
||||
|
||||
await provider.shutdown();
|
||||
});
|
||||
});
|
||||
@@ -215,6 +215,20 @@ describe("Query orderBy", () => {
|
||||
expect(results[2].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should combine repeated where clauses with AND", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 1.0")
|
||||
.where("score < 3.0")
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
// Only rows matching both predicates should be returned, rather than the
|
||||
// second where() silently replacing the first.
|
||||
expect(results.length).toBe(2);
|
||||
expect(results[0].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(2.8, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with limit", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
|
||||
@@ -16,6 +16,7 @@ import {
|
||||
PhraseQuery,
|
||||
Table,
|
||||
connect,
|
||||
tokenize,
|
||||
} from "../lancedb";
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
@@ -2307,6 +2308,75 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
expect(results2[0].text).toBe(data[1].text);
|
||||
});
|
||||
|
||||
test("tokenizes FTS queries by column or index name", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
{
|
||||
text: "Running in cafés",
|
||||
japanese: "Hello, こんにちは世界!",
|
||||
vector: [0.1, 0.2, 0.3],
|
||||
},
|
||||
];
|
||||
const table = await db.createTable("test", data);
|
||||
await table.createIndex("text", {
|
||||
config: Index.fts({ baseTokenizer: "simple" }),
|
||||
});
|
||||
await table.createIndex("japanese", {
|
||||
config: Index.fts({
|
||||
baseTokenizer: "icu",
|
||||
stem: false,
|
||||
removeStopWords: false,
|
||||
}),
|
||||
name: "japanese_icu_idx",
|
||||
});
|
||||
|
||||
await expect(table.tokenize("hello", {} as never)).rejects.toThrow(
|
||||
"Specify exactly one",
|
||||
);
|
||||
await expect(
|
||||
table.tokenize("hello", {
|
||||
column: "text",
|
||||
indexName: "text_idx",
|
||||
} as never),
|
||||
).rejects.toThrow("Specify exactly one");
|
||||
|
||||
const simpleTokens = await table.tokenize("Running in cafés", {
|
||||
column: "text",
|
||||
});
|
||||
expect(simpleTokens).toEqual([
|
||||
{ text: "run", position: 0 },
|
||||
{ text: "cafe", position: 2 },
|
||||
]);
|
||||
|
||||
const icuTokens = await table.tokenize("Hello, こんにちは世界!", {
|
||||
indexName: "japanese_icu_idx",
|
||||
});
|
||||
expect(icuTokens).toEqual([
|
||||
{ text: "hello", position: 0 },
|
||||
{ text: "こんにちは", position: 1 },
|
||||
{ text: "世界", position: 2 },
|
||||
]);
|
||||
|
||||
const directSimpleTokens = await tokenize("Running in cafés", {
|
||||
baseTokenizer: "simple",
|
||||
});
|
||||
expect(directSimpleTokens).toEqual([
|
||||
{ text: "run", position: 0 },
|
||||
{ text: "cafe", position: 2 },
|
||||
]);
|
||||
|
||||
const directIcuTokens = await tokenize("Hello, こんにちは世界!", {
|
||||
baseTokenizer: "icu",
|
||||
stem: false,
|
||||
removeStopWords: false,
|
||||
});
|
||||
expect(directIcuTokens).toEqual([
|
||||
{ text: "hello", position: 0 },
|
||||
{ text: "こんにちは", position: 1 },
|
||||
{ text: "世界", position: 2 },
|
||||
]);
|
||||
});
|
||||
|
||||
test("full text search fast search", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [{ text: "hello world", vector: [0.1, 0.2, 0.3], id: 1 }];
|
||||
@@ -2992,6 +3062,56 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => {
|
||||
}),
|
||||
).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("reads back the installed spec via getLsmWriteSpec", async () => {
|
||||
const conn = await connect(tmpDir.name);
|
||||
const table = await makeTable(conn);
|
||||
await table.setUnenforcedPrimaryKey("id");
|
||||
|
||||
// Nothing installed yet.
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// A real scalar index is needed to name it as a maintained index.
|
||||
await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]);
|
||||
await table.createIndex("id");
|
||||
const indexName = (await table.listIndices())[0].name;
|
||||
|
||||
// Bucket spec round-trips, including maintained indexes and writer config
|
||||
// defaults. Lance writer-config keys are canonically snake_case.
|
||||
// biome-ignore lint/style/useNamingConvention: Lance writer-config keys are snake_case
|
||||
const writerConfigDefaults = { durable_write: "false" };
|
||||
await table.setLsmWriteSpec({
|
||||
specType: "bucket",
|
||||
column: "id",
|
||||
numBuckets: 4,
|
||||
maintainedIndexes: [indexName],
|
||||
writerConfigDefaults,
|
||||
});
|
||||
const spec = await table.getLsmWriteSpec();
|
||||
expect(spec).toBeDefined();
|
||||
expect(spec?.specType).toBe("bucket");
|
||||
expect(spec?.column).toBe("id");
|
||||
expect(spec?.numBuckets).toBe(4);
|
||||
expect(spec?.maintainedIndexes).toEqual([indexName]);
|
||||
expect(spec?.writerConfigDefaults).toEqual(writerConfigDefaults);
|
||||
|
||||
// After unset, undefined again.
|
||||
await table.unsetLsmWriteSpec();
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// Identity round-trips (column recovered from the schema).
|
||||
await table.setLsmWriteSpec({ specType: "identity", column: "id" });
|
||||
const identity = await table.getLsmWriteSpec();
|
||||
expect(identity?.specType).toBe("identity");
|
||||
expect(identity?.column).toBe("id");
|
||||
await table.unsetLsmWriteSpec();
|
||||
|
||||
// Unsharded round-trips (no routing column).
|
||||
await table.setLsmWriteSpec({ specType: "unsharded" });
|
||||
const unsharded = await table.getLsmWriteSpec();
|
||||
expect(unsharded?.specType).toBe("unsharded");
|
||||
expect(unsharded?.column).toBeFalsy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("LSM merge insert", () => {
|
||||
|
||||
@@ -13,13 +13,21 @@ import {
|
||||
Connection as LanceDbConnection,
|
||||
JsHeaderProvider as NativeJsHeaderProvider,
|
||||
Session,
|
||||
tokenize as nativeTokenize,
|
||||
} from "./native.js";
|
||||
|
||||
import { HeaderProvider } from "./header";
|
||||
import type { BaseTokenizer } from "./indices";
|
||||
import type { FtsToken } from "./table";
|
||||
|
||||
// Re-export native header provider for use with connectWithHeaderProvider
|
||||
export { JsHeaderProvider as NativeJsHeaderProvider } from "./native.js";
|
||||
|
||||
// OpenTelemetry metrics bridge. Only the high-level entry point is public; the
|
||||
// underlying recorder/catalog/snapshot functions remain internal plumbing that
|
||||
// `otel.ts` consumes from the native module.
|
||||
export { instrumentLanceDbMetrics } from "./otel";
|
||||
|
||||
export {
|
||||
AddColumnsSql,
|
||||
ConnectionOptions,
|
||||
@@ -52,6 +60,7 @@ export {
|
||||
SplitHashOptions,
|
||||
SplitSequentialOptions,
|
||||
ShuffleOptions,
|
||||
OAuthConfig as NativeOAuthConfig,
|
||||
} from "./native.js";
|
||||
|
||||
export {
|
||||
@@ -108,6 +117,7 @@ export {
|
||||
HnswPqOptions,
|
||||
HnswSqOptions,
|
||||
FtsOptions,
|
||||
BaseTokenizer,
|
||||
} from "./indices";
|
||||
|
||||
export {
|
||||
@@ -118,6 +128,8 @@ export {
|
||||
OptimizeOptions,
|
||||
Version,
|
||||
WriteProgress,
|
||||
FtsToken,
|
||||
TokenizeTableOptions,
|
||||
LsmWriteSpec,
|
||||
ColumnAlteration,
|
||||
FieldMetadataUpdate,
|
||||
@@ -130,6 +142,8 @@ export {
|
||||
TokenResponse,
|
||||
} from "./header";
|
||||
|
||||
export { OAuthConfig, OAuthFlowType } from "./oauth";
|
||||
|
||||
export { MergeInsertBuilder, WriteExecutionOptions } from "./merge";
|
||||
|
||||
export * as embedding from "./embedding";
|
||||
@@ -147,6 +161,68 @@ export {
|
||||
} from "./arrow";
|
||||
export { IntoSql, packBits } from "./util";
|
||||
|
||||
/**
|
||||
* Options for tokenizing a full-text search query without a table index.
|
||||
*/
|
||||
export interface TokenizeOptions {
|
||||
/**
|
||||
* The tokenizer to use. The default is "simple".
|
||||
*/
|
||||
baseTokenizer?: BaseTokenizer;
|
||||
|
||||
/** Language for stemming and stop words. */
|
||||
language?: string;
|
||||
|
||||
/** Maximum token length; tokens longer than this are ignored. */
|
||||
maxTokenLength?: number;
|
||||
|
||||
/** Whether to lowercase tokens. */
|
||||
lowercase?: boolean;
|
||||
|
||||
/** Whether to stem tokens. */
|
||||
stem?: boolean;
|
||||
|
||||
/** Whether to remove stop words. */
|
||||
removeStopWords?: boolean;
|
||||
|
||||
/** Whether to fold ASCII characters. */
|
||||
asciiFolding?: boolean;
|
||||
|
||||
/** N-gram minimum length. */
|
||||
ngramMinLength?: number;
|
||||
|
||||
/** N-gram maximum length. */
|
||||
ngramMaxLength?: number;
|
||||
|
||||
/** Whether to only emit token prefixes for the n-gram tokenizer. */
|
||||
prefixOnly?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize a full-text search query using an explicit tokenizer.
|
||||
*
|
||||
* This does not require a table or FTS index. The tokenizer options match
|
||||
* {@link Index.fts}.
|
||||
*/
|
||||
export async function tokenize(
|
||||
query: string,
|
||||
options?: Partial<TokenizeOptions>,
|
||||
): Promise<FtsToken[]> {
|
||||
return await nativeTokenize(
|
||||
query,
|
||||
options?.baseTokenizer,
|
||||
options?.language,
|
||||
options?.maxTokenLength,
|
||||
options?.lowercase,
|
||||
options?.stem,
|
||||
options?.removeStopWords,
|
||||
options?.asciiFolding,
|
||||
options?.ngramMinLength,
|
||||
options?.ngramMaxLength,
|
||||
options?.prefixOnly,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Connect to a LanceDB instance at the given URI.
|
||||
*
|
||||
|
||||
@@ -486,6 +486,16 @@ export interface IvfFlatOptions {
|
||||
sampleRate?: number;
|
||||
}
|
||||
|
||||
export type BaseTokenizer =
|
||||
| "simple"
|
||||
| "whitespace"
|
||||
| "raw"
|
||||
| "ngram"
|
||||
| "icu"
|
||||
| "icu/split"
|
||||
| `jieba/${string}`
|
||||
| `lindera/${string}`;
|
||||
|
||||
/**
|
||||
* Options to create a full text search index
|
||||
*/
|
||||
@@ -509,8 +519,12 @@ export interface FtsOptions {
|
||||
* "whitespace" - Whitespace tokenizer. This tokenizer splits the text into tokens using whitespace as a delimiter.
|
||||
*
|
||||
* "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token.
|
||||
*
|
||||
* "icu" - ICU dictionary-based word segmentation.
|
||||
*
|
||||
* "icu/split" - ICU segmentation with simple-style delimiter splitting.
|
||||
*/
|
||||
baseTokenizer?: "simple" | "whitespace" | "raw" | "ngram";
|
||||
baseTokenizer?: BaseTokenizer;
|
||||
|
||||
/**
|
||||
* language for stemming and stop words
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
/**
|
||||
* OAuth authentication flow types.
|
||||
*/
|
||||
export enum OAuthFlowType {
|
||||
/** Client Credentials grant (service-to-service / M2M). */
|
||||
ClientCredentials = "client_credentials",
|
||||
/** Azure Managed Identity via IMDS. */
|
||||
AzureManagedIdentity = "azure_managed_identity",
|
||||
}
|
||||
|
||||
/**
|
||||
* OAuth configuration for LanceDB authentication.
|
||||
*
|
||||
* This is the public TypeScript OAuth configuration type. The generated
|
||||
* `NativeOAuthConfig` type has the same runtime shape but is an implementation
|
||||
* detail of the napi-rs binding.
|
||||
*
|
||||
* All token acquisition and refresh is handled in the Rust layer.
|
||||
* This config is passed through to Rust via napi-rs.
|
||||
*
|
||||
* @example Client Credentials (service-to-service):
|
||||
* ```typescript
|
||||
* const config: OAuthConfig = {
|
||||
* issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
* clientId: "app-id",
|
||||
* clientSecret: "secret",
|
||||
* scopes: ["api://lancedb-api/.default"],
|
||||
* };
|
||||
* ```
|
||||
*
|
||||
* @example Azure Managed Identity:
|
||||
* ```typescript
|
||||
* const config: OAuthConfig = {
|
||||
* issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
* clientId: "app-id",
|
||||
* scopes: ["api://lancedb-api/.default"],
|
||||
* flow: OAuthFlowType.AzureManagedIdentity,
|
||||
* };
|
||||
* ```
|
||||
*/
|
||||
export interface OAuthConfig {
|
||||
/**
|
||||
* OIDC issuer URL or OAuth authority URL.
|
||||
* For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
|
||||
*/
|
||||
issuerUrl: string;
|
||||
|
||||
/** Application / Client ID. */
|
||||
clientId: string;
|
||||
|
||||
/**
|
||||
* OAuth scopes to request.
|
||||
* For Azure managed identity, exactly one scope or resource is required.
|
||||
* For example: `["api://{app_id}/.default"]`
|
||||
*/
|
||||
scopes: string[];
|
||||
|
||||
/** Authentication flow (default: ClientCredentials). */
|
||||
flow?: OAuthFlowType;
|
||||
|
||||
/** Client secret (required for ClientCredentials). */
|
||||
clientSecret?: string;
|
||||
|
||||
/** Client ID for user-assigned managed identity (AzureManagedIdentity). */
|
||||
managedIdentityClientId?: string;
|
||||
|
||||
/**
|
||||
* Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
* Keep this well below the token TTL; if it is greater than or equal to
|
||||
* the TTL, each request refreshes the token.
|
||||
*/
|
||||
refreshBufferSecs?: number;
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
type Attributes,
|
||||
type MeterProvider,
|
||||
type ObservableResult,
|
||||
metrics,
|
||||
} from "@opentelemetry/api";
|
||||
|
||||
import {
|
||||
lancedbMetricsCatalog,
|
||||
registerLancedbMetricsRecorder,
|
||||
snapshotLancedbMetrics,
|
||||
} from "./native";
|
||||
|
||||
let instrumented = false;
|
||||
|
||||
/**
|
||||
* Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
*
|
||||
* Installs a process-global metrics recorder and creates one observable
|
||||
* instrument per LanceDB metric (currently object store request counts, bytes,
|
||||
* latency, errors, and throttles) on the given (or global) `MeterProvider`. The
|
||||
* configured `MetricReader` then collects them on its own schedule.
|
||||
*
|
||||
* Counters and gauges map directly to observable counters/gauges. Because
|
||||
* OpenTelemetry has no asynchronous histogram instrument, each histogram is
|
||||
* exported Prometheus-style as cumulative `le` bucket counts (`<name>_bucket`,
|
||||
* with an `le` attribute) plus `<name>_count` and `<name>_sum`.
|
||||
*
|
||||
* Requires `@opentelemetry/api` (a dependency) and, to actually export, an
|
||||
* OpenTelemetry SDK such as `@opentelemetry/sdk-metrics`.
|
||||
*
|
||||
* @param meterProvider The provider to register instruments on. Defaults to the
|
||||
* global provider from `@opentelemetry/api`.
|
||||
* @returns `true` if the recorder is installed and instruments are registered.
|
||||
* `false` if a different `metrics` recorder is already installed in this
|
||||
* process (only one global recorder is permitted), in which case a warning is
|
||||
* emitted and no instruments are created. Calling this more than once is safe;
|
||||
* instruments are created only on the first successful call.
|
||||
*/
|
||||
export function instrumentLanceDbMetrics(
|
||||
meterProvider?: MeterProvider,
|
||||
): boolean {
|
||||
if (!registerLancedbMetricsRecorder()) {
|
||||
console.warn(
|
||||
"Could not install the LanceDB metrics recorder: another `metrics` " +
|
||||
"recorder is already installed in this process. LanceDB metrics will " +
|
||||
"not be exported via OpenTelemetry.",
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (instrumented) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const provider = meterProvider ?? metrics.getMeterProvider();
|
||||
const meter = provider.getMeter("lancedb");
|
||||
|
||||
const scalarCallback = (metricName: string) => (result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name === metricName && point.value != null) {
|
||||
result.observe(point.value, point.attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const bucketCallback = (metricName: string) => (result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name !== metricName || point.buckets == null) {
|
||||
continue;
|
||||
}
|
||||
for (const bucket of point.buckets) {
|
||||
const attributes: Attributes = {
|
||||
...point.attributes,
|
||||
le: bucket.le,
|
||||
};
|
||||
result.observe(bucket.cumulativeCount, attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const fieldCallback =
|
||||
(metricName: string, field: "count" | "sum") =>
|
||||
(result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name !== metricName) {
|
||||
continue;
|
||||
}
|
||||
const value = point[field];
|
||||
if (value != null) {
|
||||
result.observe(value, point.attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
for (const desc of lancedbMetricsCatalog()) {
|
||||
const unit = desc.unit ?? "";
|
||||
if (desc.kind === "counter") {
|
||||
const counter = meter.createObservableCounter(desc.name, {
|
||||
unit,
|
||||
description: desc.description,
|
||||
});
|
||||
counter.addCallback(scalarCallback(desc.name));
|
||||
} else if (desc.kind === "gauge") {
|
||||
const gauge = meter.createObservableGauge(desc.name, {
|
||||
unit,
|
||||
description: desc.description,
|
||||
});
|
||||
gauge.addCallback(scalarCallback(desc.name));
|
||||
} else if (desc.kind === "histogram") {
|
||||
// `_bucket` and `_count` observe cumulative sample counts, not the
|
||||
// histogram's measured quantity, so they are unitless; only `_sum`
|
||||
// carries the histogram's unit.
|
||||
const bucket = meter.createObservableCounter(`${desc.name}_bucket`, {
|
||||
description: `${desc.description} (cumulative buckets)`,
|
||||
});
|
||||
bucket.addCallback(bucketCallback(desc.name));
|
||||
|
||||
const count = meter.createObservableCounter(`${desc.name}_count`, {
|
||||
description: `${desc.description} (count)`,
|
||||
});
|
||||
count.addCallback(fieldCallback(desc.name, "count"));
|
||||
|
||||
const sum = meter.createObservableCounter(`${desc.name}_sum`, {
|
||||
unit,
|
||||
description: `${desc.description} (sum)`,
|
||||
});
|
||||
sum.addCallback(fieldCallback(desc.name, "sum"));
|
||||
}
|
||||
}
|
||||
|
||||
instrumented = true;
|
||||
return true;
|
||||
}
|
||||
@@ -362,6 +362,9 @@ export class StandardQueryBase<
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
where(predicate: string): this {
|
||||
this.doCall((inner: NativeQueryType) => inner.onlyIf(predicate));
|
||||
|
||||
@@ -158,6 +158,26 @@ export interface Version {
|
||||
metadata: Record<string, string>;
|
||||
}
|
||||
|
||||
/** Token produced by the tokenizer configured on a full-text search index. */
|
||||
export interface FtsToken {
|
||||
/** Token text after tokenizer filters have been applied. */
|
||||
text: string;
|
||||
/** Token position used by full-text query matching. */
|
||||
position: number;
|
||||
}
|
||||
|
||||
export type TokenizeTableOptions =
|
||||
| {
|
||||
/** FTS-indexed column whose tokenizer should be used. */
|
||||
column: string;
|
||||
indexName?: never;
|
||||
}
|
||||
| {
|
||||
/** Name of the FTS index whose tokenizer should be used. */
|
||||
indexName: string;
|
||||
column?: never;
|
||||
};
|
||||
|
||||
/**
|
||||
* Specification selecting Lance's MemWAL LSM-style write path for
|
||||
* `mergeInsert`.
|
||||
@@ -585,6 +605,17 @@ export abstract class Table {
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
abstract unsetLsmWriteSpec(): Promise<void>;
|
||||
/**
|
||||
* Read the {@link LsmWriteSpec} currently installed on this table.
|
||||
*
|
||||
* Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
* spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}).
|
||||
* The returned spec — including its `maintainedIndexes` and
|
||||
* `writerConfigDefaults` — mirrors what was passed to
|
||||
* {@link Table#setLsmWriteSpec}.
|
||||
* @returns {Promise<LsmWriteSpec | undefined>}
|
||||
*/
|
||||
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
|
||||
/**
|
||||
* Drain and close any cached MemWAL shard writers held for this table.
|
||||
*
|
||||
@@ -705,6 +736,19 @@ export abstract class Table {
|
||||
abstract optimize(options?: Partial<OptimizeOptions>): Promise<OptimizeStats>;
|
||||
/** List all indices that have been created with {@link Table.createIndex} */
|
||||
abstract listIndices(): Promise<IndexConfig[]>;
|
||||
/**
|
||||
* Tokenize a full-text search query using the tokenizer configured on an FTS index.
|
||||
*
|
||||
* Specify exactly one of `column` or `indexName`.
|
||||
*
|
||||
* Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
* the client process from index metadata. For remote tables, this means the
|
||||
* same tokenizer model files must also exist locally.
|
||||
*/
|
||||
abstract tokenize(
|
||||
query: string,
|
||||
options: TokenizeTableOptions,
|
||||
): Promise<FtsToken[]>;
|
||||
/** Return the table as an arrow table */
|
||||
abstract toArrow(): Promise<ArrowTable>;
|
||||
|
||||
@@ -1091,6 +1135,15 @@ export class LocalTable extends Table {
|
||||
return await this.inner.unsetLsmWriteSpec();
|
||||
}
|
||||
|
||||
async getLsmWriteSpec(): Promise<LsmWriteSpec | undefined> {
|
||||
// The native binding types `specType` as a plain `string`; narrow it back
|
||||
// to the public union. The Rust `From` impl only ever emits one of the
|
||||
// three valid values, so the cast is safe.
|
||||
return ((await this.inner.getLsmWriteSpec()) ?? undefined) as
|
||||
| LsmWriteSpec
|
||||
| undefined;
|
||||
}
|
||||
|
||||
async closeLsmWriters(): Promise<void> {
|
||||
return await this.inner.closeLsmWriters();
|
||||
}
|
||||
@@ -1153,6 +1206,17 @@ export class LocalTable extends Table {
|
||||
return await this.inner.listIndices();
|
||||
}
|
||||
|
||||
async tokenize(
|
||||
query: string,
|
||||
options: TokenizeTableOptions,
|
||||
): Promise<FtsToken[]> {
|
||||
return await this.inner.tokenize(
|
||||
query,
|
||||
options?.column,
|
||||
options?.indexName,
|
||||
);
|
||||
}
|
||||
|
||||
async toArrow(): Promise<ArrowTable> {
|
||||
return await this.query().toArrow();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+73
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
@@ -18,6 +18,7 @@
|
||||
"win32"
|
||||
],
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"devDependencies": {
|
||||
@@ -27,6 +28,7 @@
|
||||
"@biomejs/biome": "^1.7.3",
|
||||
"@jest/globals": "^29.7.0",
|
||||
"@napi-rs/cli": "3.7.0",
|
||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
||||
"@types/axios": "^0.14.0",
|
||||
"@types/jest": "^29.1.2",
|
||||
"@types/node": "22.7.4",
|
||||
@@ -4148,6 +4150,75 @@
|
||||
"@octokit/openapi-types": "^27.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/api": {
|
||||
"version": "1.9.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.1.tgz",
|
||||
"integrity": "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==",
|
||||
"license": "Apache-2.0",
|
||||
"engines": {
|
||||
"node": ">=8.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/core": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-1.30.1.tgz",
|
||||
"integrity": "sha512-OOCM2C/QIURhJMuKaekP3TRBxBKxG/TWWA0TL2J6nXUtDnuCtccy49LUJF8xPFXMX+0LMcxFpCo8M9cGY1W6rQ==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/semantic-conventions": "1.28.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.0.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/resources": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-1.30.1.tgz",
|
||||
"integrity": "sha512-5UxZqiAgLYGFjS4s9qm5mBVo433u+dSPUFWVWXmLAD4wB65oMCoXaJP1KJa9DIYYMeHu3z4BZcStG3LC593cWA==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/core": "1.30.1",
|
||||
"@opentelemetry/semantic-conventions": "1.28.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.0.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/sdk-metrics": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-1.30.1.tgz",
|
||||
"integrity": "sha512-q9zcZ0Okl8jRgmy7eNW3Ku1XSgg3sDLa5evHZpCwjspw7E8Is4K/haRPDJrBcX3YSn/Y7gUvFnByNYEKQNbNog==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/core": "1.30.1",
|
||||
"@opentelemetry/resources": "1.30.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.3.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/semantic-conventions": {
|
||||
"version": "1.28.0",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.28.0.tgz",
|
||||
"integrity": "sha512-lp4qAiMTD4sNWW4DbKLBkfiMZ4jbAboJIGOQr5DvciMRI494OapieI9qiODpOt0XBr1LjIDy1xAGAnVs5supTA==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
}
|
||||
},
|
||||
"node_modules/@protobufjs/aspromise": {
|
||||
"version": "1.1.2",
|
||||
"resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz",
|
||||
|
||||
+3
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.31.0-beta.2",
|
||||
"version": "0.32.0-beta.1",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
@@ -44,6 +44,7 @@
|
||||
"@biomejs/biome": "^1.7.3",
|
||||
"@jest/globals": "^29.7.0",
|
||||
"@napi-rs/cli": "3.7.0",
|
||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
||||
"@types/axios": "^0.14.0",
|
||||
"@types/jest": "^29.1.2",
|
||||
"@types/node": "22.7.4",
|
||||
@@ -92,6 +93,7 @@
|
||||
"version": "napi version"
|
||||
},
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
|
||||
Generated
+53
@@ -8,6 +8,9 @@ importers:
|
||||
|
||||
.:
|
||||
dependencies:
|
||||
'@opentelemetry/api':
|
||||
specifier: ^1.9.0
|
||||
version: 1.9.1
|
||||
apache-arrow:
|
||||
specifier: '>=15.0.0 <=18.1.0'
|
||||
version: 18.1.0
|
||||
@@ -33,6 +36,9 @@ importers:
|
||||
'@napi-rs/cli':
|
||||
specifier: 3.7.0
|
||||
version: 3.7.0(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)(@types/node@22.7.4)
|
||||
'@opentelemetry/sdk-metrics':
|
||||
specifier: ^1.30.0
|
||||
version: 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@types/axios':
|
||||
specifier: ^0.14.0
|
||||
version: 0.14.4
|
||||
@@ -1307,6 +1313,32 @@ packages:
|
||||
'@octokit/types@16.0.0':
|
||||
resolution: {integrity: sha512-sKq+9r1Mm4efXW1FCk7hFSeJo4QKreL/tTbR0rz/qx/r1Oa2VV83LTA/H/MuCOX7uCIJmQVRKBcbmWoySjAnSg==}
|
||||
|
||||
'@opentelemetry/api@1.9.1':
|
||||
resolution: {integrity: sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==}
|
||||
engines: {node: '>=8.0.0'}
|
||||
|
||||
'@opentelemetry/core@1.30.1':
|
||||
resolution: {integrity: sha512-OOCM2C/QIURhJMuKaekP3TRBxBKxG/TWWA0TL2J6nXUtDnuCtccy49LUJF8xPFXMX+0LMcxFpCo8M9cGY1W6rQ==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.0.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/resources@1.30.1':
|
||||
resolution: {integrity: sha512-5UxZqiAgLYGFjS4s9qm5mBVo433u+dSPUFWVWXmLAD4wB65oMCoXaJP1KJa9DIYYMeHu3z4BZcStG3LC593cWA==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.0.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/sdk-metrics@1.30.1':
|
||||
resolution: {integrity: sha512-q9zcZ0Okl8jRgmy7eNW3Ku1XSgg3sDLa5evHZpCwjspw7E8Is4K/haRPDJrBcX3YSn/Y7gUvFnByNYEKQNbNog==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.3.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.28.0':
|
||||
resolution: {integrity: sha512-lp4qAiMTD4sNWW4DbKLBkfiMZ4jbAboJIGOQr5DvciMRI494OapieI9qiODpOt0XBr1LjIDy1xAGAnVs5supTA==}
|
||||
engines: {node: '>=14'}
|
||||
|
||||
'@protobufjs/aspromise@1.1.2':
|
||||
resolution: {integrity: sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==}
|
||||
|
||||
@@ -4925,6 +4957,27 @@ snapshots:
|
||||
dependencies:
|
||||
'@octokit/openapi-types': 27.0.0
|
||||
|
||||
'@opentelemetry/api@1.9.1': {}
|
||||
|
||||
'@opentelemetry/core@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/semantic-conventions': 1.28.0
|
||||
|
||||
'@opentelemetry/resources@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/core': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@opentelemetry/semantic-conventions': 1.28.0
|
||||
|
||||
'@opentelemetry/sdk-metrics@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/core': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@opentelemetry/resources': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.28.0': {}
|
||||
|
||||
'@protobufjs/aspromise@1.1.2':
|
||||
optional: true
|
||||
|
||||
|
||||
@@ -112,6 +112,12 @@ impl Connection {
|
||||
|
||||
builder = builder.client_config(rust_config);
|
||||
|
||||
if let Some(oauth_config) = options.oauth_config {
|
||||
let config: lancedb::remote::oauth::OAuthConfig =
|
||||
oauth_config.try_into().default_error()?;
|
||||
builder = builder.oauth_config(config);
|
||||
}
|
||||
|
||||
if let Some(api_key) = options.api_key {
|
||||
builder = builder.api_key(&api_key);
|
||||
}
|
||||
|
||||
@@ -9,8 +9,11 @@ use lancedb::index::vector::{
|
||||
IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
|
||||
IvfRqIndexBuilder,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use napi_derive::napi;
|
||||
|
||||
use crate::error::NapiErrorExt;
|
||||
use crate::table::FtsToken;
|
||||
use crate::util::parse_distance_type;
|
||||
|
||||
#[napi]
|
||||
@@ -30,6 +33,65 @@ impl Index {
|
||||
}
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
#[allow(dead_code, clippy::too_many_arguments)]
|
||||
pub fn tokenize(
|
||||
query: String,
|
||||
base_tokenizer: Option<String>,
|
||||
language: Option<String>,
|
||||
max_token_length: Option<u32>,
|
||||
lower_case: Option<bool>,
|
||||
stem: Option<bool>,
|
||||
remove_stop_words: Option<bool>,
|
||||
ascii_folding: Option<bool>,
|
||||
ngram_min_length: Option<u32>,
|
||||
ngram_max_length: Option<u32>,
|
||||
prefix_only: Option<bool>,
|
||||
) -> napi::Result<Vec<FtsToken>> {
|
||||
let mut opts = FtsIndexBuilder::default();
|
||||
if let Some(base_tokenizer) = base_tokenizer {
|
||||
opts = opts.base_tokenizer(base_tokenizer);
|
||||
}
|
||||
if let Some(language) = language {
|
||||
opts = opts.language(&language).map_err(|_| {
|
||||
napi::Error::from_reason(format!(
|
||||
"LanceDB does not support the requested language: '{}'",
|
||||
language
|
||||
))
|
||||
})?;
|
||||
}
|
||||
if let Some(max_token_length) = max_token_length {
|
||||
opts = opts.max_token_length(Some(max_token_length as usize));
|
||||
}
|
||||
if let Some(lower_case) = lower_case {
|
||||
opts = opts.lower_case(lower_case);
|
||||
}
|
||||
if let Some(stem) = stem {
|
||||
opts = opts.stem(stem);
|
||||
}
|
||||
if let Some(remove_stop_words) = remove_stop_words {
|
||||
opts = opts.remove_stop_words(remove_stop_words);
|
||||
}
|
||||
if let Some(ascii_folding) = ascii_folding {
|
||||
opts = opts.ascii_folding(ascii_folding);
|
||||
}
|
||||
if let Some(ngram_min_length) = ngram_min_length {
|
||||
opts = opts.ngram_min_length(ngram_min_length);
|
||||
}
|
||||
if let Some(ngram_max_length) = ngram_max_length {
|
||||
opts = opts.ngram_max_length(ngram_max_length);
|
||||
}
|
||||
if let Some(prefix_only) = prefix_only {
|
||||
opts = opts.ngram_prefix_only(prefix_only);
|
||||
}
|
||||
|
||||
Ok(lancedb_tokenize(&query, &opts)
|
||||
.default_error()?
|
||||
.into_iter()
|
||||
.map(FtsToken::from)
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl Index {
|
||||
#[napi(factory)]
|
||||
|
||||
@@ -12,6 +12,7 @@ mod header;
|
||||
mod index;
|
||||
mod iterator;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
mod query;
|
||||
pub mod remote;
|
||||
@@ -65,6 +66,11 @@ pub struct ConnectionOptions {
|
||||
/// (For LanceDB cloud only): the host to use for LanceDB cloud. Used
|
||||
/// for testing purposes.
|
||||
pub host_override: Option<String>,
|
||||
/// (For LanceDB cloud only): OAuth configuration for IdP-based
|
||||
/// authentication (e.g., Azure Entra ID). When set, token acquisition
|
||||
/// and refresh are handled entirely in Rust. TypeScript users should pass
|
||||
/// the public `OAuthConfig` type exported from `@lancedb/lancedb`.
|
||||
pub oauth_config: Option<remote::OAuthConfig>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
|
||||
+4
-6
@@ -3,7 +3,7 @@
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::{arrow::IntoArrow, ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder};
|
||||
use lancedb::{ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi_derive::napi;
|
||||
|
||||
@@ -66,11 +66,9 @@ impl NativeMergeInsertBuilder {
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn execute(&self, buf: Buffer) -> napi::Result<MergeResult> {
|
||||
let data = ipc_file_to_batches(buf.to_vec())
|
||||
.and_then(IntoArrow::into_arrow)
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e)))
|
||||
})?;
|
||||
let data = ipc_file_to_batches(buf.to_vec()).map_err(|e| {
|
||||
napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e)))
|
||||
})?;
|
||||
|
||||
let this = self.clone();
|
||||
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Node.js bindings over [`lancedb::metrics_otel`].
|
||||
//!
|
||||
//! The aggregation, catalog, and histogram bucketing all live in the LanceDB
|
||||
//! core crate; this module only converts the core snapshot types into napi
|
||||
//! objects and exposes the three entry points to JavaScript, where
|
||||
//! `lancedb/otel.ts` bridges them into the user's OpenTelemetry `MeterProvider`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use lancedb::metrics_otel::{MetricPoint as CoreMetricPoint, MetricValue};
|
||||
use napi_derive::napi;
|
||||
|
||||
/// One cumulative histogram bucket: all samples with value `<= le`.
|
||||
#[napi(object)]
|
||||
pub struct MetricBucket {
|
||||
/// The inclusive upper bound of the bucket, or `"+Inf"` for the final bucket.
|
||||
pub le: String,
|
||||
/// Cumulative number of samples less than or equal to `le`.
|
||||
pub cumulative_count: f64,
|
||||
}
|
||||
|
||||
/// One aggregated metric data point. For counters and gauges only `value` is
|
||||
/// set; for histograms `buckets` (cumulative `le` counts), `count`, and `sum`
|
||||
/// are set.
|
||||
#[napi(object)]
|
||||
pub struct MetricPoint {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub attributes: HashMap<String, String>,
|
||||
pub value: Option<f64>,
|
||||
pub buckets: Option<Vec<MetricBucket>>,
|
||||
pub count: Option<f64>,
|
||||
pub sum: Option<f64>,
|
||||
}
|
||||
|
||||
impl From<CoreMetricPoint> for MetricPoint {
|
||||
fn from(point: CoreMetricPoint) -> Self {
|
||||
let kind = point.kind.as_str().to_string();
|
||||
let (value, buckets, count, sum) = match point.value {
|
||||
MetricValue::Scalar(v) => (Some(v), None, None, None),
|
||||
MetricValue::Histogram {
|
||||
buckets,
|
||||
count,
|
||||
sum,
|
||||
} => (
|
||||
None,
|
||||
Some(
|
||||
buckets
|
||||
.into_iter()
|
||||
// Counts stay well within the f64-exact integer range
|
||||
// (2^53), so this cast is lossless in practice and keeps
|
||||
// the values plain JS numbers for OpenTelemetry.
|
||||
.map(|(le, cumulative_count)| MetricBucket {
|
||||
le,
|
||||
cumulative_count: cumulative_count as f64,
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
Some(count as f64),
|
||||
Some(sum),
|
||||
),
|
||||
};
|
||||
Self {
|
||||
name: point.name,
|
||||
kind,
|
||||
attributes: point.attributes,
|
||||
value,
|
||||
buckets,
|
||||
count,
|
||||
sum,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A described metric, used by the JavaScript layer to create instruments up front.
|
||||
#[napi(object)]
|
||||
pub struct MetricDescription {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub unit: Option<String>,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// Install the LanceDB metrics recorder as the process-global `metrics` recorder.
|
||||
///
|
||||
/// Returns `true` if the recorder is installed (now or previously). Returns
|
||||
/// `false` if a *different* recorder is already installed — `metrics` allows
|
||||
/// only one global recorder per process, so LanceDB cannot coexist with another.
|
||||
#[napi]
|
||||
pub fn register_lancedb_metrics_recorder() -> bool {
|
||||
lancedb::metrics_otel::register_metrics_recorder()
|
||||
}
|
||||
|
||||
/// The catalog of described LanceDB metrics. Empty until the recorder is installed.
|
||||
#[napi]
|
||||
pub fn lancedb_metrics_catalog() -> Vec<MetricDescription> {
|
||||
lancedb::metrics_otel::metrics_catalog()
|
||||
.into_iter()
|
||||
.map(|desc| MetricDescription {
|
||||
name: desc.name,
|
||||
kind: desc.kind.as_str().to_string(),
|
||||
unit: desc.unit,
|
||||
description: desc.description,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A point-in-time snapshot of every recorded metric. Empty until the recorder
|
||||
/// is installed.
|
||||
#[napi]
|
||||
pub fn snapshot_lancedb_metrics() -> Vec<MetricPoint> {
|
||||
lancedb::metrics_otel::snapshot_metrics()
|
||||
.into_iter()
|
||||
.map(MetricPoint::from)
|
||||
.collect()
|
||||
}
|
||||
@@ -16,6 +16,7 @@ pub struct SplitRandomOptions {
|
||||
pub counts: Option<Vec<i64>>,
|
||||
pub fixed: Option<i64>,
|
||||
pub seed: Option<i64>,
|
||||
pub clump_size: Option<i64>,
|
||||
pub split_names: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -125,10 +126,15 @@ impl PermutationBuilder {
|
||||
};
|
||||
|
||||
let seed = options.seed.map(|s| s as u64);
|
||||
let clump_size = options.clump_size.map(|c| c as u64);
|
||||
|
||||
self.modify(|builder| {
|
||||
builder.with_split_strategy(
|
||||
SplitStrategy::Random { seed, sizes },
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
},
|
||||
options.split_names.clone(),
|
||||
)
|
||||
})
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use lancedb::error::Error;
|
||||
use napi_derive::*;
|
||||
|
||||
/// Timeout configuration for remote HTTP client.
|
||||
@@ -140,6 +141,84 @@ impl From<TlsConfig> for lancedb::remote::TlsConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// OAuth configuration for LanceDB authentication.
|
||||
///
|
||||
/// This is the generated napi-rs binding shape. TypeScript users should prefer
|
||||
/// the public `OAuthConfig` type exported from `@lancedb/lancedb`.
|
||||
///
|
||||
/// All token acquisition and refresh is handled in the Rust layer.
|
||||
#[napi(object)]
|
||||
#[derive(Clone)]
|
||||
pub struct OAuthConfig {
|
||||
/// OIDC issuer URL or OAuth authority URL.
|
||||
/// For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
|
||||
pub issuer_url: String,
|
||||
/// Application / Client ID.
|
||||
pub client_id: String,
|
||||
/// OAuth scopes to request. For Azure managed identity, exactly one scope
|
||||
/// or resource is required. For example: `["api://{app_id}/.default"]`
|
||||
pub scopes: Vec<String>,
|
||||
/// Authentication flow: "client_credentials" or "azure_managed_identity"
|
||||
pub flow: Option<String>,
|
||||
/// Client secret (required for client_credentials).
|
||||
pub client_secret: Option<String>,
|
||||
/// Client ID for user-assigned managed identity (azure_managed_identity).
|
||||
pub managed_identity_client_id: Option<String>,
|
||||
/// Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
/// Keep this well below the token TTL; if it is greater than or equal to
|
||||
/// the TTL, each request refreshes the token.
|
||||
pub refresh_buffer_secs: Option<u32>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for OAuthConfig {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("OAuthConfig")
|
||||
.field("issuer_url", &self.issuer_url)
|
||||
.field("client_id", &self.client_id)
|
||||
.field("scopes", &self.scopes)
|
||||
.field("flow", &self.flow)
|
||||
.field(
|
||||
"client_secret",
|
||||
&self.client_secret.as_deref().map(|_| "<redacted>"),
|
||||
)
|
||||
.field(
|
||||
"managed_identity_client_id",
|
||||
&self.managed_identity_client_id,
|
||||
)
|
||||
.field("refresh_buffer_secs", &self.refresh_buffer_secs)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(config: OAuthConfig) -> Result<Self, Self::Error> {
|
||||
use lancedb::remote::oauth::OAuthFlow;
|
||||
|
||||
let flow = match config.flow.as_deref().unwrap_or("client_credentials") {
|
||||
"client_credentials" => OAuthFlow::ClientCredentials,
|
||||
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
|
||||
client_id: config.managed_identity_client_id,
|
||||
},
|
||||
other => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("Unknown OAuth flow type: {other}"),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
issuer_url: config.issuer_url,
|
||||
client_id: config.client_id,
|
||||
client_secret: config.client_secret,
|
||||
scopes: config.scopes,
|
||||
flow,
|
||||
refresh_buffer_secs: config.refresh_buffer_secs.map(|v| v as u64),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ClientConfig> for lancedb::remote::ClientConfig {
|
||||
fn from(config: ClientConfig) -> Self {
|
||||
Self {
|
||||
@@ -156,3 +235,45 @@ impl From<ClientConfig> for lancedb::remote::ClientConfig {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_unknown_oauth_flow_returns_invalid_input() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("typo".to_string()),
|
||||
client_secret: None,
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
};
|
||||
|
||||
let err = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::InvalidInput { message }
|
||||
if message == "Unknown OAuth flow type: typo"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oauth_config_debug_redacts_client_secret() {
|
||||
let config = OAuthConfig {
|
||||
issuer_url: "https://issuer.example.com".to_string(),
|
||||
client_id: "client-id".to_string(),
|
||||
scopes: vec!["scope".to_string()],
|
||||
flow: Some("client_credentials".to_string()),
|
||||
client_secret: Some("super-secret".to_string()),
|
||||
managed_identity_client_id: None,
|
||||
refresh_buffer_secs: None,
|
||||
};
|
||||
|
||||
let debug = format!("{config:?}");
|
||||
assert!(!debug.contains("super-secret"));
|
||||
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
|
||||
}
|
||||
}
|
||||
|
||||
+92
-2
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
|
||||
use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema};
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration,
|
||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, NewColumnTransform, OptimizeAction,
|
||||
OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
|
||||
@@ -411,6 +411,16 @@ impl Table {
|
||||
.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn get_lsm_write_spec(&self) -> napi::Result<Option<LsmWriteSpec>> {
|
||||
let spec = self
|
||||
.inner_ref()?
|
||||
.get_lsm_write_spec()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn close_lsm_writers(&self) -> napi::Result<()> {
|
||||
self.inner_ref()?.close_lsm_writers().await.default_error()
|
||||
@@ -564,6 +574,27 @@ impl Table {
|
||||
.collect::<Vec<_>>())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn tokenize(
|
||||
&self,
|
||||
query: String,
|
||||
column: Option<String>,
|
||||
index_name: Option<String>,
|
||||
) -> napi::Result<Vec<FtsToken>> {
|
||||
let table = self.inner_ref()?;
|
||||
let tokens = match (column.as_deref(), index_name.as_deref()) {
|
||||
(Some(_), Some(_)) | (None, None) => {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Specify exactly one of 'column' or 'indexName'",
|
||||
));
|
||||
}
|
||||
(Some(column), None) => table.tokenize_with_column(&query, column).await,
|
||||
(None, Some(index_name)) => table.tokenize(&query, index_name).await,
|
||||
}
|
||||
.default_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn index_stats(&self, index_name: String) -> napi::Result<Option<IndexStatistics>> {
|
||||
let tbl = self.inner_ref()?;
|
||||
@@ -671,6 +702,24 @@ impl From<lancedb::index::IndexConfig> for IndexConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
/// A token produced by the tokenizer configured on a full-text search index.
|
||||
pub struct FtsToken {
|
||||
/// The token text after the index tokenizer has applied its filters.
|
||||
pub text: String,
|
||||
/// The token position used by full-text query matching.
|
||||
pub position: u32,
|
||||
}
|
||||
|
||||
impl From<LanceDbFtsToken> for FtsToken {
|
||||
fn from(token: LanceDbFtsToken) -> Self {
|
||||
Self {
|
||||
text: token.text,
|
||||
position: token.position,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Specification selecting Lance's MemWAL LSM-style write path for
|
||||
/// `mergeInsert`.
|
||||
///
|
||||
@@ -728,6 +777,47 @@ impl TryFrom<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(spec: lancedb::table::LsmWriteSpec) -> Self {
|
||||
use lancedb::table::LsmWriteSpec as Native;
|
||||
match spec {
|
||||
Native::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "bucket".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: Some(num_buckets),
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Identity {
|
||||
column,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "identity".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Unsharded {
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "unsharded".to_string(),
|
||||
column: None,
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.34.0-beta.3"
|
||||
current_version = "0.35.0-beta.2"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.34.0-beta.3"
|
||||
version = "0.35.0-beta.2"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -47,6 +47,6 @@ pyo3-build-config = { version = "0.28", features = [
|
||||
] }
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Benchmark for StreamingDataset throughput.
|
||||
|
||||
Sweeps read_batch_size from 1 to 16384 to show how amortising the per-request
|
||||
overhead scales. Each row at each chunk size is timed via the real
|
||||
StreamingDataset so the numbers reflect production code.
|
||||
|
||||
Run with:
|
||||
cd python
|
||||
uv run --extra tests benchmarks/bench_streaming_dataloader.py
|
||||
|
||||
Optional env vars:
|
||||
BENCH_NUM_ROWS — total rows in the table (default 49152 = 24 × 2048)
|
||||
BENCH_NUM_SPLITS — number of splits (default 24)
|
||||
BENCH_STEPS — round-robin cycles to time per chunk size (default 100)
|
||||
BENCH_ROW_BYTES — approximate bytes per row padded with a binary column
|
||||
(default 4096, mimics a small embedding/image patch)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
|
||||
from lancedb.streaming import StreamingDataset
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NUM_SPLITS = int(os.environ.get("BENCH_NUM_SPLITS", 24))
|
||||
# Default: 2048 rows per split so every chunk size up to 16Ki has ≥1 full
|
||||
# chunk (except 16Ki itself which gets a single full-split fetch — still valid).
|
||||
NUM_ROWS = int(os.environ.get("BENCH_NUM_ROWS", NUM_SPLITS * 2048))
|
||||
STEPS = int(os.environ.get("BENCH_STEPS", 100))
|
||||
ROW_BYTES = int(os.environ.get("BENCH_ROW_BYTES", 4096))
|
||||
|
||||
assert NUM_ROWS % NUM_SPLITS == 0, "NUM_ROWS must be divisible by NUM_SPLITS"
|
||||
|
||||
CHUNK_SIZES = [1, 4, 16, 64, 256, 1024, 4096, 16384]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Table helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_table(db_path: str) -> lancedb.table.Table:
|
||||
db = lancedb.connect(db_path)
|
||||
payload = b"x" * ROW_BYTES
|
||||
data = pa.table(
|
||||
{
|
||||
"id": pa.array(range(NUM_ROWS), type=pa.int32()),
|
||||
"payload": pa.array([payload] * NUM_ROWS, type=pa.large_binary()),
|
||||
}
|
||||
)
|
||||
return db.create_table("bench", data, mode="overwrite")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_chunk(table, chunk_size: int, steps: int) -> tuple[int, float]:
|
||||
"""Return (rows_drained, elapsed_seconds) for one timed run."""
|
||||
total_rows = steps * NUM_SPLITS
|
||||
ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk_size
|
||||
)
|
||||
count = 0
|
||||
t0 = time.perf_counter()
|
||||
for _ in ds:
|
||||
count += 1
|
||||
if count >= total_rows:
|
||||
break
|
||||
return count, time.perf_counter() - t0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows_per_split = NUM_ROWS // NUM_SPLITS
|
||||
print("Benchmark config:")
|
||||
print(
|
||||
f" NUM_ROWS={NUM_ROWS} NUM_SPLITS={NUM_SPLITS} "
|
||||
f"rows/split={rows_per_split} STEPS={STEPS} ROW_BYTES={ROW_BYTES}"
|
||||
)
|
||||
print(f" ~{NUM_ROWS * ROW_BYTES / 1024 / 1024:.1f} MB total table size")
|
||||
print()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
print("Creating table...", flush=True)
|
||||
table = make_table(tmp)
|
||||
|
||||
cols = (
|
||||
f"{'chunk':>6} {'rows':>6} {'elapsed':>8} {'rows/s':>10} {'ms/step':>9}"
|
||||
)
|
||||
print(f"\n{cols}")
|
||||
print("-" * 52)
|
||||
|
||||
for chunk in CHUNK_SIZES:
|
||||
# Warm-up pass (one step's worth of rows)
|
||||
warmup_ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk
|
||||
)
|
||||
warmup_count = 0
|
||||
for _ in warmup_ds:
|
||||
warmup_count += 1
|
||||
if warmup_count >= NUM_SPLITS:
|
||||
break
|
||||
|
||||
drained, elapsed = bench_chunk(table, chunk, STEPS)
|
||||
rows_per_sec = drained / elapsed if elapsed > 0 else float("inf")
|
||||
ms_per_step = elapsed / STEPS * 1000
|
||||
|
||||
print(
|
||||
f"{chunk:>6} {drained:>6} {elapsed:>7.3f}s "
|
||||
f"{rows_per_sec:>10.0f} {ms_per_step:>8.1f}ms"
|
||||
)
|
||||
|
||||
print()
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -47,6 +47,10 @@ repository = "https://github.com/lancedb/lancedb"
|
||||
pylance = [
|
||||
"pylance>=5.0.0b5",
|
||||
]
|
||||
# A library only needs the OpenTelemetry API; the application supplies and
|
||||
# configures the SDK (the actual exporter/reader). See
|
||||
# https://opentelemetry.io/docs/languages/python/instrumentation/
|
||||
otel = ["opentelemetry-api"]
|
||||
tests = [
|
||||
"aiohttp>=3.9.0",
|
||||
"boto3>=1.28.57",
|
||||
@@ -61,6 +65,7 @@ tests = [
|
||||
"pylance>=5.0.0b5",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=52,<53",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
]
|
||||
dev = [
|
||||
"ruff>=0.3.0",
|
||||
|
||||
@@ -6,19 +6,22 @@ import importlib.metadata
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional, Union, Any, List
|
||||
from typing import Dict, Optional, Union, Any, List, Iterable
|
||||
|
||||
__version__ = importlib.metadata.version("lancedb")
|
||||
|
||||
from ._lancedb import connect as lancedb_connect
|
||||
from ._lancedb import FtsToken
|
||||
from ._lancedb import tokenize as _tokenize
|
||||
from .common import URI, sanitize_uri
|
||||
from urllib.parse import urlparse
|
||||
from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||
from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .schema import vector
|
||||
from .schema import blob, vector, BlobType
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
from .namespace import (
|
||||
connect_namespace,
|
||||
@@ -89,6 +92,8 @@ def connect(
|
||||
If presented, connect to LanceDB cloud.
|
||||
Otherwise, connect to a database on file system or cloud storage.
|
||||
Can be set via environment variable `LANCEDB_API_KEY`.
|
||||
OAuth configuration is currently supported only by ``connect_async``;
|
||||
synchronous LanceDB Cloud connections require an API key.
|
||||
region: str, default "us-east-1"
|
||||
The region to use for LanceDB Cloud.
|
||||
host_override: str, optional
|
||||
@@ -147,8 +152,14 @@ def connect(
|
||||
|
||||
For object storage, use a URI prefix:
|
||||
|
||||
>>> db = lancedb.connect("s3://my-bucket/lancedb",
|
||||
... storage_options={"aws_access_key_id": "***"})
|
||||
>>> db = lancedb.connect( # doctest: +SKIP
|
||||
... "s3://my-bucket/lancedb",
|
||||
... storage_options={
|
||||
... "aws_access_key_id": "***",
|
||||
... "aws_secret_access_key": "***",
|
||||
... "aws_region": "us-east-1",
|
||||
... },
|
||||
... )
|
||||
|
||||
For tests and temporary data, use an in-memory database:
|
||||
|
||||
@@ -238,6 +249,40 @@ def connect(
|
||||
)
|
||||
|
||||
|
||||
def tokenize(
|
||||
query: str,
|
||||
*,
|
||||
base_tokenizer: BaseTokenizerType = "simple",
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
) -> Iterable[FtsToken]:
|
||||
"""Tokenize a full-text search query using an explicit tokenizer.
|
||||
|
||||
This does not require a table or FTS index. The tokenizer options match
|
||||
:class:`lancedb.index.FTS`.
|
||||
"""
|
||||
return _tokenize(
|
||||
query,
|
||||
base_tokenizer=base_tokenizer,
|
||||
language=language,
|
||||
max_token_length=max_token_length,
|
||||
lower_case=lower_case,
|
||||
stem=stem,
|
||||
remove_stop_words=remove_stop_words,
|
||||
ascii_folding=ascii_folding,
|
||||
ngram_min_length=ngram_min_length,
|
||||
ngram_max_length=ngram_max_length,
|
||||
prefix_only=prefix_only,
|
||||
)
|
||||
|
||||
|
||||
WORKER_PROPERTY_PREFIX = "_lancedb_worker_"
|
||||
|
||||
|
||||
@@ -340,6 +385,7 @@ async def connect_async(
|
||||
session: Optional[Session] = None,
|
||||
manifest_enabled: bool = False,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
oauth_config=None,
|
||||
) -> AsyncConnection:
|
||||
"""Connect to a LanceDB database.
|
||||
|
||||
@@ -389,6 +435,10 @@ async def connect_async(
|
||||
namespace_client_properties : dict, optional
|
||||
Additional directory namespace client properties to use with
|
||||
``manifest_enabled=True``.
|
||||
oauth_config : OAuthConfig, optional
|
||||
OAuth configuration for LanceDB Cloud/Enterprise. This is supported by
|
||||
``connect_async`` only; synchronous ``connect`` uses API key
|
||||
authentication for ``db://`` URIs.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -435,6 +485,7 @@ async def connect_async(
|
||||
session,
|
||||
manifest_enabled,
|
||||
namespace_client_properties,
|
||||
oauth_config,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -442,17 +493,21 @@ async def connect_async(
|
||||
__all__ = [
|
||||
"connect",
|
||||
"connect_async",
|
||||
"tokenize",
|
||||
"connect_namespace",
|
||||
"connect_namespace_async",
|
||||
"AsyncConnection",
|
||||
"AsyncLanceNamespaceDBConnection",
|
||||
"AsyncTable",
|
||||
"FtsToken",
|
||||
"col",
|
||||
"Expr",
|
||||
"func",
|
||||
"lit",
|
||||
"URI",
|
||||
"sanitize_uri",
|
||||
"blob",
|
||||
"BlobType",
|
||||
"vector",
|
||||
"DBConnection",
|
||||
"LanceDBConnection",
|
||||
|
||||
@@ -0,0 +1,420 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Blob fetch API and v2 projection helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from typing import TYPE_CHECKING, Optional, Union
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from .expr import Expr
|
||||
from .schema import blob_v2_column_paths
|
||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||
from .util import get_uri_scheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import WriteableBuffer
|
||||
|
||||
from .remote.table import RemoteTable
|
||||
from .table import AsyncTable, Table
|
||||
|
||||
BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"bytes": "all_binary",
|
||||
"descriptions": "blobs_descriptions",
|
||||
}
|
||||
|
||||
ROW_ID_FIELD_NAME = "_lance_row_id"
|
||||
|
||||
FetchBlobsSync = Callable[[str, pa.Table], pa.Array | pa.ChunkedArray]
|
||||
FetchBlobsAsync = Callable[[str, pa.Table], Awaitable[pa.Array | pa.ChunkedArray]]
|
||||
|
||||
|
||||
class BlobFile(io.RawIOBase):
|
||||
"""Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`.
|
||||
|
||||
Bytes load on ``read`` or ``read_range``, not when the handle is opened.
|
||||
Use :meth:`aread` from async code.
|
||||
"""
|
||||
|
||||
def __init__(self, inner) -> None:
|
||||
self._inner = inner
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
return await self._inner.read()
|
||||
|
||||
def close(self) -> None:
|
||||
self._inner.close()
|
||||
|
||||
@property
|
||||
def closed(self) -> bool:
|
||||
return self._inner.is_closed()
|
||||
|
||||
def readable(self) -> bool:
|
||||
return True
|
||||
|
||||
def seekable(self) -> bool:
|
||||
return True
|
||||
|
||||
def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
|
||||
if whence == io.SEEK_SET:
|
||||
self._inner.seek(offset)
|
||||
elif whence == io.SEEK_CUR:
|
||||
self._inner.seek(self._inner.tell() + offset)
|
||||
elif whence == io.SEEK_END:
|
||||
self._inner.seek(self._inner.size() + offset)
|
||||
else:
|
||||
raise ValueError(f"invalid whence: {whence}")
|
||||
return self._inner.tell()
|
||||
|
||||
def tell(self) -> int:
|
||||
return self._inner.tell()
|
||||
|
||||
def size(self) -> int:
|
||||
return self._inner.size()
|
||||
|
||||
def readall(self) -> bytes:
|
||||
return self._inner.read_bytes()
|
||||
|
||||
def read(self, size: int = -1) -> bytes:
|
||||
if size == -1:
|
||||
return self._inner.read_bytes()
|
||||
return super().read(size)
|
||||
|
||||
def read_range(self, offset: int, length: int) -> bytes:
|
||||
return self._inner.read_range(offset, length)
|
||||
|
||||
def readinto(self, b: WriteableBuffer) -> int:
|
||||
view = memoryview(b).cast("B")
|
||||
chunk = self._inner.read_up_to(len(view))
|
||||
view[: len(chunk)] = chunk
|
||||
return len(chunk)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<BlobFile size={self.size()}>"
|
||||
|
||||
|
||||
def validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in BLOB_MODE_TO_HANDLING:
|
||||
modes = ", ".join(repr(mode) for mode in BLOB_MODE_TO_HANDLING)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
|
||||
|
||||
def supports_blob_auto_row_id(table: Table | AsyncTable | RemoteTable) -> bool:
|
||||
"""Blob auto row-id applies to native tables, not LanceDB Cloud."""
|
||||
from .remote.table import RemoteTable
|
||||
|
||||
if isinstance(table, RemoteTable):
|
||||
return False
|
||||
|
||||
inner = getattr(table, "_inner", None)
|
||||
if inner is not None:
|
||||
uri = inner.database().uri
|
||||
if isinstance(uri, str) and get_uri_scheme(uri) == "db":
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def projection_includes_blob_column(
|
||||
projection: QueryProjection,
|
||||
blob_columns: Iterable[str],
|
||||
) -> bool:
|
||||
columns = set(blob_columns)
|
||||
if not columns:
|
||||
return False
|
||||
if projection is None:
|
||||
return True
|
||||
for output, source in _iter_projection_pairs(projection):
|
||||
if output in columns or source in columns:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def blob_v2_projection_sources(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
) -> dict[str, str]:
|
||||
blob_columns = blob_v2_column_paths(schema)
|
||||
if not blob_columns:
|
||||
return {}
|
||||
columns = set(blob_columns)
|
||||
if projection is None:
|
||||
return {column: column for column in blob_columns}
|
||||
return {
|
||||
output: source
|
||||
for output, source in _iter_projection_pairs(projection)
|
||||
if source in columns
|
||||
}
|
||||
|
||||
|
||||
def v2_projection_needs_row_id(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
*,
|
||||
with_row_id: bool,
|
||||
) -> bool:
|
||||
if with_row_id:
|
||||
return False
|
||||
return projection_includes_blob_column(projection, blob_v2_column_paths(schema))
|
||||
|
||||
|
||||
def blob_auto_row_id_for_scan(
|
||||
table: Table | AsyncTable | RemoteTable,
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
*,
|
||||
with_row_id: bool | None,
|
||||
) -> bool:
|
||||
if with_row_id is not None:
|
||||
return False
|
||||
if not supports_blob_auto_row_id(table):
|
||||
return False
|
||||
return v2_projection_needs_row_id(schema, projection, with_row_id=False)
|
||||
|
||||
|
||||
def finalize_blob_query_table(
|
||||
tbl: pa.Table,
|
||||
*,
|
||||
user_requested_row_id: bool,
|
||||
blob_auto_row_id: bool,
|
||||
blob_paths: Iterable[str] = (),
|
||||
) -> pa.Table:
|
||||
if user_requested_row_id or not blob_auto_row_id:
|
||||
return tbl
|
||||
return stash_auto_row_ids(tbl, blob_paths)
|
||||
|
||||
|
||||
async def replace_v2_blob_columns_with_bytes(
|
||||
tbl: pa.Table,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsAsync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if output_name not in tbl.column_names:
|
||||
continue
|
||||
blobs = await fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
return tbl
|
||||
|
||||
|
||||
def replace_v2_blob_columns_with_bytes_sync(
|
||||
tbl: pa.Table,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsSync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if output_name not in tbl.column_names:
|
||||
continue
|
||||
blobs = fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
return tbl
|
||||
|
||||
|
||||
def stash_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
if "_rowid" not in tbl.column_names:
|
||||
raise ValueError("query result has no '_rowid' column to hide")
|
||||
|
||||
present_paths = [p for p in blob_paths if p.split(".")[0] in tbl.column_names]
|
||||
if not present_paths:
|
||||
raise ValueError("query result has no blob v2 column to carry a row id")
|
||||
|
||||
row_ids = tbl["_rowid"]
|
||||
if isinstance(row_ids, pa.ChunkedArray):
|
||||
row_ids = row_ids.combine_chunks()
|
||||
row_ids = row_ids.cast(pa.uint64())
|
||||
|
||||
for path in present_paths:
|
||||
tbl = _embed_row_id_in_column(tbl, path, row_ids)
|
||||
return tbl.drop_columns(["_rowid"])
|
||||
|
||||
|
||||
def read_row_ids_from_hits(hits: pa.Table, blob_column: str) -> list[int]:
|
||||
if "_rowid" in hits.column_names:
|
||||
return hits["_rowid"].to_pylist()
|
||||
|
||||
try:
|
||||
leaf = _leaf_struct_column(hits, blob_column)
|
||||
if ROW_ID_FIELD_NAME in leaf.type.names:
|
||||
return leaf.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
# blob_column is the source name; aliased projections use the output name in hits.
|
||||
row_ids = _find_row_id_in_any_column(hits)
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
|
||||
raise ValueError(
|
||||
f"query result has no '_rowid' column and no '{ROW_ID_FIELD_NAME}' "
|
||||
f"field on blob column '{blob_column}'. Pass fresh blob query "
|
||||
"results, call .with_row_id(True), or pass a list of row ids."
|
||||
)
|
||||
|
||||
|
||||
def _find_row_id_in_any_column(tbl: pa.Table) -> Optional[list[int]]:
|
||||
for name in tbl.column_names:
|
||||
column = tbl.column(name)
|
||||
if isinstance(column, pa.ChunkedArray):
|
||||
column = column.combine_chunks()
|
||||
row_ids = _find_row_id_in_struct(column)
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
return None
|
||||
|
||||
|
||||
def _find_row_id_in_struct(array: pa.Array) -> Optional[list[int]]:
|
||||
if not pa.types.is_struct(array.type):
|
||||
return None
|
||||
if ROW_ID_FIELD_NAME in array.type.names:
|
||||
return array.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
for i in range(array.type.num_fields):
|
||||
row_ids = _find_row_id_in_struct(array.field(i))
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
return None
|
||||
|
||||
|
||||
def _iter_projection_pairs(
|
||||
projection: QueryProjectionSpec,
|
||||
) -> Iterable[tuple[str, str]]:
|
||||
if isinstance(projection, dict):
|
||||
for name, expr in projection.items():
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
return
|
||||
for column in projection:
|
||||
if isinstance(column, str):
|
||||
yield column, column
|
||||
elif isinstance(column, tuple) and len(column) == 2:
|
||||
name, expr = column
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
|
||||
|
||||
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
|
||||
index = tbl.schema.get_field_index(output_name)
|
||||
return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs])
|
||||
|
||||
|
||||
def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table:
|
||||
def add_row_id(children: list, child_fields: list) -> None:
|
||||
children.append(row_ids)
|
||||
child_fields.append(pa.field(ROW_ID_FIELD_NAME, pa.uint64(), nullable=False))
|
||||
|
||||
return _transform_struct_column(tbl, path, add_row_id)
|
||||
|
||||
|
||||
def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
"""Remove any `_lance_row_id` field embedded in blob descriptor structs.
|
||||
|
||||
For read-only descriptor views (`blob_mode="descriptions"`) that never
|
||||
fetch bytes, so have no use for the row id.
|
||||
"""
|
||||
|
||||
def drop_row_id(children: list, child_fields: list) -> None:
|
||||
for i, field in enumerate(child_fields):
|
||||
if field.name == ROW_ID_FIELD_NAME:
|
||||
del children[i], child_fields[i]
|
||||
return
|
||||
|
||||
for path in blob_paths:
|
||||
if path.split(".")[0] not in tbl.column_names:
|
||||
continue
|
||||
tbl = _transform_struct_column(tbl, path, drop_row_id)
|
||||
return tbl
|
||||
|
||||
|
||||
def _transform_struct_column(
|
||||
tbl: pa.Table, path: str, leaf_transform: Callable[[list, list], None]
|
||||
) -> pa.Table:
|
||||
top_name, *rest = path.split(".")
|
||||
top_index = tbl.schema.get_field_index(top_name)
|
||||
top_field = tbl.schema.field(top_index)
|
||||
top_array = tbl.column(top_name)
|
||||
if isinstance(top_array, pa.ChunkedArray):
|
||||
top_array = top_array.combine_chunks()
|
||||
|
||||
new_array, new_field = _rebuild_struct(top_array, top_field, rest, leaf_transform)
|
||||
return tbl.set_column(top_index, new_field, new_array)
|
||||
|
||||
|
||||
def _rebuild_struct(
|
||||
struct_array: pa.StructArray,
|
||||
struct_field: pa.Field,
|
||||
remaining_path: list[str],
|
||||
leaf_transform: Callable[[list, list], None],
|
||||
) -> tuple[pa.StructArray, pa.Field]:
|
||||
null_mask = struct_array.is_null()
|
||||
if not remaining_path:
|
||||
children = [struct_array.field(i) for i in range(struct_array.type.num_fields)]
|
||||
child_fields = list(struct_array.type)
|
||||
leaf_transform(children, child_fields)
|
||||
new_array = pa.StructArray.from_arrays(
|
||||
children, fields=child_fields, mask=null_mask
|
||||
)
|
||||
else:
|
||||
child_name = remaining_path[0]
|
||||
child_index = struct_array.type.get_field_index(child_name)
|
||||
child_array = struct_array.field(child_index)
|
||||
child_field = struct_array.type.field(child_index)
|
||||
new_child_array, new_child_field = _rebuild_struct(
|
||||
child_array, child_field, remaining_path[1:], leaf_transform
|
||||
)
|
||||
|
||||
children = []
|
||||
child_fields = []
|
||||
for i in range(struct_array.type.num_fields):
|
||||
field = struct_array.type.field(i)
|
||||
if field.name == child_name:
|
||||
children.append(new_child_array)
|
||||
child_fields.append(new_child_field)
|
||||
else:
|
||||
children.append(struct_array.field(i))
|
||||
child_fields.append(field)
|
||||
new_array = pa.StructArray.from_arrays(
|
||||
children, fields=child_fields, mask=null_mask
|
||||
)
|
||||
|
||||
new_field = pa.field(
|
||||
struct_field.name,
|
||||
new_array.type,
|
||||
nullable=struct_field.nullable,
|
||||
metadata=struct_field.metadata,
|
||||
)
|
||||
return new_array, new_field
|
||||
|
||||
|
||||
def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray:
|
||||
parts = path.split(".")
|
||||
column = tbl.column(parts[0])
|
||||
if isinstance(column, pa.ChunkedArray):
|
||||
column = column.combine_chunks()
|
||||
for part in parts[1:]:
|
||||
column = column.field(part)
|
||||
return column
|
||||
|
||||
|
||||
def _normalize_blob_row_ids(
|
||||
row_ids: Union[list[int], pa.Table], blob_column: str
|
||||
) -> list[int]:
|
||||
if isinstance(row_ids, pa.Table):
|
||||
return read_row_ids_from_hits(row_ids, blob_column)
|
||||
if isinstance(row_ids, (pa.Array, pa.ChunkedArray)):
|
||||
raise ValueError(
|
||||
"pass a query table with _rowid, not a column array "
|
||||
"(use fetch_blobs('image', hits), not fetch_blobs('image', hits['image']))"
|
||||
)
|
||||
return list(row_ids)
|
||||
|
||||
|
||||
def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]:
|
||||
return [BlobFile(handle) if handle is not None else None for handle in handles]
|
||||
@@ -1,4 +1,5 @@
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -24,11 +25,45 @@ from lance_namespace import (
|
||||
ListTablesResponse,
|
||||
)
|
||||
from .remote import ClientConfig
|
||||
from .types import BaseTokenizerType
|
||||
|
||||
IvfHnswPq: type[HnswPq] = HnswPq
|
||||
IvfHnswSq: type[HnswSq] = HnswSq
|
||||
IvfHnswFlat: type[HnswFlat] = HnswFlat
|
||||
|
||||
class MetricPoint:
|
||||
name: str
|
||||
kind: str
|
||||
attributes: Dict[str, str]
|
||||
value: Optional[float]
|
||||
buckets: Optional[List[Tuple[str, int]]]
|
||||
count: Optional[int]
|
||||
sum: Optional[float]
|
||||
|
||||
class MetricDescription:
|
||||
name: str
|
||||
kind: str
|
||||
unit: Optional[str]
|
||||
description: str
|
||||
|
||||
def register_lancedb_metrics_recorder() -> bool: ...
|
||||
def lancedb_metrics_catalog() -> List[MetricDescription]: ...
|
||||
def snapshot_lancedb_metrics() -> List[MetricPoint]: ...
|
||||
def tokenize(
|
||||
query: str,
|
||||
*,
|
||||
base_tokenizer: BaseTokenizerType = "simple",
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
) -> List["FtsToken"]: ...
|
||||
|
||||
class PyExpr:
|
||||
"""A type-safe DataFusion expression node (Rust-side handle)."""
|
||||
|
||||
@@ -53,7 +88,9 @@ class PyExpr:
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ...
|
||||
def expr_lit(
|
||||
value: Union[bool, int, float, str, bytes, date, datetime, Decimal],
|
||||
) -> PyExpr: ...
|
||||
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
||||
|
||||
class Session:
|
||||
@@ -159,6 +196,17 @@ class Connection(object):
|
||||
self,
|
||||
) -> Dict[str, Any]: ...
|
||||
|
||||
class BlobFile:
|
||||
async def read(self) -> bytes: ...
|
||||
def read_bytes(self) -> bytes: ...
|
||||
def close(self) -> None: ...
|
||||
def is_closed(self) -> bool: ...
|
||||
def seek(self, position: int) -> None: ...
|
||||
def tell(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||
def read_up_to(self, length: int) -> bytes: ...
|
||||
|
||||
class Table:
|
||||
def name(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
@@ -205,6 +253,13 @@ class Table:
|
||||
async def prewarm_index(self, index_name: str) -> None: ...
|
||||
async def prewarm_data(self, columns: Optional[List[str]] = None) -> None: ...
|
||||
async def list_indices(self) -> list[IndexConfig]: ...
|
||||
async def tokenize(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
column: Optional[str] = None,
|
||||
index_name: Optional[str] = None,
|
||||
) -> list[FtsToken]: ...
|
||||
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
|
||||
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
|
||||
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
|
||||
@@ -226,6 +281,7 @@ class Table:
|
||||
async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
|
||||
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
||||
async def unset_lsm_write_spec(self) -> None: ...
|
||||
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
|
||||
async def close_lsm_writers(self) -> None: ...
|
||||
@property
|
||||
def tags(self) -> Tags: ...
|
||||
@@ -235,6 +291,13 @@ class Table:
|
||||
def query(self) -> Query: ...
|
||||
def take_offsets(self, offsets: list[int]) -> TakeQuery: ...
|
||||
def take_row_ids(self, row_ids: list[int]) -> TakeQuery: ...
|
||||
async def blob_columns(self) -> list[str]: ...
|
||||
async def fetch_blobs(
|
||||
self, column: str, row_ids: list[int]
|
||||
) -> pa.LargeBinaryArray: ...
|
||||
async def fetch_blob_files(
|
||||
self, column: str, row_ids: list[int]
|
||||
) -> list[Optional[BlobFile]]: ...
|
||||
def vector_search(self) -> VectorQuery: ...
|
||||
|
||||
class Tags:
|
||||
@@ -280,6 +343,24 @@ async def connect(
|
||||
session: Optional[Session],
|
||||
manifest_enabled: bool = False,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
oauth_config: Optional[Any] = None,
|
||||
) -> Connection: ...
|
||||
def connect_namespace(
|
||||
namespace_client_impl: str,
|
||||
namespace_client_properties: Dict[str, str],
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
) -> Connection: ...
|
||||
def connect_namespace_client(
|
||||
namespace_client: Any,
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
) -> Connection: ...
|
||||
|
||||
class RecordBatchStream:
|
||||
@@ -452,6 +533,10 @@ class MergeResult:
|
||||
num_attempts: int
|
||||
num_rows: int
|
||||
|
||||
class FtsToken:
|
||||
text: str
|
||||
position: int
|
||||
|
||||
class LsmWriteSpec:
|
||||
"""Specification selecting Lance's MemWAL LSM-style write path for
|
||||
`merge_insert`."""
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
import os
|
||||
from functools import cached_property
|
||||
from typing import List, Union
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -15,6 +15,8 @@ from .base import TextEmbeddingFunction
|
||||
from .registry import register
|
||||
from .utils import TEXT, api_key_not_found_help
|
||||
|
||||
EMBEDDING_BATCH_SIZE = 100
|
||||
|
||||
|
||||
@register("gemini-text")
|
||||
class GeminiText(TextEmbeddingFunction):
|
||||
@@ -81,6 +83,7 @@ class GeminiText(TextEmbeddingFunction):
|
||||
"""
|
||||
|
||||
name: str = "gemini-embedding-001"
|
||||
dim: Optional[int] = None
|
||||
query_task_type: str = "retrieval_query"
|
||||
source_task_type: str = "retrieval_document"
|
||||
|
||||
@@ -93,6 +96,8 @@ class GeminiText(TextEmbeddingFunction):
|
||||
model_config["ignored_types"] = (cached_property,)
|
||||
|
||||
def ndims(self):
|
||||
if self.dim:
|
||||
return self.dim
|
||||
# TODO: fix hardcoding
|
||||
return 768
|
||||
|
||||
@@ -133,22 +138,22 @@ class GeminiText(TextEmbeddingFunction):
|
||||
contents.append({"parts": [{"text": text}]})
|
||||
|
||||
# Build config
|
||||
config_kwargs = {}
|
||||
config_kwargs = {"output_dimensionality": self.ndims()}
|
||||
if task_type:
|
||||
config_kwargs["task_type"] = task_type.upper() # API expects uppercase
|
||||
|
||||
# Call embed_content for each content
|
||||
config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
# Call embed_content in groups of at most EMBEDDING_BATCH_SIZE docs at a time
|
||||
embeddings = []
|
||||
for content in contents:
|
||||
config = (
|
||||
types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
)
|
||||
for i in range(0, len(contents), EMBEDDING_BATCH_SIZE):
|
||||
chunk = contents[i : i + EMBEDDING_BATCH_SIZE]
|
||||
response = self.client.models.embed_content(
|
||||
model=self.name,
|
||||
contents=content,
|
||||
contents=chunk,
|
||||
config=config,
|
||||
)
|
||||
embeddings.append(response.embeddings[0].values)
|
||||
embeddings.extend([np.array(e.values) for e in response.embeddings])
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -160,5 +165,13 @@ class GeminiText(TextEmbeddingFunction):
|
||||
api_key_not_found_help("google")
|
||||
|
||||
from google import genai as genai_module
|
||||
from lancedb import __version__
|
||||
|
||||
return genai_module.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
|
||||
return genai_module.Client(
|
||||
api_key=os.environ.get("GOOGLE_API_KEY"),
|
||||
http_options={
|
||||
"headers": {
|
||||
"x-goog-api-client": f"lancedb/{__version__}",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -56,6 +56,16 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
embeddings = self._compute_embedding(texts)
|
||||
return list(embeddings)
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
state = super().__getstate__()
|
||||
state["__dict__"] = {
|
||||
k: v for k, v in state["__dict__"].items() if k != "_ollama_client"
|
||||
}
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: dict[str, Any]) -> None:
|
||||
super().__setstate__(state)
|
||||
|
||||
@cached_property
|
||||
def _ollama_client(self) -> "ollama.Client":
|
||||
ollama = attempt_import_or_raise("ollama")
|
||||
|
||||
@@ -19,6 +19,8 @@ operators::
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Iterable, Union
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -63,7 +65,7 @@ def _coerce(value: "ExprLike") -> "Expr":
|
||||
|
||||
|
||||
# Type alias used in annotations.
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes]
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes, date, datetime, Decimal]
|
||||
|
||||
|
||||
class Expr:
|
||||
@@ -118,10 +120,18 @@ class Expr:
|
||||
"""Logical AND (``expr_a & expr_b``)."""
|
||||
return Expr(self._inner.and_(_coerce(other)._inner))
|
||||
|
||||
def __rand__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical AND (``True & expr``)."""
|
||||
return Expr(_coerce(other)._inner.and_(self._inner))
|
||||
|
||||
def __or__(self, other: "Expr") -> "Expr":
|
||||
"""Logical OR (``expr_a | expr_b``)."""
|
||||
return Expr(self._inner.or_(_coerce(other)._inner))
|
||||
|
||||
def __ror__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical OR (``False | expr``)."""
|
||||
return Expr(_coerce(other)._inner.or_(self._inner))
|
||||
|
||||
def __invert__(self) -> "Expr":
|
||||
"""Logical NOT (``~expr``)."""
|
||||
return Expr(self._inner.not_())
|
||||
@@ -266,13 +276,14 @@ def col(name: str) -> Expr:
|
||||
return Expr(expr_col(name))
|
||||
|
||||
|
||||
def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr:
|
||||
"""Create a literal (constant) value expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value:
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``.
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``,
|
||||
``datetime``, or ``Decimal``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -280,6 +291,9 @@ def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
>>> col("price") * lit(1.1)
|
||||
Expr((price * 1.1))
|
||||
"""
|
||||
if not isinstance(value, (bool, int, float, str, bytes, date, datetime, Decimal)):
|
||||
raise TypeError(f"Unsupported literal type: {type(value).__name__}")
|
||||
|
||||
return Expr(expr_lit(value))
|
||||
|
||||
|
||||
|
||||
@@ -127,6 +127,8 @@ class FTS:
|
||||
- "whitespace": Split text by whitespace, but not punctuation.
|
||||
- "raw": No tokenization. The entire text is treated as a single token.
|
||||
- "ngram": N-gram tokenizer for substring-style matching.
|
||||
- "icu": ICU dictionary-based word segmentation.
|
||||
- "icu/split": ICU segmentation with simple-style delimiter splitting.
|
||||
- "jieba/*": Jieba tokenizer loaded from Lance's language model home.
|
||||
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
||||
language : str, default "English"
|
||||
|
||||
@@ -51,6 +51,15 @@ class LanceMergeInsertBuilder(object):
|
||||
If there are multiple matches then the behavior is undefined.
|
||||
Currently this causes multiple copies of the row to be created
|
||||
but that behavior is subject to change.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: Optional[str], default None
|
||||
An optional filter to limit which rows are updated. Column
|
||||
references in this expression must be prefixed with "target."
|
||||
to refer to the existing table data. For example, to only
|
||||
update rows where the existing color is red, use:
|
||||
``where="target.color = 'red'"``
|
||||
"""
|
||||
self._when_matched_update_all = True
|
||||
self._when_matched_update_all_condition = where
|
||||
|
||||
+258
-169
@@ -38,15 +38,13 @@ from lance_namespace_urllib3_client.models.query_table_request_vector import (
|
||||
QueryTableRequestVector,
|
||||
)
|
||||
from lance_namespace_urllib3_client.models.string_fts_query import StringFtsQuery
|
||||
from lance_namespace.errors import TableNotFoundError
|
||||
from lancedb._lancedb import connect_namespace_client as _connect_namespace_client
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb._lancedb import (
|
||||
connect_namespace as _connect_namespace,
|
||||
connect_namespace_client as _connect_namespace_client,
|
||||
)
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
_normalize_drop_namespace_mode,
|
||||
_normalize_drop_namespace_behavior,
|
||||
)
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
connect as namespace_connect,
|
||||
@@ -55,13 +53,6 @@ from lance_namespace import (
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
ListTablesRequest,
|
||||
DescribeNamespaceRequest,
|
||||
DropTableRequest,
|
||||
RenameTableRequest,
|
||||
ListNamespacesRequest,
|
||||
CreateNamespaceRequest,
|
||||
DropNamespaceRequest,
|
||||
)
|
||||
from lancedb.table import AsyncTable, LanceTable, Table
|
||||
from lancedb.util import validate_table_name
|
||||
@@ -373,6 +364,23 @@ def _convert_pyarrow_schema_to_json(schema: pa.Schema) -> JsonArrowSchema:
|
||||
return JsonArrowSchema(fields=fields, metadata=meta)
|
||||
|
||||
|
||||
def _builds_namespace_natively(
|
||||
namespace_client_impl: Optional[str],
|
||||
namespace_client_properties: Optional[Dict[str, str]],
|
||||
) -> bool:
|
||||
"""Whether ``connect_namespace_client`` builds the namespace client natively
|
||||
in Rust (installing the read-freshness context provider) rather than wrapping
|
||||
the pre-built Python client.
|
||||
|
||||
Must mirror Rust ``build_namespace_natively`` in ``python/src/connection.rs``.
|
||||
"""
|
||||
return namespace_client_impl == "rest" and bool(namespace_client_properties)
|
||||
|
||||
|
||||
def _supports_native_namespace(namespace_client_impl: str) -> bool:
|
||||
return namespace_client_impl in {"dir", "rest"}
|
||||
|
||||
|
||||
class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
A LanceDB connection that uses a namespace for table management.
|
||||
@@ -383,7 +391,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace_client: LanceNamespace,
|
||||
namespace_client: Optional[LanceNamespace] = None,
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
@@ -391,6 +399,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
):
|
||||
"""
|
||||
Initialize a namespace-based LanceDB connection.
|
||||
@@ -432,23 +441,36 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
# When the namespace connection or client is built natively in Rust, the
|
||||
# underlying Rust table performs QueryTable pushdown through the
|
||||
# read-freshness context provider, which the pure-Python ``query_table``
|
||||
# path bypasses.
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
else:
|
||||
if namespace_client is None:
|
||||
raise ValueError("namespace_client is required without a native _inner")
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
)
|
||||
self._uri = self._inner.uri
|
||||
|
||||
@override
|
||||
def serialize(self) -> str:
|
||||
@@ -494,11 +516,11 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
return LOOP.run(
|
||||
self._inner.table_names(
|
||||
namespace_path=namespace_path, start_after=page_token, limit=limit
|
||||
)
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return response.tables if response.tables else []
|
||||
|
||||
@override
|
||||
def create_table(
|
||||
@@ -543,6 +565,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_path=namespace_path,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
_async=async_table,
|
||||
)
|
||||
|
||||
@@ -568,8 +591,8 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Table not found" in str(e):
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if "Table not found" in str(e) or "was not found" in str(e):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
|
||||
raise
|
||||
@@ -580,6 +603,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_path=namespace_path,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
_async=async_table,
|
||||
)
|
||||
if branch is not None:
|
||||
@@ -590,12 +614,9 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
@override
|
||||
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
# Use namespace drop_table directly
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
table_id = namespace_path + [name]
|
||||
request = DropTableRequest(id=table_id)
|
||||
self._namespace_client.drop_table(request)
|
||||
LOOP.run(self._inner.drop_table(name, namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def rename_table(
|
||||
@@ -609,14 +630,19 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
cur_namespace_path = []
|
||||
if new_namespace_path is None:
|
||||
new_namespace_path = []
|
||||
cur_table_id = cur_namespace_path + [cur_name]
|
||||
new_namespace_id = new_namespace_path if new_namespace_path else None
|
||||
request = RenameTableRequest(
|
||||
id=cur_table_id,
|
||||
new_table_name=new_name,
|
||||
new_namespace_id=new_namespace_id,
|
||||
)
|
||||
self._namespace_client.rename_table(request)
|
||||
try:
|
||||
LOOP.run(
|
||||
self._inner.rename_table(
|
||||
cur_name,
|
||||
new_name,
|
||||
cur_namespace_path=cur_namespace_path,
|
||||
new_namespace_path=new_namespace_path,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "rename_table not implemented" in str(e):
|
||||
raise NotImplementedError("rename_table not implemented") from e
|
||||
raise
|
||||
|
||||
@override
|
||||
def drop_database(self):
|
||||
@@ -628,8 +654,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
def drop_all_tables(self, namespace_path: Optional[List[str]] = None):
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
for table_name in self.table_names(namespace_path=namespace_path):
|
||||
self.drop_table(table_name, namespace_path=namespace_path)
|
||||
LOOP.run(self._inner.drop_all_tables(namespace_path=namespace_path))
|
||||
|
||||
@override
|
||||
def list_namespaces(
|
||||
@@ -659,13 +684,10 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListNamespacesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_namespaces(request)
|
||||
return ListNamespacesResponse(
|
||||
namespaces=response.namespaces if response.namespaces else [],
|
||||
page_token=response.page_token,
|
||||
return LOOP.run(
|
||||
self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -693,14 +715,12 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
CreateNamespaceResponse
|
||||
Response containing the properties of the created namespace.
|
||||
"""
|
||||
request = CreateNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_create_namespace_mode(mode),
|
||||
properties=properties,
|
||||
)
|
||||
response = self._namespace_client.create_namespace(request)
|
||||
return CreateNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
return LOOP.run(
|
||||
self._inner.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
)
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -728,20 +748,18 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
DropNamespaceResponse
|
||||
Response containing properties and transaction_id if applicable.
|
||||
"""
|
||||
request = DropNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_drop_namespace_mode(mode),
|
||||
behavior=_normalize_drop_namespace_behavior(behavior),
|
||||
)
|
||||
response = self._namespace_client.drop_namespace(request)
|
||||
return DropNamespaceResponse(
|
||||
properties=(
|
||||
response.properties if hasattr(response, "properties") else None
|
||||
),
|
||||
transaction_id=(
|
||||
response.transaction_id if hasattr(response, "transaction_id") else None
|
||||
),
|
||||
)
|
||||
try:
|
||||
return LOOP.run(
|
||||
self._inner.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
|
||||
@override
|
||||
def describe_namespace(
|
||||
@@ -760,11 +778,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
DescribeNamespaceResponse
|
||||
Response containing the namespace properties.
|
||||
"""
|
||||
request = DescribeNamespaceRequest(id=namespace_path)
|
||||
response = self._namespace_client.describe_namespace(request)
|
||||
return DescribeNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
return LOOP.run(self._inner.describe_namespace(namespace_path))
|
||||
|
||||
@override
|
||||
def list_tables(
|
||||
@@ -794,13 +808,10 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return ListTablesResponse(
|
||||
tables=response.tables if response.tables else [],
|
||||
page_token=response.page_token,
|
||||
return LOOP.run(
|
||||
self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
)
|
||||
|
||||
def _lance_table_from_uri(
|
||||
@@ -856,6 +867,18 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
LanceNamespace
|
||||
The namespace client for this connection.
|
||||
"""
|
||||
if self._namespace_client is None:
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot construct a Python namespace client without "
|
||||
"namespace implementation properties"
|
||||
)
|
||||
self._namespace_client = namespace_connect(
|
||||
self._namespace_client_impl, self._namespace_client_properties
|
||||
)
|
||||
return self._namespace_client
|
||||
|
||||
|
||||
@@ -869,12 +892,15 @@ class AsyncLanceNamespaceDBConnection:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace_client: LanceNamespace,
|
||||
namespace_client: Optional[LanceNamespace] = None,
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
):
|
||||
"""
|
||||
Initialize an async namespace-based LanceDB connection.
|
||||
@@ -900,6 +926,12 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace.create_table() instead of using declare_table + local write.
|
||||
|
||||
Default is None (no pushdown, all operations run locally).
|
||||
namespace_client_impl : Optional[str]
|
||||
The namespace implementation name used to create this connection.
|
||||
Required (with ``namespace_client_properties``) for the Rust client to
|
||||
be built natively and install the read-freshness provider.
|
||||
namespace_client_properties : Optional[Dict[str, str]]
|
||||
The namespace properties used to create this connection.
|
||||
"""
|
||||
self._namespace_client = namespace_client
|
||||
self.read_consistency_interval = read_consistency_interval
|
||||
@@ -908,23 +940,37 @@ class AsyncLanceNamespaceDBConnection:
|
||||
self._namespace_client_pushdown_operations = set(
|
||||
namespace_client_pushdown_operations or []
|
||||
)
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=None,
|
||||
namespace_client_properties=None,
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# See LanceNamespaceDBConnection: when Rust owns the namespace
|
||||
# connection/client, its table performs QueryTable pushdown through the
|
||||
# read-freshness provider, so defer to it rather than the urllib3 client
|
||||
# path (which omits x-lancedb-min-timestamp).
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
else:
|
||||
if namespace_client is None:
|
||||
raise ValueError("namespace_client is required without a native _inner")
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
)
|
||||
|
||||
async def table_names(
|
||||
self,
|
||||
@@ -948,11 +994,9 @@ class AsyncLanceNamespaceDBConnection:
|
||||
)
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
return await self._inner.table_names(
|
||||
namespace_path=namespace_path, start_after=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return response.tables if response.tables else []
|
||||
|
||||
async def create_table(
|
||||
self,
|
||||
@@ -992,6 +1036,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_path=namespace_path,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
)
|
||||
|
||||
async def open_table(
|
||||
@@ -1014,8 +1059,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Table not found" in str(e):
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if "Table not found" in str(e) or "was not found" in str(e):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
|
||||
raise
|
||||
@@ -1029,15 +1074,14 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_path=namespace_path,
|
||||
namespace_client=self._namespace_client,
|
||||
pushdown_operations=self._namespace_client_pushdown_operations,
|
||||
route_pushdown_to_rust=self._route_pushdown_to_rust,
|
||||
)
|
||||
|
||||
async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
"""Drop a table from the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
table_id = namespace_path + [name]
|
||||
request = DropTableRequest(id=table_id)
|
||||
self._namespace_client.drop_table(request)
|
||||
await self._inner.drop_table(name, namespace_path=namespace_path)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
@@ -1051,14 +1095,17 @@ class AsyncLanceNamespaceDBConnection:
|
||||
cur_namespace_path = []
|
||||
if new_namespace_path is None:
|
||||
new_namespace_path = []
|
||||
cur_table_id = cur_namespace_path + [cur_name]
|
||||
new_namespace_id = new_namespace_path if new_namespace_path else None
|
||||
request = RenameTableRequest(
|
||||
id=cur_table_id,
|
||||
new_table_name=new_name,
|
||||
new_namespace_id=new_namespace_id,
|
||||
)
|
||||
self._namespace_client.rename_table(request)
|
||||
try:
|
||||
await self._inner.rename_table(
|
||||
cur_name,
|
||||
new_name,
|
||||
cur_namespace_path=cur_namespace_path,
|
||||
new_namespace_path=new_namespace_path,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "rename_table not implemented" in str(e):
|
||||
raise NotImplementedError("rename_table not implemented") from e
|
||||
raise
|
||||
|
||||
async def drop_database(self):
|
||||
"""Deprecated method."""
|
||||
@@ -1070,9 +1117,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""Drop all tables in the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
table_names = await self.table_names(namespace_path=namespace_path)
|
||||
for table_name in table_names:
|
||||
await self.drop_table(table_name, namespace_path=namespace_path)
|
||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||
|
||||
async def list_namespaces(
|
||||
self,
|
||||
@@ -1101,13 +1146,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListNamespacesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_namespaces(request)
|
||||
return ListNamespacesResponse(
|
||||
namespaces=response.namespaces if response.namespaces else [],
|
||||
page_token=response.page_token,
|
||||
return await self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
async def create_namespace(
|
||||
@@ -1134,15 +1174,11 @@ class AsyncLanceNamespaceDBConnection:
|
||||
CreateNamespaceResponse
|
||||
Response containing the properties of the created namespace.
|
||||
"""
|
||||
request = CreateNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_create_namespace_mode(mode),
|
||||
return await self._inner.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
)
|
||||
response = self._namespace_client.create_namespace(request)
|
||||
return CreateNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
async def drop_namespace(
|
||||
self,
|
||||
@@ -1168,20 +1204,16 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DropNamespaceResponse
|
||||
Response containing properties and transaction_id if applicable.
|
||||
"""
|
||||
request = DropNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_drop_namespace_mode(mode),
|
||||
behavior=_normalize_drop_namespace_behavior(behavior),
|
||||
)
|
||||
response = self._namespace_client.drop_namespace(request)
|
||||
return DropNamespaceResponse(
|
||||
properties=(
|
||||
response.properties if hasattr(response, "properties") else None
|
||||
),
|
||||
transaction_id=(
|
||||
response.transaction_id if hasattr(response, "transaction_id") else None
|
||||
),
|
||||
)
|
||||
try:
|
||||
return await self._inner.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
|
||||
async def describe_namespace(
|
||||
self, namespace_path: List[str]
|
||||
@@ -1199,11 +1231,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DescribeNamespaceResponse
|
||||
Response containing the namespace properties.
|
||||
"""
|
||||
request = DescribeNamespaceRequest(id=namespace_path)
|
||||
response = self._namespace_client.describe_namespace(request)
|
||||
return DescribeNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
return await self._inner.describe_namespace(namespace_path)
|
||||
|
||||
async def list_tables(
|
||||
self,
|
||||
@@ -1232,13 +1260,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return ListTablesResponse(
|
||||
tables=response.tables if response.tables else [],
|
||||
page_token=response.page_token,
|
||||
return await self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
async def namespace_client(self) -> LanceNamespace:
|
||||
@@ -1252,6 +1275,18 @@ class AsyncLanceNamespaceDBConnection:
|
||||
LanceNamespace
|
||||
The namespace client for this connection.
|
||||
"""
|
||||
if self._namespace_client is None:
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot construct a Python namespace client without "
|
||||
"namespace implementation properties"
|
||||
)
|
||||
self._namespace_client = namespace_connect(
|
||||
self._namespace_client_impl, self._namespace_client_properties
|
||||
)
|
||||
return self._namespace_client
|
||||
|
||||
|
||||
@@ -1302,6 +1337,32 @@ def connect_namespace(
|
||||
LanceNamespaceDBConnection
|
||||
A namespace-based connection to LanceDB
|
||||
"""
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
)
|
||||
)
|
||||
return LanceNamespaceDBConnection(
|
||||
namespace_client=None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
@@ -1377,6 +1438,32 @@ def connect_namespace_async(
|
||||
... tables = await db.table_names()
|
||||
... table = await db.create_table("my_table", schema=schema)
|
||||
"""
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
)
|
||||
)
|
||||
return AsyncLanceNamespaceDBConnection(
|
||||
namespace_client=None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
@@ -1387,4 +1474,6 @@ def connect_namespace_async(
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Bridge LanceDB's internal metrics into OpenTelemetry.
|
||||
|
||||
LanceDB (through Lance core) publishes metrics (currently object store request
|
||||
counts, bytes, latency, errors, and throttles) through the Rust ``metrics``
|
||||
facade. This module installs a process-global recorder that aggregates them and
|
||||
registers OpenTelemetry observable instruments that report the aggregated values
|
||||
into the user's ``MeterProvider``.
|
||||
|
||||
The bridge is generic: every metric LanceDB describes is surfaced automatically,
|
||||
with no per-metric Python code. Histograms have no asynchronous OpenTelemetry
|
||||
instrument, so each is exported Prometheus-style as cumulative ``le`` buckets
|
||||
plus ``_count`` and ``_sum`` observable counters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from ._lancedb import (
|
||||
lancedb_metrics_catalog,
|
||||
register_lancedb_metrics_recorder,
|
||||
snapshot_lancedb_metrics,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.metrics import MeterProvider
|
||||
|
||||
_INSTRUMENTED = False
|
||||
|
||||
|
||||
def instrument_lancedb_metrics(
|
||||
meter_provider: Optional["MeterProvider"] = None,
|
||||
) -> bool:
|
||||
"""Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
|
||||
Installs a process-global metrics recorder and creates one observable
|
||||
instrument per LanceDB metric on the given (or global) ``MeterProvider``. The
|
||||
user's configured ``MetricReader`` then collects them on its own schedule.
|
||||
|
||||
Counters and gauges map directly to observable counters/gauges. Each
|
||||
histogram is exported as cumulative ``le`` bucket counts (``<name>_bucket``,
|
||||
with an ``le`` attribute) plus ``<name>_count`` and ``<name>_sum``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
meter_provider : opentelemetry.metrics.MeterProvider, optional
|
||||
The provider to register instruments on. Defaults to the global provider
|
||||
from ``opentelemetry.metrics.get_meter_provider()``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
``True`` if the recorder is installed and instruments are registered.
|
||||
``False`` if a different ``metrics`` recorder is already installed in
|
||||
this process (``metrics`` permits only one global recorder), in which
|
||||
case a warning is emitted and no instruments are created.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Requires the OpenTelemetry API (``pip install lancedb[otel]``) and, to
|
||||
actually export, an OpenTelemetry SDK (``pip install opentelemetry-sdk``)
|
||||
configured by the application. Calling this more than once is safe;
|
||||
instruments are created only on the first successful call.
|
||||
"""
|
||||
global _INSTRUMENTED
|
||||
|
||||
try:
|
||||
from opentelemetry.metrics import Observation, get_meter_provider
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"instrument_lancedb_metrics requires the OpenTelemetry API/SDK. "
|
||||
"Install it with `pip install lancedb[otel]` or "
|
||||
"`pip install opentelemetry-sdk`."
|
||||
) from exc
|
||||
|
||||
if not register_lancedb_metrics_recorder():
|
||||
warnings.warn(
|
||||
"Could not install the LanceDB metrics recorder: another `metrics` "
|
||||
"recorder is already installed in this process. LanceDB metrics will "
|
||||
"not be exported via OpenTelemetry.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return False
|
||||
|
||||
if _INSTRUMENTED:
|
||||
return True
|
||||
|
||||
provider = meter_provider or get_meter_provider()
|
||||
meter = provider.get_meter("lancedb")
|
||||
|
||||
def scalar_callback(metric_name: str):
|
||||
def callback(_options):
|
||||
return [
|
||||
Observation(point.value, point.attributes)
|
||||
for point in snapshot_lancedb_metrics()
|
||||
if point.name == metric_name and point.value is not None
|
||||
]
|
||||
|
||||
return callback
|
||||
|
||||
def bucket_callback(metric_name: str):
|
||||
def callback(_options):
|
||||
observations = []
|
||||
for point in snapshot_lancedb_metrics():
|
||||
if point.name != metric_name or point.buckets is None:
|
||||
continue
|
||||
for le, cumulative in point.buckets:
|
||||
attributes = dict(point.attributes)
|
||||
attributes["le"] = le
|
||||
observations.append(Observation(cumulative, attributes))
|
||||
return observations
|
||||
|
||||
return callback
|
||||
|
||||
def field_callback(metric_name: str, field: str):
|
||||
def callback(_options):
|
||||
observations = []
|
||||
for point in snapshot_lancedb_metrics():
|
||||
if point.name != metric_name:
|
||||
continue
|
||||
value = getattr(point, field)
|
||||
if value is not None:
|
||||
observations.append(Observation(value, point.attributes))
|
||||
return observations
|
||||
|
||||
return callback
|
||||
|
||||
for desc in lancedb_metrics_catalog():
|
||||
unit = desc.unit or ""
|
||||
if desc.kind == "counter":
|
||||
meter.create_observable_counter(
|
||||
desc.name,
|
||||
callbacks=[scalar_callback(desc.name)],
|
||||
unit=unit,
|
||||
description=desc.description,
|
||||
)
|
||||
elif desc.kind == "gauge":
|
||||
meter.create_observable_gauge(
|
||||
desc.name,
|
||||
callbacks=[scalar_callback(desc.name)],
|
||||
unit=unit,
|
||||
description=desc.description,
|
||||
)
|
||||
elif desc.kind == "histogram":
|
||||
# `_bucket` and `_count` observe cumulative sample counts, not the
|
||||
# histogram's measured quantity, so they are unitless; only `_sum`
|
||||
# carries the histogram's unit.
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_bucket",
|
||||
callbacks=[bucket_callback(desc.name)],
|
||||
description=f"{desc.description} (cumulative buckets)",
|
||||
)
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_count",
|
||||
callbacks=[field_callback(desc.name, "count")],
|
||||
description=f"{desc.description} (count)",
|
||||
)
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_sum",
|
||||
callbacks=[field_callback(desc.name, "sum")],
|
||||
unit=unit,
|
||||
description=f"{desc.description} (sum)",
|
||||
)
|
||||
|
||||
_INSTRUMENTED = True
|
||||
return True
|
||||
@@ -11,7 +11,7 @@ import pyarrow as pa
|
||||
from ._lancedb import async_permutation_builder, PermutationReader
|
||||
from .table import LanceTable, Table
|
||||
from .background_loop import LOOP
|
||||
from .util import batch_to_tensor, batch_to_tensor_rows
|
||||
from .util import batch_to_tensor, batch_to_tensor_dict, batch_to_tensor_rows
|
||||
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -48,6 +48,14 @@ class PermutationBuilder:
|
||||
By default, the permutation builder will create a single split that contains all
|
||||
rows in the same order as the base table.
|
||||
"""
|
||||
if not hasattr(table, "_inner"):
|
||||
raise TypeError(
|
||||
f"PermutationBuilder requires a local LanceTable, "
|
||||
f"got {type(table).__name__}. "
|
||||
"The permutation API is not supported on remote tables. "
|
||||
"Remote tables connect to LanceDB Cloud or Enterprise and do not have "
|
||||
"direct access to the underlying Lance dataset needed for permutations."
|
||||
)
|
||||
self._async = async_permutation_builder(table)
|
||||
|
||||
def split_random(
|
||||
@@ -57,6 +65,7 @@ class PermutationBuilder:
|
||||
counts: Optional[list[int]] = None,
|
||||
fixed: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
clump_size: Optional[int] = None,
|
||||
split_names: Optional[list[str]] = None,
|
||||
) -> "PermutationBuilder":
|
||||
"""
|
||||
@@ -79,6 +88,9 @@ class PermutationBuilder:
|
||||
Rows will be randomly assigned to splits. The optional seed can be provided to
|
||||
make the assignment deterministic.
|
||||
|
||||
If clump_size is provided, rows are shuffled as contiguous groups of that size,
|
||||
preserving I/O locality while still randomising the split assignment.
|
||||
|
||||
The optional split_names can be provided to name the splits. If not provided,
|
||||
the splits can only be referenced by their index.
|
||||
"""
|
||||
@@ -87,6 +99,7 @@ class PermutationBuilder:
|
||||
counts=counts,
|
||||
fixed=fixed,
|
||||
seed=seed,
|
||||
clump_size=clump_size,
|
||||
split_names=split_names,
|
||||
)
|
||||
return self
|
||||
@@ -933,6 +946,7 @@ class Permutation:
|
||||
"pandas",
|
||||
"arrow",
|
||||
"torch",
|
||||
"torch_row",
|
||||
"torch_col",
|
||||
"polars",
|
||||
],
|
||||
@@ -948,15 +962,19 @@ class Permutation:
|
||||
- "python_col" - the batch will be a dict of lists (one entry per column)
|
||||
- "pandas" - the batch will be a pandas DataFrame
|
||||
- "arrow" - the batch will be a pyarrow RecordBatch
|
||||
- "torch" - the batch will be a list of tensors, one per row
|
||||
- "torch" - the batch will be a list of per-row dicts mapping column
|
||||
name to a 0-D torch tensor. Works with the default
|
||||
``torch.utils.data.DataLoader`` collate, which stacks the per-row
|
||||
dicts back into a dict of batched tensors.
|
||||
- "torch_row" - the batch will be a list of tensors, one per row
|
||||
- "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns)
|
||||
- "polars" - the batch will be a polars DataFrame
|
||||
|
||||
Conversion may or may not involve a data copy. Lance uses Arrow internally
|
||||
and so it is able to zero-copy to the arrow and polars formats.
|
||||
|
||||
Conversion to torch_col will be zero-copy but will only support a subset of data
|
||||
types (numeric types).
|
||||
Conversion to torch and torch_col will be zero-copy but will only support a
|
||||
subset of data types (numeric types).
|
||||
|
||||
Conversion to numpy and/or pandas will typically be zero-copy for numeric
|
||||
types. Conversion of strings, lists, and structs will require creating python
|
||||
@@ -977,6 +995,8 @@ class Permutation:
|
||||
elif format == "arrow":
|
||||
return self.with_transform(Transforms.arrow2arrow)
|
||||
elif format == "torch":
|
||||
return self.with_transform(batch_to_tensor_dict)
|
||||
elif format == "torch_row":
|
||||
return self.with_transform(batch_to_tensor_rows)
|
||||
elif format == "torch_col":
|
||||
return self.with_transform(batch_to_tensor)
|
||||
|
||||
+336
-81
@@ -15,10 +15,12 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
import deprecation
|
||||
@@ -39,15 +41,21 @@ from .expr import Expr
|
||||
from .rerankers.base import Reranker
|
||||
from .rerankers.rrf import RRFReranker
|
||||
from .rerankers.util import check_reranker_result
|
||||
from .schema import is_blob_like_field, schema_has_blob_field
|
||||
from .util import flatten_columns
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
_BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"bytes": "all_binary",
|
||||
"descriptions": "blobs_descriptions",
|
||||
}
|
||||
from ._blob import (
|
||||
BLOB_MODE_TO_HANDLING,
|
||||
FetchBlobsAsync,
|
||||
FetchBlobsSync,
|
||||
blob_auto_row_id_for_scan,
|
||||
blob_v2_projection_sources,
|
||||
finalize_blob_query_table,
|
||||
replace_v2_blob_columns_with_bytes,
|
||||
replace_v2_blob_columns_with_bytes_sync,
|
||||
supports_blob_auto_row_id,
|
||||
validate_blob_mode,
|
||||
)
|
||||
from .types import BlobMode, QueryProjection
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import sys
|
||||
@@ -73,25 +81,22 @@ if TYPE_CHECKING:
|
||||
T = TypeVar("T", bound="LanceModel")
|
||||
|
||||
|
||||
def _validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in _BLOB_MODE_TO_HANDLING:
|
||||
modes = ", ".join(repr(mode) for mode in _BLOB_MODE_TO_HANDLING)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
@runtime_checkable
|
||||
class _LanceScanner(Protocol):
|
||||
projected_schema: pa.Schema | None
|
||||
schema: pa.Schema | None
|
||||
|
||||
def to_pandas(self, blob_mode: BlobMode | None = ..., **kwargs) -> pd.DataFrame: ...
|
||||
|
||||
def _field_is_blob(field: pa.Field) -> bool:
|
||||
metadata = field.metadata or {}
|
||||
return metadata.get(b"lance-encoding:blob") == b"true" or (
|
||||
metadata.get("lance-encoding:blob") == "true"
|
||||
)
|
||||
def to_pyarrow(self): ...
|
||||
|
||||
def to_table(self) -> pa.Table: ...
|
||||
|
||||
def _schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return any(_field_is_blob(field) for field in schema)
|
||||
def to_reader(self): ...
|
||||
|
||||
|
||||
def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool:
|
||||
return blob_mode in _BLOB_MODE_TO_HANDLING and _schema_has_blob_field(schema)
|
||||
return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema)
|
||||
|
||||
|
||||
def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
|
||||
@@ -119,13 +124,28 @@ def _filter_to_sql(filter: Optional[Union[str, Expr]]) -> Optional[str]:
|
||||
return filter
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(
|
||||
columns: Optional[
|
||||
Union[
|
||||
List[str], List[Tuple[str, Union[str, Expr]]], Dict[str, Union[str, Expr]]
|
||||
]
|
||||
],
|
||||
) -> Dict[str, Any]:
|
||||
def _combine_where(
|
||||
existing: Optional[Union[str, Expr]], new: Union[str, Expr]
|
||||
) -> Union[str, Expr]:
|
||||
"""Combine a new filter with an existing one using a logical AND.
|
||||
|
||||
Calling ``where`` more than once composes the filters with AND instead of
|
||||
replacing the previous filter. Two :class:`~lancedb.expr.Expr` filters are
|
||||
combined as an expression; otherwise both filters are lowered to SQL strings
|
||||
and combined as SQL.
|
||||
"""
|
||||
if existing is None:
|
||||
return new
|
||||
existing_is_expr = isinstance(existing, Expr)
|
||||
new_is_expr = isinstance(new, Expr)
|
||||
if existing_is_expr and new_is_expr:
|
||||
return existing & new
|
||||
existing_sql = existing.to_sql() if existing_is_expr else existing
|
||||
new_sql = new.to_sql() if new_is_expr else new
|
||||
return f"({existing_sql}) AND ({new_sql})"
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
if columns is None:
|
||||
return {}
|
||||
if isinstance(columns, list):
|
||||
@@ -150,7 +170,11 @@ def _projection_to_scanner_kwargs(
|
||||
|
||||
|
||||
def _scanner_kwargs_for_query(
|
||||
query: Query, blob_mode: BlobMode, dataset: Optional[Any] = None
|
||||
query: Query,
|
||||
blob_mode: BlobMode,
|
||||
dataset: Optional[Any] = None,
|
||||
*,
|
||||
with_row_id: Optional[bool] = None,
|
||||
) -> Dict[str, Any]:
|
||||
fragments = _scanner_fragments_for_query(query, dataset)
|
||||
kwargs = {
|
||||
@@ -158,10 +182,10 @@ def _scanner_kwargs_for_query(
|
||||
"filter": _filter_to_sql(query.filter),
|
||||
"limit": query.limit,
|
||||
"offset": query.offset,
|
||||
"with_row_id": query.with_row_id,
|
||||
"with_row_id": with_row_id if with_row_id is not None else query.with_row_id,
|
||||
"with_row_address": query.with_row_address,
|
||||
"fast_search": query.fast_search,
|
||||
"blob_handling": _BLOB_MODE_TO_HANDLING[blob_mode],
|
||||
"blob_handling": BLOB_MODE_TO_HANDLING[blob_mode],
|
||||
"fragments": fragments,
|
||||
}
|
||||
return {key: value for key, value in kwargs.items() if value is not None}
|
||||
@@ -194,11 +218,11 @@ def _scanner_fragments_for_query(query: Query, dataset: Optional[Any]) -> Option
|
||||
def _ensure_lazy_blob_frame(
|
||||
df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode
|
||||
) -> "pd.DataFrame":
|
||||
if blob_mode != "lazy" or not _schema_has_blob_field(schema) or len(df) == 0:
|
||||
if blob_mode != "lazy" or not schema_has_blob_field(schema) or len(df) == 0:
|
||||
return df
|
||||
|
||||
for field in schema:
|
||||
if not _field_is_blob(field) or field.name not in df.columns:
|
||||
if not is_blob_like_field(field) or field.name not in df.columns:
|
||||
continue
|
||||
value = df[field.name].iloc[0]
|
||||
if value is not None and not hasattr(value, "readall"):
|
||||
@@ -208,7 +232,7 @@ def _ensure_lazy_blob_frame(
|
||||
return df
|
||||
|
||||
|
||||
def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||
def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
||||
if hasattr(scanner, "to_pyarrow"):
|
||||
reader = scanner.to_pyarrow()
|
||||
return reader.read_all()
|
||||
@@ -218,7 +242,9 @@ def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||
return reader.read_all()
|
||||
|
||||
|
||||
def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataFrame":
|
||||
def _scanner_to_pandas(
|
||||
scanner: _LanceScanner, blob_mode: BlobMode, **kwargs
|
||||
) -> pd.DataFrame:
|
||||
schema = getattr(scanner, "projected_schema", None)
|
||||
if schema is None:
|
||||
schema = getattr(scanner, "schema", None)
|
||||
@@ -239,13 +265,71 @@ def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataF
|
||||
return df
|
||||
|
||||
tbl = _scanner_to_table(scanner)
|
||||
if blob_mode == "lazy" and _schema_has_blob_field(tbl.schema):
|
||||
if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"the Lance scanner does not expose to_pandas"
|
||||
)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
|
||||
|
||||
def _finish_plain_scan_pandas(
|
||||
scanner: _LanceScanner,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsSync,
|
||||
strip_auto_row_id: bool,
|
||||
flatten: Optional[Union[int, bool]],
|
||||
**kwargs,
|
||||
) -> pd.DataFrame:
|
||||
if blob_sources:
|
||||
tbl = _scanner_to_table(scanner)
|
||||
tbl = replace_v2_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(tbl, flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
return tbl.to_pandas(**kwargs)
|
||||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
if strip_auto_row_id and "_rowid" in df.columns:
|
||||
return df.drop(columns=["_rowid"])
|
||||
return df
|
||||
|
||||
|
||||
async def _finish_plain_scan_pandas_async(
|
||||
scanner: _LanceScanner,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsAsync,
|
||||
strip_auto_row_id: bool,
|
||||
flatten: Optional[Union[int, bool]],
|
||||
**kwargs,
|
||||
) -> pd.DataFrame:
|
||||
if blob_sources:
|
||||
tbl = _scanner_to_table(scanner)
|
||||
tbl = await replace_v2_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(tbl, flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
return tbl.to_pandas(**kwargs)
|
||||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
if strip_auto_row_id and "_rowid" in df.columns:
|
||||
return df.drop(columns=["_rowid"])
|
||||
return df
|
||||
|
||||
|
||||
# Pydantic validation function for vector queries
|
||||
def ensure_vector_query(
|
||||
val: Any,
|
||||
@@ -653,7 +737,7 @@ class Query(pydantic.BaseModel):
|
||||
distance_type: Optional[str] = None
|
||||
|
||||
# which columns to return in the results (dict values may be str or Expr)
|
||||
columns: Optional[Union[List[str], Dict[str, Union[str, Expr]]]] = None
|
||||
columns: QueryProjection = None
|
||||
|
||||
# minimum number of IVF partitions to search
|
||||
#
|
||||
@@ -937,7 +1021,7 @@ class LanceQueryBuilder(ABC):
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
validate_blob_mode(blob_mode)
|
||||
output_schema = getattr(self, "output_schema", None)
|
||||
if output_schema is not None:
|
||||
schema = output_schema()
|
||||
@@ -996,6 +1080,11 @@ class LanceQueryBuilder(ABC):
|
||||
Execute the query and return the results as a pyarrow
|
||||
[RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html)
|
||||
|
||||
For v2 blob projections, ``to_batches`` keeps the auto ``_rowid``
|
||||
column visible so batch consumers can call ``fetch_blobs``. Use
|
||||
``to_arrow``, ``to_list``, or ``to_pandas`` if you want LanceDB to hide
|
||||
auto row ids in the final collected result.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
batch_size: int
|
||||
@@ -1148,8 +1237,13 @@ class LanceQueryBuilder(ABC):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = where
|
||||
self._where = _combine_where(self._where, where)
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
|
||||
@@ -1169,6 +1263,42 @@ class LanceQueryBuilder(ABC):
|
||||
self._with_row_id = with_row_id
|
||||
return self
|
||||
|
||||
def _user_requested_row_id(self) -> bool:
|
||||
return self._with_row_id is True
|
||||
|
||||
def _blob_auto_row_id_enabled(self) -> bool:
|
||||
if not supports_blob_auto_row_id(self._table):
|
||||
return False
|
||||
return blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
self._table.schema,
|
||||
self._columns,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
|
||||
def _scan_needs_row_id(self) -> bool:
|
||||
return self._user_requested_row_id() or self._blob_auto_row_id_enabled()
|
||||
|
||||
def _query_for_scan(self) -> Query:
|
||||
query = self.to_query_object()
|
||||
if self._scan_needs_row_id():
|
||||
query.with_row_id = True
|
||||
return query
|
||||
|
||||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
blob_paths = (
|
||||
blob_v2_projection_sources(self._table.schema, self._columns).keys()
|
||||
if blob_auto_row_id
|
||||
else ()
|
||||
)
|
||||
return finalize_blob_query_table(
|
||||
tbl,
|
||||
user_requested_row_id=self._user_requested_row_id(),
|
||||
blob_auto_row_id=blob_auto_row_id,
|
||||
blob_paths=blob_paths,
|
||||
)
|
||||
|
||||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||||
"""Set whether to return row addresses.
|
||||
|
||||
@@ -1345,13 +1475,29 @@ class LanceQueryBuilder(ABC):
|
||||
return None
|
||||
|
||||
dataset = self._table.to_lance()
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
blob_sources = (
|
||||
blob_v2_projection_sources(self._table.schema, query.columns)
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(
|
||||
query,
|
||||
"descriptions" if blob_sources else blob_mode,
|
||||
dataset,
|
||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||||
)
|
||||
)
|
||||
return _finish_plain_scan_pandas(
|
||||
scanner,
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def to_query_object(self) -> Query:
|
||||
@@ -1599,7 +1745,9 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
"""
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
"""
|
||||
@@ -1659,7 +1807,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
vector = self._query if isinstance(self._query, list) else self._query.tolist()
|
||||
if isinstance(vector[0], np.ndarray):
|
||||
vector = [v.tolist() for v in vector]
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
result_set = self._table._execute_query(
|
||||
query, batch_size=batch_size, timeout=timeout
|
||||
)
|
||||
@@ -1693,8 +1841,13 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = where
|
||||
self._where = _combine_where(self._where, where)
|
||||
if prefilter is not None:
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
@@ -1798,8 +1951,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1809,6 +1961,21 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
self._phrase_query = phrase_query
|
||||
return self
|
||||
|
||||
def _query_with_phrase_semantics(self) -> str | FullTextQuery:
|
||||
query = self._query
|
||||
if not self._phrase_query:
|
||||
return query
|
||||
if isinstance(query, str):
|
||||
if not query.startswith('"') or not query.endswith('"'):
|
||||
return f'"{query}"'
|
||||
return query
|
||||
if isinstance(query, PhraseQuery):
|
||||
return query
|
||||
raise TypeError(
|
||||
"phrase_query() requires a string or PhraseQuery, "
|
||||
f"got {type(query).__name__}"
|
||||
)
|
||||
|
||||
def fast_search(self) -> LanceFtsQueryBuilder:
|
||||
"""
|
||||
Skip a flat search of unindexed data. This will improve
|
||||
@@ -1833,7 +2000,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
full_text_query=FullTextSearchQuery(
|
||||
query=self._query, columns=self._fts_columns
|
||||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
||||
),
|
||||
offset=self._offset,
|
||||
fast_search=self._fast_search,
|
||||
@@ -1851,22 +2018,13 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
self._table._ensure_no_legacy_fts_index()
|
||||
|
||||
query = self._query
|
||||
if self._phrase_query:
|
||||
if isinstance(query, str):
|
||||
if not query.startswith('"') or not query.endswith('"'):
|
||||
self._query = f'"{query}"'
|
||||
elif isinstance(query, FullTextQuery) and not isinstance(
|
||||
query, PhraseQuery
|
||||
):
|
||||
raise TypeError("Please use PhraseQuery for phrase queries.")
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
results = self._table._execute_query(query, timeout=timeout)
|
||||
results = results.read_all()
|
||||
if self._reranker is not None:
|
||||
results = self._reranker.rerank_fts(self._query, results)
|
||||
check_reranker_result(results)
|
||||
return results
|
||||
return self._finalize_blob_query_table(results)
|
||||
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
@@ -1894,7 +2052,9 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
return Query(
|
||||
@@ -1916,7 +2076,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
) -> pa.RecordBatchReader:
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
||||
|
||||
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
||||
@@ -1988,14 +2148,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
return vector_query, text_query
|
||||
|
||||
def phrase_query(self, phrase_query: bool = None) -> LanceHybridQueryBuilder:
|
||||
def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
|
||||
"""Set whether to use phrase query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -2020,15 +2179,25 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
fts_results = fts_future.result()
|
||||
vector_results = vector_future.result()
|
||||
|
||||
return self._combine_hybrid_results(
|
||||
results = self._combine_hybrid_results(
|
||||
fts_results=fts_results,
|
||||
vector_results=vector_results,
|
||||
norm=self._norm,
|
||||
fts_query=self._fts_query._query,
|
||||
reranker=self._reranker,
|
||||
limit=self._limit,
|
||||
with_row_ids=self._with_row_id,
|
||||
with_row_ids=True,
|
||||
)
|
||||
return self._finish_hybrid_results(results)
|
||||
|
||||
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
|
||||
if self._user_requested_row_id():
|
||||
return results
|
||||
if self._blob_auto_row_id_enabled():
|
||||
return self._finalize_blob_query_table(results)
|
||||
if "_rowid" in results.column_names:
|
||||
return results.drop(["_rowid"])
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _combine_hybrid_results(
|
||||
@@ -2469,7 +2638,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._vector_query.ef(self._ef)
|
||||
if self._bypass_vector_index:
|
||||
self._vector_query.bypass_vector_index()
|
||||
if self._lower_bound or self._upper_bound:
|
||||
if self._lower_bound is not None or self._upper_bound is not None:
|
||||
self._vector_query.distance_range(
|
||||
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
||||
)
|
||||
@@ -2499,6 +2668,9 @@ class AsyncQueryBase(object):
|
||||
self._with_row_address = None
|
||||
self._fragments = None
|
||||
self._fragment_ids = None
|
||||
self._with_row_id = None
|
||||
self._blob_auto_row_id = False
|
||||
self._blob_paths: tuple[str, ...] = ()
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
"""
|
||||
@@ -2508,11 +2680,46 @@ class AsyncQueryBase(object):
|
||||
python and more easily serializable.
|
||||
"""
|
||||
query = Query.from_inner(self._inner.to_query_request())
|
||||
query.with_row_id = self._user_requested_row_id()
|
||||
query.with_row_address = self._with_row_address
|
||||
query.fragments = self._fragments
|
||||
query.fragment_ids = self._fragment_ids
|
||||
return query
|
||||
|
||||
def _user_requested_row_id(self) -> bool:
|
||||
return self._with_row_id is True
|
||||
|
||||
def _blob_auto_row_id_enabled(self) -> bool:
|
||||
return self._blob_auto_row_id
|
||||
|
||||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||||
return finalize_blob_query_table(
|
||||
tbl,
|
||||
user_requested_row_id=self._user_requested_row_id(),
|
||||
blob_auto_row_id=self._blob_auto_row_id_enabled(),
|
||||
blob_paths=self._blob_paths,
|
||||
)
|
||||
|
||||
async def _maybe_add_blob_row_id(self) -> None:
|
||||
if self._table is None or not supports_blob_auto_row_id(self._table):
|
||||
self._blob_auto_row_id = False
|
||||
self._blob_paths = ()
|
||||
return
|
||||
|
||||
req = self._inner.to_query_request()
|
||||
schema = await self._table.schema()
|
||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
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._inner.with_row_id()
|
||||
|
||||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||||
"""
|
||||
Return only the specified columns.
|
||||
@@ -2565,6 +2772,7 @@ class AsyncQueryBase(object):
|
||||
"""
|
||||
Include the _rowid column in the results.
|
||||
"""
|
||||
self._with_row_id = True
|
||||
self._inner.with_row_id()
|
||||
return self
|
||||
|
||||
@@ -2611,6 +2819,7 @@ class AsyncQueryBase(object):
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
"""
|
||||
await self._maybe_add_blob_row_id()
|
||||
return AsyncRecordBatchReader(
|
||||
await self._inner.execute(
|
||||
max_batch_length=max_batch_length, timeout=timeout
|
||||
@@ -2641,8 +2850,8 @@ class AsyncQueryBase(object):
|
||||
complete within the specified time, an error will be raised.
|
||||
"""
|
||||
batch_iter = await self.to_batches(timeout=timeout)
|
||||
return pa.Table.from_batches(
|
||||
await batch_iter.read_all(), schema=batch_iter.schema
|
||||
return self._finalize_blob_query_table(
|
||||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
||||
)
|
||||
|
||||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
@@ -2709,7 +2918,7 @@ class AsyncQueryBase(object):
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
validate_blob_mode(blob_mode)
|
||||
if hasattr(self._inner, "output_schema"):
|
||||
schema = await self.output_schema()
|
||||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||||
@@ -2750,14 +2959,36 @@ class AsyncQueryBase(object):
|
||||
if not _query_is_plain_scan(query):
|
||||
return None
|
||||
|
||||
schema = await self._table.schema()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
query.columns,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
blob_sources = (
|
||||
blob_v2_projection_sources(schema, query.columns)
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
dataset = await self._table._to_lance()
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
**_scanner_kwargs_for_query(
|
||||
query,
|
||||
"descriptions" if blob_sources else blob_mode,
|
||||
dataset,
|
||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||||
)
|
||||
)
|
||||
return await _finish_plain_scan_pandas_async(
|
||||
scanner,
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
|
||||
async def to_polars(
|
||||
self,
|
||||
@@ -2894,6 +3125,9 @@ class AsyncStandardQuery(AsyncQueryBase):
|
||||
|
||||
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.
|
||||
"""
|
||||
if isinstance(predicate, Expr):
|
||||
self._inner.where_expr(predicate._inner)
|
||||
@@ -3539,9 +3773,24 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
||||
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||||
|
||||
# save the row ID choice that was made on the query builder and force it
|
||||
# to actually fetch the row ids because we need this for reranking
|
||||
with_row_ids = self._inner.get_with_row_id()
|
||||
req = fts_query._inner.to_query_request()
|
||||
blob_auto_row_id = False
|
||||
blob_paths: tuple[str, ...] = ()
|
||||
if self._table is not None and supports_blob_auto_row_id(self._table):
|
||||
schema = await self._table.schema()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if blob_auto_row_id:
|
||||
blob_paths = tuple(
|
||||
blob_v2_projection_sources(schema, req.select).keys()
|
||||
)
|
||||
self._blob_auto_row_id = blob_auto_row_id
|
||||
self._blob_paths = blob_paths
|
||||
|
||||
fts_query.with_row_id()
|
||||
vec_query.with_row_id()
|
||||
|
||||
@@ -3557,8 +3806,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query=fts_query.get_query(),
|
||||
reranker=self._reranker,
|
||||
limit=self._inner.get_limit(),
|
||||
with_row_ids=with_row_ids,
|
||||
with_row_ids=True,
|
||||
)
|
||||
if (
|
||||
not self._user_requested_row_id()
|
||||
and not blob_auto_row_id
|
||||
and "_rowid" in result.column_names
|
||||
):
|
||||
result = result.drop(["_rowid"])
|
||||
|
||||
return AsyncRecordBatchReader(result, max_batch_length=max_batch_length)
|
||||
|
||||
|
||||
@@ -9,6 +9,7 @@ from typing import List, Optional
|
||||
from lancedb import __version__
|
||||
|
||||
from .header import HeaderProvider
|
||||
from .oauth import OAuthConfig, OAuthFlowType
|
||||
|
||||
__all__ = [
|
||||
"TimeoutConfig",
|
||||
@@ -16,6 +17,8 @@ __all__ = [
|
||||
"TlsConfig",
|
||||
"ClientConfig",
|
||||
"HeaderProvider",
|
||||
"OAuthConfig",
|
||||
"OAuthFlowType",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
class OAuthFlowType(str, Enum):
|
||||
"""OAuth authentication flow types."""
|
||||
|
||||
CLIENT_CREDENTIALS = "client_credentials"
|
||||
"""Client Credentials grant (service-to-service / M2M)."""
|
||||
|
||||
AZURE_MANAGED_IDENTITY = "azure_managed_identity"
|
||||
"""Azure Managed Identity via IMDS."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class OAuthConfig:
|
||||
"""OAuth configuration for LanceDB authentication.
|
||||
|
||||
All token acquisition and refresh is handled in the Rust layer.
|
||||
This config is passed through to Rust via PyO3.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
issuer_url : str
|
||||
OIDC issuer URL or OAuth authority URL.
|
||||
For Azure: ``https://login.microsoftonline.com/{tenant_id}/v2.0``
|
||||
client_id : str
|
||||
Application / Client ID.
|
||||
scopes : List[str]
|
||||
OAuth scopes to request.
|
||||
For Azure managed identity, exactly one scope or resource is required.
|
||||
For example: ``["api://{app_id}/.default"]``
|
||||
flow : OAuthFlowType
|
||||
Authentication flow to use. Default: CLIENT_CREDENTIALS.
|
||||
client_secret : Optional[str]
|
||||
Client secret (required for CLIENT_CREDENTIALS).
|
||||
managed_identity_client_id : Optional[str]
|
||||
Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
|
||||
refresh_buffer_secs : Optional[int]
|
||||
Seconds before expiry to trigger proactive refresh (default: 300).
|
||||
Keep this well below the token TTL; if it is greater than or equal to
|
||||
the TTL, each request refreshes the token.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Client Credentials (service-to-service):
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... client_secret="secret",
|
||||
... scopes=["api://lancedb-api/.default"],
|
||||
... )
|
||||
|
||||
Azure Managed Identity:
|
||||
|
||||
>>> config = OAuthConfig(
|
||||
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
|
||||
... client_id="app-id",
|
||||
... scopes=["api://lancedb-api/.default"],
|
||||
... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY,
|
||||
... )
|
||||
"""
|
||||
|
||||
issuer_url: str
|
||||
client_id: str
|
||||
scopes: List[str]
|
||||
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
|
||||
client_secret: Optional[str] = field(default=None, repr=False)
|
||||
managed_identity_client_id: Optional[str] = None
|
||||
refresh_buffer_secs: Optional[int] = None
|
||||
@@ -28,6 +28,7 @@ from lancedb._lancedb import (
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
IndexConfig,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
@@ -244,6 +245,23 @@ class RemoteTable(Table):
|
||||
"""List all the indices on the table"""
|
||||
return LOOP.run(self._table.list_indices())
|
||||
|
||||
def tokenize(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
column: Optional[str] = None,
|
||||
index_name: Optional[str] = None,
|
||||
) -> Iterable[FtsToken]:
|
||||
"""Tokenize a query using the tokenizer configured on an FTS index.
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata, so the same tokenizer
|
||||
model files must exist locally.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.tokenize(query, column=column, index_name=index_name)
|
||||
)
|
||||
|
||||
def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]:
|
||||
"""List all the stats of a specified index"""
|
||||
return LOOP.run(self._table.index_stats(index_uuid))
|
||||
@@ -912,6 +930,10 @@ class RemoteTable(Table):
|
||||
"""Not supported on LanceDB Cloud."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""No-op on LanceDB Cloud (no local shard writers)."""
|
||||
return LOOP.run(self._table.close_lsm_writers())
|
||||
@@ -990,6 +1012,19 @@ class RemoteTable(Table):
|
||||
"migrate_v2_manifest_paths() is not supported on the LanceDB Cloud"
|
||||
)
|
||||
|
||||
def blob_columns(self) -> list[str]:
|
||||
raise NotImplementedError(
|
||||
"blob_columns() is not yet supported on the LanceDB Cloud"
|
||||
)
|
||||
|
||||
def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray:
|
||||
raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud")
|
||||
|
||||
def fetch_blob_files(self, column: str, row_ids):
|
||||
raise NotImplementedError(
|
||||
"fetch_blob_files() is not supported on LanceDB Cloud"
|
||||
)
|
||||
|
||||
def head(self, n=5) -> pa.Table:
|
||||
"""
|
||||
Return the first `n` rows of the table.
|
||||
|
||||
@@ -12,6 +12,7 @@ from .rrf import RRFReranker
|
||||
from .mrr import MRRReranker
|
||||
from .answerdotai import AnswerdotaiRerankers
|
||||
from .voyageai import VoyageAIReranker
|
||||
from .watsonx import WatsonxReranker
|
||||
|
||||
__all__ = [
|
||||
"Reranker",
|
||||
@@ -25,4 +26,5 @@ __all__ = [
|
||||
"AnswerdotaiRerankers",
|
||||
"VoyageAIReranker",
|
||||
"MRRReranker",
|
||||
"WatsonxReranker",
|
||||
]
|
||||
|
||||
@@ -156,9 +156,16 @@ class MRRReranker(Reranker):
|
||||
reciprocal_rank = 1.0 / rank
|
||||
mrr_score_map[result_id].append(reciprocal_rank)
|
||||
|
||||
# MRR averages the reciprocal rank across *all* ranking systems, treating
|
||||
# a system in which a document does not appear as a reciprocal rank of 0.
|
||||
# We therefore divide by the total number of systems, not by the number of
|
||||
# systems the document happens to appear in -- otherwise a document found
|
||||
# by a single ranking would outrank one ranked highly by every system,
|
||||
# defeating the purpose of fusing the rankings.
|
||||
num_systems = len(vector_results)
|
||||
final_mrr_scores = {}
|
||||
for result_id, reciprocal_ranks in mrr_score_map.items():
|
||||
mean_rr = np.mean(reciprocal_ranks)
|
||||
mean_rr = float(np.sum(reciprocal_ranks)) / num_systems
|
||||
final_mrr_scores[result_id] = mean_rr
|
||||
|
||||
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import os
|
||||
from functools import cached_property
|
||||
from typing import Dict, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from ..util import attempt_import_or_raise
|
||||
from .base import Reranker
|
||||
|
||||
DEFAULT_WATSONX_URL = "https://us-south.ml.cloud.ibm.com"
|
||||
|
||||
|
||||
class WatsonxReranker(Reranker):
|
||||
"""
|
||||
Reranks the results using the IBM watsonx.ai Rerank API.
|
||||
|
||||
Uses the ``ibm_watsonx_ai`` SDK (``Rerank.generate``) under the hood.
|
||||
|
||||
API Docs:
|
||||
https://cloud.ibm.com/docs/apis/watsonx-ai#text-rerank
|
||||
|
||||
Supported rerank models:
|
||||
https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#rerank
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model_name : str, default "cross-encoder/ms-marco-minilm-l-12-v2"
|
||||
The ID of the rerank model to use.
|
||||
column : str, default "text"
|
||||
The name of the column to use as input to the reranker.
|
||||
top_n : int, optional
|
||||
Return only the top-n results. If ``None``, all results are returned.
|
||||
return_score : str, default "relevance"
|
||||
Options are ``"relevance"`` or ``"all"``.
|
||||
api_key : str, optional
|
||||
IBM Cloud API key. Falls back to the ``WATSONX_API_KEY`` environment
|
||||
variable when not provided.
|
||||
project_id : str, optional
|
||||
watsonx.ai project ID. Falls back to the ``WATSONX_PROJECT_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
``space_id`` — exactly one must be supplied.
|
||||
space_id : str, optional
|
||||
watsonx.ai deployment space ID. Falls back to the ``WATSONX_SPACE_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
``project_id`` — exactly one must be supplied.
|
||||
url : str, optional
|
||||
watsonx.ai service URL. Defaults to
|
||||
``"https://us-south.ml.cloud.ibm.com"``.
|
||||
truncate_input_tokens : int, optional
|
||||
Truncate each input to this many tokens before scoring. Passed
|
||||
directly to the ``parameters`` dict of ``Rerank.generate``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "cross-encoder/ms-marco-minilm-l-12-v2",
|
||||
column: str = "text",
|
||||
top_n: Optional[int] = None,
|
||||
return_score: str = "relevance",
|
||||
api_key: Optional[str] = None,
|
||||
project_id: Optional[str] = None,
|
||||
space_id: Optional[str] = None,
|
||||
url: Optional[str] = None,
|
||||
truncate_input_tokens: Optional[int] = None,
|
||||
):
|
||||
super().__init__(return_score)
|
||||
self.model_name = model_name
|
||||
self.column = column
|
||||
self.top_n = top_n
|
||||
self.api_key = api_key
|
||||
self.project_id = project_id
|
||||
self.space_id = space_id
|
||||
self.url = url
|
||||
self.truncate_input_tokens = truncate_input_tokens
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"WatsonxReranker(model_name={self.model_name})"
|
||||
|
||||
@cached_property
|
||||
def _client(self):
|
||||
ibm_watsonx_ai = attempt_import_or_raise("ibm_watsonx_ai")
|
||||
ibm_watsonx_ai_foundation_models = attempt_import_or_raise(
|
||||
"ibm_watsonx_ai.foundation_models"
|
||||
)
|
||||
|
||||
# --- credentials ---
|
||||
api_key = self.api_key or os.environ.get("WATSONX_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"WATSONX_API_KEY not set. Either set it in your environment or "
|
||||
"pass it as `api_key` argument to WatsonxReranker."
|
||||
)
|
||||
credentials = ibm_watsonx_ai.Credentials(
|
||||
api_key=api_key,
|
||||
url=self.url or DEFAULT_WATSONX_URL,
|
||||
)
|
||||
|
||||
# --- project_id / space_id (exactly one required) ---
|
||||
project_id = self.project_id or os.environ.get("WATSONX_PROJECT_ID")
|
||||
space_id = self.space_id or os.environ.get("WATSONX_SPACE_ID")
|
||||
|
||||
if project_id and space_id:
|
||||
raise ValueError("Provide either `project_id` or `space_id`, not both.")
|
||||
if not project_id and not space_id:
|
||||
raise ValueError(
|
||||
"Either WATSONX_PROJECT_ID or WATSONX_SPACE_ID must be set. "
|
||||
"Pass one as an argument to WatsonxReranker or set the corresponding "
|
||||
"environment variable."
|
||||
)
|
||||
|
||||
kwargs: Dict = dict(model_id=self.model_name, credentials=credentials)
|
||||
if project_id:
|
||||
kwargs["project_id"] = project_id
|
||||
else:
|
||||
kwargs["space_id"] = space_id
|
||||
|
||||
return ibm_watsonx_ai_foundation_models.Rerank(**kwargs)
|
||||
|
||||
def _build_params(self) -> Dict:
|
||||
"""Build the ``parameters`` dict forwarded to ``Rerank.generate``."""
|
||||
return_options: Dict = {"inputs": True}
|
||||
if self.top_n is not None:
|
||||
return_options["top_n"] = self.top_n
|
||||
params: Dict = {"return_options": return_options}
|
||||
if self.truncate_input_tokens is not None:
|
||||
params["truncate_input_tokens"] = self.truncate_input_tokens
|
||||
return params
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str) -> pa.Table:
|
||||
result_set = self._handle_empty_results(result_set)
|
||||
if len(result_set) == 0:
|
||||
return result_set
|
||||
|
||||
docs = result_set[self.column].to_pylist()
|
||||
response = self._client.generate(
|
||||
query=query,
|
||||
inputs=docs,
|
||||
params=self._build_params(),
|
||||
)
|
||||
results = response["results"]
|
||||
|
||||
indices, scores = zip(
|
||||
*[(result["index"], result["score"]) for result in results]
|
||||
)
|
||||
result_set = result_set.take(list(indices))
|
||||
result_set = result_set.append_column(
|
||||
"_relevance_score", pa.array(scores, type=pa.float32())
|
||||
)
|
||||
return result_set
|
||||
|
||||
def rerank_hybrid(
|
||||
self,
|
||||
query: str,
|
||||
vector_results: pa.Table,
|
||||
fts_results: pa.Table,
|
||||
) -> pa.Table:
|
||||
if self.score == "all":
|
||||
combined_results = self._merge_and_keep_scores(vector_results, fts_results)
|
||||
else:
|
||||
combined_results = self.merge_results(vector_results, fts_results)
|
||||
combined_results = self._rerank(combined_results, query)
|
||||
if self.score == "relevance":
|
||||
combined_results = self._keep_relevance_score(combined_results)
|
||||
return combined_results
|
||||
|
||||
def rerank_vector(self, query: str, vector_results: pa.Table) -> pa.Table:
|
||||
vector_results = self._rerank(vector_results, query)
|
||||
if self.score == "relevance":
|
||||
vector_results = vector_results.drop_columns(["_distance"])
|
||||
return vector_results
|
||||
|
||||
def rerank_fts(self, query: str, fts_results: pa.Table) -> pa.Table:
|
||||
fts_results = self._rerank(fts_results, query)
|
||||
if self.score == "relevance":
|
||||
fts_results = fts_results.drop_columns(["_score"])
|
||||
return fts_results
|
||||
@@ -2,10 +2,134 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
"""Schema related utilities."""
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
_BLOB_EXTENSION_NAME = "lance.blob.v2"
|
||||
_BLOB_V1_KEY = "lance-encoding:blob"
|
||||
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
|
||||
|
||||
|
||||
class BlobType(pa.ExtensionType):
|
||||
"""PyArrow extension type for a Lance blob v2 column.
|
||||
|
||||
Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files`
|
||||
for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
storage_type = pa.struct(
|
||||
[
|
||||
pa.field("data", pa.large_binary(), nullable=True),
|
||||
pa.field("uri", pa.utf8(), nullable=True),
|
||||
pa.field("position", pa.uint64(), nullable=True),
|
||||
pa.field("size", pa.uint64(), nullable=True),
|
||||
]
|
||||
)
|
||||
super().__init__(storage_type, _BLOB_EXTENSION_NAME)
|
||||
|
||||
def __arrow_ext_serialize__(self) -> bytes:
|
||||
return b""
|
||||
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(
|
||||
cls, storage_type: pa.DataType, serialized: bytes
|
||||
) -> "BlobType":
|
||||
return cls()
|
||||
|
||||
def __reduce__(self):
|
||||
# Ensure pickle round-trips on older pyarrow (apache/arrow#35599).
|
||||
return type(self).__arrow_ext_deserialize__, (
|
||||
self.storage_type,
|
||||
self.__arrow_ext_serialize__(),
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
pa.register_extension_type(BlobType()) # type: ignore[arg-type]
|
||||
except pa.ArrowKeyError:
|
||||
pass
|
||||
|
||||
|
||||
def _metadata_value(metadata: dict, key: str):
|
||||
return metadata.get(key.encode()) or metadata.get(key)
|
||||
|
||||
|
||||
def _metadata_marks_blob_v2(metadata: dict) -> bool:
|
||||
if not metadata:
|
||||
return False
|
||||
|
||||
extension_name = _metadata_value(metadata, _ARROW_EXT_NAME_KEY)
|
||||
return extension_name in (_BLOB_EXTENSION_NAME, _BLOB_EXTENSION_NAME.encode())
|
||||
|
||||
|
||||
def _metadata_marks_legacy_blob(metadata: dict) -> bool:
|
||||
if not metadata:
|
||||
return False
|
||||
|
||||
return _metadata_value(metadata, _BLOB_V1_KEY) in ("true", b"true")
|
||||
|
||||
|
||||
def is_blob_v2_field(field: pa.Field) -> bool:
|
||||
"""Return True if `field` declares a blob v2 extension column."""
|
||||
field_type = field.type
|
||||
if (
|
||||
isinstance(field_type, pa.ExtensionType)
|
||||
and field_type.extension_name == _BLOB_EXTENSION_NAME
|
||||
):
|
||||
return True
|
||||
return _metadata_marks_blob_v2(field.metadata or {})
|
||||
|
||||
|
||||
def is_blob_like_field(field: pa.Field) -> bool:
|
||||
"""Blob detection for ``to_pandas(blob_mode=...)`` and scanner paths only.
|
||||
|
||||
Matches v2 extension fields on table schema, legacy ``lance-encoding:blob``
|
||||
storage columns, and v2 query descriptor fields (the engine tags those with
|
||||
the same metadata). Not used for fetch or auto ``_rowid``.
|
||||
"""
|
||||
return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {})
|
||||
|
||||
|
||||
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]:
|
||||
paths: list[str] = []
|
||||
|
||||
def walk(fields, prefix: str) -> None:
|
||||
for field in fields:
|
||||
path = f"{prefix}.{field.name}" if prefix else field.name
|
||||
if is_blob(field):
|
||||
paths.append(path)
|
||||
elif pa.types.is_struct(field.type):
|
||||
walk(field.type, path)
|
||||
elif (
|
||||
pa.types.is_list(field.type)
|
||||
or pa.types.is_large_list(field.type)
|
||||
or pa.types.is_fixed_size_list(field.type)
|
||||
):
|
||||
walk([field.type.value_field], path)
|
||||
|
||||
walk(schema, "")
|
||||
return paths
|
||||
|
||||
|
||||
def blob_column_paths(schema: pa.Schema) -> list[str]:
|
||||
"""Dotted paths of blob-like columns (v2 extension or legacy metadata)."""
|
||||
return _collect_blob_paths(schema, is_blob_like_field)
|
||||
|
||||
|
||||
def blob_v2_column_paths(schema: pa.Schema) -> list[str]:
|
||||
return _collect_blob_paths(schema, is_blob_v2_field)
|
||||
|
||||
|
||||
def schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return bool(blob_column_paths(schema))
|
||||
|
||||
|
||||
def blob(name: str, nullable: bool = True) -> pa.Field:
|
||||
"""Create a Lance blob v2 column field."""
|
||||
return pa.field(name, BlobType(), nullable=nullable)
|
||||
|
||||
|
||||
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||
"""A help function to create a vector type.
|
||||
|
||||
@@ -0,0 +1,607 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Elastic streaming dataloader for PyTorch.
|
||||
|
||||
Provides StreamingDataset, a PyTorch IterableDataset that guarantees:
|
||||
|
||||
- **Elastic determinism**: for a fixed (num_splits, shuffle_seed, epoch) the set
|
||||
of samples that forms each global training step is identical regardless of
|
||||
world_size or num_workers.
|
||||
- **Resumability**: state_dict / load_state_dict capture per-split consumption
|
||||
counts so training can resume from an exact mid-epoch position even when the
|
||||
distributed topology changes between runs.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from multiprocessing import RawArray
|
||||
from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
from torch.utils.data import IterableDataset, get_worker_info
|
||||
|
||||
from .permutation import (
|
||||
Permutation,
|
||||
Transforms,
|
||||
permutation_builder,
|
||||
_table_from_pickle_state,
|
||||
_table_to_pickle_state,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Multiplier used to combine shuffle_seed and epoch into a single permutation
|
||||
# seed. Chosen to be a large prime so different (seed, epoch) pairs produce
|
||||
# distinct seeds for any practically encountered epoch count.
|
||||
_EPOCH_PRIME = 100003
|
||||
|
||||
DEFAULT_READ_BATCH_SIZE = 64
|
||||
DEFAULT_PREFETCH_BATCHES = 4
|
||||
|
||||
|
||||
class StreamingDataset(IterableDataset):
|
||||
"""An elastic, resumable PyTorch IterableDataset backed by a LanceDB table.
|
||||
|
||||
The table is partitioned into ``num_splits`` fixed splits using a
|
||||
deterministic random shuffle controlled by ``shuffle_seed`` and ``epoch``.
|
||||
Each rank is assigned a contiguous block of splits, and within a rank each
|
||||
DataLoader worker is assigned a contiguous sub-block. Samples are yielded
|
||||
by round-robining over the assigned splits, one sample per split per cycle.
|
||||
|
||||
Internally ``__iter__`` runs a two-stage pipeline:
|
||||
|
||||
- **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches``
|
||||
workers fetches raw ``RecordBatch`` objects from LanceDB in parallel
|
||||
across all splits and places them in a per-split raw-batch queue.
|
||||
- **Stage 2 (transform)**: a second thread pool with ``os.cpu_count()``
|
||||
workers picks up raw batches, applies the transform, and places the
|
||||
results in a per-split cooked-row queue.
|
||||
|
||||
The main thread round-robins over the cooked queues, yielding one row per
|
||||
split per cycle.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
table:
|
||||
LanceDB table to stream from.
|
||||
num_splits:
|
||||
Number of fixed splits to partition the table into. Must be divisible
|
||||
by ``world_size``. When used with DataLoader workers it must also be
|
||||
divisible by ``world_size * num_workers``. Defaults to ``world_size``.
|
||||
If the row count (after any ``filter``) is not evenly divisible by
|
||||
``num_splits``, the surplus rows — at most ``num_splits - 1`` per epoch
|
||||
— are silently dropped to keep all splits the same length.
|
||||
shuffle:
|
||||
Whether to randomly assign rows to splits. When ``True`` (the
|
||||
default) rows are shuffled using ``shuffle_seed`` and ``epoch``.
|
||||
When ``False`` rows are divided into splits sequentially in storage
|
||||
order, which can be useful for deterministic debugging or evaluation.
|
||||
shuffle_seed:
|
||||
Base seed for the random permutation. Combined with ``epoch`` so
|
||||
each epoch produces a different ordering. Pass ``None`` to generate
|
||||
a random seed at construction time.
|
||||
epoch:
|
||||
Current training epoch. Combined with ``shuffle_seed`` so that each
|
||||
epoch produces a different sample ordering.
|
||||
rank:
|
||||
This process's rank in the distributed training group.
|
||||
world_size:
|
||||
Total number of processes in the distributed training group.
|
||||
read_batch_size:
|
||||
Number of rows fetched from each split in a single ``take_offsets``
|
||||
call. Larger values amortise per-request overhead (critical on object
|
||||
storage) at the cost of higher memory usage per split buffer. Defaults
|
||||
to ``DEFAULT_READ_BATCH_SIZE`` (64).
|
||||
prefetch_batches:
|
||||
Number of I/O batches to keep in flight per split. Higher values
|
||||
overlap storage latency with transform and training compute at the cost
|
||||
of more memory and threads. Defaults to ``DEFAULT_PREFETCH_BATCHES``
|
||||
(4).
|
||||
columns:
|
||||
Optional list of column names to read. When set, only those columns
|
||||
are fetched from storage; all others are omitted. ``None`` (the
|
||||
default) reads every column.
|
||||
shuffle_clump_size:
|
||||
When set, rows are shuffled in contiguous groups of this size rather
|
||||
than individually. Larger clumps improve I/O locality (important on
|
||||
object storage) at the cost of reduced randomness. ``None`` (the
|
||||
default) shuffles rows individually.
|
||||
filter:
|
||||
Optional SQL filter expression (e.g. ``"label = 'dog'"``). Only rows
|
||||
that satisfy the predicate are included in the permutation. The filter
|
||||
is applied during permutation construction so split sizes reflect the
|
||||
filtered row count.
|
||||
transform:
|
||||
Optional callable applied to each ``pyarrow.RecordBatch`` before rows
|
||||
are yielded. Receives one batch at a time and must return an iterable
|
||||
whose length equals the number of rows in the batch. When ``None``
|
||||
(the default) rows are returned as plain Python dicts.
|
||||
worker_info_override:
|
||||
If set, used in place of ``torch.utils.data.get_worker_info()`` to
|
||||
determine the DataLoader worker assignment. Intended for unit tests
|
||||
that need to simulate multiple workers without spawning real processes.
|
||||
If both this and the real worker info are non-None a warning is logged
|
||||
and the override takes precedence.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
table,
|
||||
*,
|
||||
num_splits: Optional[int] = None,
|
||||
shuffle: bool = True,
|
||||
shuffle_seed: Optional[int] = 0,
|
||||
epoch: int = 0,
|
||||
rank: int = 0,
|
||||
world_size: int = 1,
|
||||
read_batch_size: int = DEFAULT_READ_BATCH_SIZE,
|
||||
prefetch_batches: int = DEFAULT_PREFETCH_BATCHES,
|
||||
columns: Optional[list[str]] = None,
|
||||
shuffle_clump_size: Optional[int] = None,
|
||||
filter: Optional[str] = None,
|
||||
transform: Optional[Callable] = None,
|
||||
connection_factory: Optional[Callable[[str], Any]] = None,
|
||||
worker_info_override=None,
|
||||
):
|
||||
super().__init__()
|
||||
if num_splits is None:
|
||||
num_splits = world_size
|
||||
if shuffle_seed is None:
|
||||
shuffle_seed = random.randrange(2**32)
|
||||
if num_splits % world_size != 0:
|
||||
raise ValueError(
|
||||
f"num_splits ({num_splits}) must be divisible by "
|
||||
f"world_size ({world_size})"
|
||||
)
|
||||
|
||||
self._table = table
|
||||
self._num_splits = num_splits
|
||||
self._shuffle = shuffle
|
||||
self._shuffle_seed = shuffle_seed
|
||||
self._epoch = epoch
|
||||
self._rank = rank
|
||||
self._world_size = world_size
|
||||
self._read_batch_size = read_batch_size
|
||||
self._prefetch_batches = prefetch_batches
|
||||
self._columns = columns
|
||||
self._shuffle_clump_size = shuffle_clump_size
|
||||
self._filter = filter
|
||||
self._transform = transform
|
||||
self._connection_factory = connection_factory
|
||||
self._worker_info_override = worker_info_override
|
||||
|
||||
# Live references to pipeline state, set only while __iter__ is running
|
||||
# in the same process. Used by the observability properties when the
|
||||
# DataLoader runs with num_workers=0.
|
||||
self._raw_batches_ref: Optional[list[deque]] = None
|
||||
self._cooked_ref: Optional[list[deque]] = None
|
||||
self._fetch_head_ref: Optional[list[int]] = None
|
||||
self._split_sizes_ref: Optional[list[int]] = None
|
||||
self._local_consumed_ref: Optional[list[int]] = None
|
||||
|
||||
# Shared-memory counters written by __iter__ (which may run in a
|
||||
# DataLoader worker process) and read by the observability properties
|
||||
# in the main process. RawArray is picklable via the forkserver
|
||||
# reduction protocol so it survives the dataset pickle round-trip.
|
||||
# Layout: [unscanned_rows, raw_rows, cooked_rows, consumed_rows,
|
||||
# bytes_loaded, fetch_time_us, transform_time_us]
|
||||
self._worker_stats: RawArray = RawArray(ctypes.c_int64, 7)
|
||||
|
||||
# Cumulative bytes of Arrow buffer data fetched across all iterations.
|
||||
self._bytes_loaded: int = 0
|
||||
# Cumulative seconds spent in LanceDB I/O and in transform functions.
|
||||
self._fetch_time: float = 0.0
|
||||
self._transform_time: float = 0.0
|
||||
|
||||
# Number of samples each split has already been consumed. At global
|
||||
# step boundaries all splits have consumed this many samples, so a
|
||||
# single scalar captures the topology-independent checkpoint state.
|
||||
self._resume_offset: int = 0
|
||||
|
||||
# Build the permutation table once, deterministically.
|
||||
builder = permutation_builder(table)
|
||||
if filter is not None:
|
||||
builder = builder.filter(filter)
|
||||
if shuffle:
|
||||
perm_seed = shuffle_seed + epoch * _EPOCH_PRIME
|
||||
self._perm_table = builder.split_random(
|
||||
fixed=num_splits, seed=perm_seed, clump_size=shuffle_clump_size
|
||||
).execute()
|
||||
else:
|
||||
self._perm_table = builder.split_sequential(fixed=num_splits).execute()
|
||||
|
||||
# Contiguous block of global split indices assigned to this rank.
|
||||
splits_per_rank = num_splits // world_size
|
||||
rank_start = rank * splits_per_rank
|
||||
self._rank_splits: list[int] = list(
|
||||
range(rank_start, rank_start + splits_per_rank)
|
||||
)
|
||||
|
||||
def _resolve_my_splits(self) -> list[int]:
|
||||
"""Return the split indices this instance should read in __iter__."""
|
||||
torch_worker_info = get_worker_info()
|
||||
if self._worker_info_override is not None:
|
||||
if torch_worker_info is not None:
|
||||
logger.warning(
|
||||
"worker_info_override is set but get_worker_info() also returned a "
|
||||
"non-None value; ignoring the real torch worker info and using the "
|
||||
"override instead. This may lead to duplicated or incorrect data "
|
||||
"from the dataset."
|
||||
)
|
||||
worker_info = self._worker_info_override
|
||||
else:
|
||||
worker_info = torch_worker_info
|
||||
|
||||
if worker_info is None:
|
||||
return self._rank_splits
|
||||
|
||||
num_workers: int = worker_info.num_workers
|
||||
worker_id: int = worker_info.id
|
||||
n_rank_splits = len(self._rank_splits)
|
||||
if n_rank_splits % num_workers != 0:
|
||||
raise ValueError(
|
||||
f"Number of rank splits ({n_rank_splits}) must be divisible by "
|
||||
f"num_workers ({num_workers})"
|
||||
)
|
||||
splits_per_worker = n_rank_splits // num_workers
|
||||
start = worker_id * splits_per_worker
|
||||
return self._rank_splits[start : start + splits_per_worker]
|
||||
|
||||
def __iter__(self) -> Iterator[dict[str, Any]]:
|
||||
if self._raw_batches_ref is not None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataset does not support concurrent iteration. "
|
||||
"Only one active iterator per dataset instance is allowed."
|
||||
)
|
||||
my_splits = self._resolve_my_splits()
|
||||
if not my_splits:
|
||||
return
|
||||
|
||||
# Set identity transform on each Permutation so __getitems__ returns
|
||||
# the raw RecordBatch. Stage 2 applies the real transform.
|
||||
permutations: list[Permutation] = []
|
||||
for split_idx in my_splits:
|
||||
perm = Permutation.from_tables(
|
||||
self._table, self._perm_table, split=split_idx
|
||||
)
|
||||
if self._columns is not None:
|
||||
perm = perm.select_columns(self._columns)
|
||||
perm = perm.with_transform(lambda batch: batch)
|
||||
if self._resume_offset > 0:
|
||||
perm = perm.with_skip(self._resume_offset)
|
||||
permutations.append(perm)
|
||||
|
||||
n = len(permutations)
|
||||
split_sizes = [perm.num_rows for perm in permutations]
|
||||
initial_offset = self._resume_offset
|
||||
local_consumed = [0] * n
|
||||
|
||||
batch_size = self._read_batch_size
|
||||
max_prefetch = self._prefetch_batches
|
||||
cpu_workers = os.cpu_count() or 1
|
||||
final_transform = (
|
||||
self._transform if self._transform is not None else Transforms.arrow2python
|
||||
)
|
||||
|
||||
# Per-split pipeline state.
|
||||
fetch_head = [0] * n
|
||||
io_pending = [deque() for _ in range(n)] # Future[RecordBatch]
|
||||
raw_batches = [deque() for _ in range(n)] # RecordBatch — fetched, awaiting tx
|
||||
tx_pending = [deque() for _ in range(n)] # Future[list[Any]]
|
||||
cooked = [deque() for _ in range(n)] # rows ready to yield
|
||||
|
||||
# Limit simultaneous transforms to cpu_workers across all splits.
|
||||
tx_semaphore = threading.Semaphore(cpu_workers)
|
||||
|
||||
# ── Stage 1 helpers ───────────────────────────────────────────────────
|
||||
|
||||
def _io_call(perm, indices):
|
||||
t0 = time.perf_counter()
|
||||
batch = perm.__getitems__(indices)
|
||||
self._bytes_loaded += batch.nbytes
|
||||
self._fetch_time += time.perf_counter() - t0
|
||||
return batch
|
||||
|
||||
def _submit_io(i: int) -> None:
|
||||
remaining = split_sizes[i] - fetch_head[i]
|
||||
if remaining <= 0:
|
||||
return
|
||||
fetch = min(batch_size, remaining)
|
||||
start = fetch_head[i]
|
||||
fetch_head[i] += fetch
|
||||
perm_i = permutations[i]
|
||||
indices = list(range(start, start + fetch))
|
||||
io_pending[i].append(io_pool.submit(_io_call, perm_i, indices))
|
||||
|
||||
def _fill_io(i: int) -> None:
|
||||
while len(io_pending[i]) < max_prefetch and fetch_head[i] < split_sizes[i]:
|
||||
_submit_io(i)
|
||||
|
||||
def _drain_io(i: int) -> None:
|
||||
"""Move completed I/O futures into raw_batches non-blockingly."""
|
||||
while io_pending[i] and io_pending[i][0].done():
|
||||
raw_batches[i].append(io_pending[i].popleft().result())
|
||||
|
||||
# ── Stage 2 helpers ───────────────────────────────────────────────────
|
||||
|
||||
def _tx_call_guarded(batch):
|
||||
try:
|
||||
t0 = time.perf_counter()
|
||||
result = final_transform(batch)
|
||||
self._transform_time += time.perf_counter() - t0
|
||||
return result
|
||||
finally:
|
||||
tx_semaphore.release()
|
||||
|
||||
def _try_submit_tx(i: int) -> None:
|
||||
"""Submit transforms for raw_batches[i] up to available capacity."""
|
||||
while raw_batches[i] and tx_semaphore.acquire(blocking=False):
|
||||
batch = raw_batches[i].popleft()
|
||||
tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch))
|
||||
|
||||
def _drain_tx(i: int) -> None:
|
||||
"""Move completed transform futures into cooked non-blockingly."""
|
||||
while tx_pending[i] and tx_pending[i][0].done():
|
||||
cooked[i].extend(tx_pending[i].popleft().result())
|
||||
|
||||
# ── Combined advance ──────────────────────────────────────────────────
|
||||
|
||||
def _advance(i: int) -> None:
|
||||
"""Non-blocking pipeline pump for split i."""
|
||||
_drain_io(i)
|
||||
_drain_tx(i)
|
||||
_try_submit_tx(i)
|
||||
_fill_io(i)
|
||||
|
||||
def _ensure_cooked(i: int) -> None:
|
||||
"""Ensure cooked[i] has at least one row, blocking if necessary."""
|
||||
_advance(i)
|
||||
while not cooked[i]:
|
||||
if tx_pending[i]:
|
||||
# Wait for the oldest in-flight transform.
|
||||
cooked[i].extend(tx_pending[i].popleft().result())
|
||||
_advance(i)
|
||||
elif raw_batches[i]:
|
||||
# Acquire a transform slot (may block briefly if all
|
||||
# cpu_workers are busy with other splits).
|
||||
tx_semaphore.acquire()
|
||||
batch = raw_batches[i].popleft()
|
||||
tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch))
|
||||
elif io_pending[i]:
|
||||
# Block on the oldest in-flight I/O fetch.
|
||||
raw_batches[i].append(io_pending[i].popleft().result())
|
||||
_advance(i)
|
||||
else:
|
||||
break # split exhausted
|
||||
|
||||
# ── Main loop ─────────────────────────────────────────────────────────
|
||||
|
||||
with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool:
|
||||
with ThreadPoolExecutor(max_workers=cpu_workers) as tx_pool:
|
||||
self._raw_batches_ref = raw_batches
|
||||
self._cooked_ref = cooked
|
||||
self._fetch_head_ref = fetch_head
|
||||
self._split_sizes_ref = split_sizes
|
||||
self._local_consumed_ref = local_consumed
|
||||
try:
|
||||
for i in range(n):
|
||||
_fill_io(i)
|
||||
|
||||
while True:
|
||||
# Stop when any split is exhausted (all exhaust
|
||||
# simultaneously: equal split sizes + round-robin).
|
||||
if any(local_consumed[i] >= split_sizes[i] for i in range(n)):
|
||||
break
|
||||
|
||||
for i in range(n):
|
||||
_ensure_cooked(i)
|
||||
row = cooked[i].popleft()
|
||||
local_consumed[i] += 1
|
||||
_advance(i)
|
||||
|
||||
# After the last split in each cycle: update the
|
||||
# global offset and refresh the shared-memory stats
|
||||
# so the main process can observe pipeline depth
|
||||
# even when __iter__ runs in a worker process.
|
||||
if i == n - 1:
|
||||
self._resume_offset = initial_offset + local_consumed[i]
|
||||
ws = self._worker_stats
|
||||
ws[0] = sum(
|
||||
split_sizes[j] - fetch_head[j] for j in range(n)
|
||||
)
|
||||
ws[1] = sum(
|
||||
batch.num_rows for q in raw_batches for batch in q
|
||||
)
|
||||
ws[2] = sum(len(q) for q in cooked)
|
||||
ws[3] = sum(local_consumed)
|
||||
ws[4] = self._bytes_loaded
|
||||
ws[5] = int(self._fetch_time * 1_000_000)
|
||||
ws[6] = int(self._transform_time * 1_000_000)
|
||||
|
||||
yield row
|
||||
finally:
|
||||
self._raw_batches_ref = None
|
||||
self._cooked_ref = None
|
||||
self._fetch_head_ref = None
|
||||
self._split_sizes_ref = None
|
||||
self._local_consumed_ref = None
|
||||
|
||||
@property
|
||||
def bytes_loaded(self) -> int:
|
||||
"""Cumulative bytes of raw Arrow buffer data fetched from storage.
|
||||
|
||||
Measured on the ``RecordBatch`` before any transform is applied, so
|
||||
the value reflects actual I/O rather than the size of transformed
|
||||
output. Accumulates across multiple iterations of the same dataset
|
||||
instance and is never reset automatically.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._bytes_loaded
|
||||
return int(self._worker_stats[4])
|
||||
|
||||
@property
|
||||
def fetch_time(self) -> float:
|
||||
"""Cumulative seconds spent waiting for data from LanceDB.
|
||||
|
||||
Measured per batch in the Stage 1 I/O threads as the total elapsed
|
||||
time of the ``take_offsets`` call. Accumulates across all splits and
|
||||
all iterations.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._fetch_time
|
||||
return self._worker_stats[5] / 1_000_000
|
||||
|
||||
@property
|
||||
def transform_time(self) -> float:
|
||||
"""Cumulative seconds spent applying the transform.
|
||||
|
||||
Measured per batch in the Stage 2 transform threads as the elapsed
|
||||
time inside the transform callable (or the default ``arrow2python``
|
||||
conversion when no transform is set). Accumulates across all splits
|
||||
and all iterations.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._transform_time
|
||||
return self._worker_stats[6] / 1_000_000
|
||||
|
||||
@property
|
||||
def raw_queue_depth(self) -> int:
|
||||
"""Number of raw rows waiting for a transform thread across all splits.
|
||||
|
||||
A persistently non-zero value means Stage 2 (transform) is the
|
||||
bottleneck: I/O is completing faster than transforms can consume
|
||||
batches. Returns 0 when not iterating.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return sum(batch.num_rows for q in self._raw_batches_ref for batch in q)
|
||||
return int(self._worker_stats[1])
|
||||
|
||||
@property
|
||||
def prefetch_queue_depth(self) -> int:
|
||||
"""Number of rows transformed and ready to yield across all splits.
|
||||
|
||||
Counts rows whose transform has completed and are sitting in memory
|
||||
waiting for the main thread — rows that can be handed off with no
|
||||
I/O or CPU wait. Returns 0 when not iterating.
|
||||
"""
|
||||
if self._cooked_ref is not None:
|
||||
return sum(len(q) for q in self._cooked_ref)
|
||||
return int(self._worker_stats[2])
|
||||
|
||||
@property
|
||||
def unscanned_rows(self) -> int:
|
||||
"""Number of rows not yet submitted to the I/O stage across all splits.
|
||||
|
||||
Decreases as the I/O stage submits fetch requests. When this reaches
|
||||
zero all data has been requested from storage (though it may not have
|
||||
arrived yet). Returns 0 when not iterating.
|
||||
"""
|
||||
if self._fetch_head_ref is not None:
|
||||
return sum(
|
||||
size - head
|
||||
for size, head in zip(self._split_sizes_ref, self._fetch_head_ref)
|
||||
)
|
||||
return int(self._worker_stats[0])
|
||||
|
||||
@property
|
||||
def consumed_rows(self) -> int:
|
||||
"""Number of rows already yielded to the caller across all splits.
|
||||
|
||||
Monotonically increases throughout iteration. Returns 0 when not
|
||||
iterating.
|
||||
"""
|
||||
if self._local_consumed_ref is not None:
|
||||
return sum(self._local_consumed_ref)
|
||||
return int(self._worker_stats[3])
|
||||
|
||||
def __getstate__(self):
|
||||
"""Support pickling for multi-worker DataLoader (forkserver / spawn).
|
||||
|
||||
The live LanceDB table object contains non-picklable connection state
|
||||
(sockets, Rust-backed PyO3 objects). If a ``connection_factory`` was
|
||||
supplied only the table name is serialised; the factory is called in
|
||||
the worker to reopen the connection without embedding any credentials.
|
||||
Without a factory the table's own picklable reopen state is captured
|
||||
via ``_table_to_pickle_state`` (mirrors the ``Permutation`` approach).
|
||||
"""
|
||||
state = self.__dict__.copy()
|
||||
# _table: replace with reconnect info (credentials must not be embedded).
|
||||
state["_table_name"] = self._table.name
|
||||
if self._connection_factory is not None:
|
||||
state["_table"] = None
|
||||
else:
|
||||
state["_table"] = _table_to_pickle_state(self._table)
|
||||
# _perm_table: always in-memory; serialise as Arrow data (mirrors
|
||||
# how Permutation.__getstate__ handles its permutation_table).
|
||||
state["_perm_table"] = (
|
||||
self._perm_table.name,
|
||||
self._perm_table.to_arrow(),
|
||||
)
|
||||
for key in (
|
||||
"_raw_batches_ref",
|
||||
"_cooked_ref",
|
||||
"_fetch_head_ref",
|
||||
"_split_sizes_ref",
|
||||
"_local_consumed_ref",
|
||||
):
|
||||
state[key] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""Reconnect to LanceDB after unpickling in a worker process."""
|
||||
from . import connect as _connect
|
||||
|
||||
table_name = state.pop("_table_name")
|
||||
table_state = state.pop("_table")
|
||||
perm_name, perm_data = state.pop("_perm_table")
|
||||
self.__dict__.update(state)
|
||||
if self._connection_factory is not None:
|
||||
self._table = self._connection_factory(table_name)
|
||||
else:
|
||||
self._table = _table_from_pickle_state(table_state)
|
||||
self._perm_table = _connect("memory://").create_table(perm_name, perm_data)
|
||||
|
||||
def state_dict(self) -> dict:
|
||||
"""Snapshot the dataset's consumption state.
|
||||
|
||||
The returned dict is topology-independent: at global step boundaries
|
||||
every split has been consumed the same number of times (by the
|
||||
round-robin design), so the per-split count is a single uniform value
|
||||
that is identical across all ranks and DataLoader workers.
|
||||
"""
|
||||
return {
|
||||
"shuffle_seed": self._shuffle_seed,
|
||||
"num_splits": self._num_splits,
|
||||
"epoch": self._epoch,
|
||||
"samples_consumed_per_split": [self._resume_offset] * self._num_splits,
|
||||
}
|
||||
|
||||
def load_state_dict(self, state: dict) -> None:
|
||||
"""Resume from a previously snapshotted state.
|
||||
|
||||
Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ
|
||||
from the checkpoint, since a different split structure or shuffle order
|
||||
makes mid-epoch resumption meaningless.
|
||||
"""
|
||||
if state["num_splits"] != self._num_splits:
|
||||
raise ValueError(
|
||||
f"num_splits mismatch: checkpoint has {state['num_splits']}, "
|
||||
f"current dataset has {self._num_splits}"
|
||||
)
|
||||
if state["shuffle_seed"] != self._shuffle_seed:
|
||||
raise ValueError(
|
||||
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
|
||||
f"current dataset has {self._shuffle_seed}"
|
||||
)
|
||||
consumed = state["samples_consumed_per_split"]
|
||||
# All entries are equal at step boundaries; use the first.
|
||||
if isinstance(consumed, list):
|
||||
self._resume_offset = consumed[0] if consumed else 0
|
||||
else:
|
||||
self._resume_offset = int(consumed)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user