mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-26 07:58:31 +00:00
Compare commits
59 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 90b3715975 | |||
| 08e9ab54a9 | |||
| 4d4e3e6d36 | |||
| 3f1a94c8d1 | |||
| f376b01b09 | |||
| fce8dfb46a | |||
| da1acc46dc | |||
| 92f1d7ca67 | |||
| b24e99d37d | |||
| 44d88b9d65 | |||
| c29264fabd | |||
| 0448706e4c | |||
| ffee382a46 | |||
| bfa20d1594 | |||
| d3e612ca9c | |||
| f18ca49491 | |||
| 8582950b6f | |||
| ffd9af7c03 | |||
| 0f8f2a42e3 | |||
| c43dbe45f8 | |||
| 339d2bee3f | |||
| cef840463a | |||
| 78884d2755 | |||
| a80d3ab082 | |||
| 87679116a6 | |||
| 4463a23535 | |||
| df911fd657 | |||
| 24cb147383 | |||
| 3d07922170 | |||
| f38d890190 | |||
| 220724e7a1 | |||
| eeffbeffb5 | |||
| 6416840c33 | |||
| b030361d86 | |||
| 325cab394b | |||
| bc8674ab22 | |||
| 37032151d3 | |||
| 00c4a7b843 | |||
| 1773fb2239 | |||
| 8a4eaaa8b9 | |||
| 3fd322a93a | |||
| d8f0982ee8 | |||
| 7276c34c51 | |||
| 1918d1a3b6 | |||
| 3b626efa47 | |||
| 137eac9b50 | |||
| 06b53c97d6 | |||
| 711e05619b | |||
| afc0e5f497 | |||
| c12a6dce9f | |||
| 8ea78e3fbc | |||
| 40238d240a | |||
| 60428e1a32 | |||
| 5b982f2f05 | |||
| cde48fad95 | |||
| 1f2068b9fe | |||
| 7527890607 | |||
| a548e59d49 | |||
| 104fc5a08e |
@@ -1,137 +0,0 @@
|
||||
---
|
||||
name: lancedb-branch-ops
|
||||
description: Branch management for LanceDB tables via the REST API. Use this skill whenever someone wants to create, delete, list, or switch branches on a LanceDB table — or needs to make sure a write (metadata update, index build, etc.) lands on a specific branch instead of main. Invoke it even without the word "branch" if context makes clear they want an experimental copy of a table, want to isolate changes, or want to confirm a mutation didn't touch main. Covers: branches/list, branches/create, branches/delete, and passing "branch" in describe/update_field_metadata/create_index to target a non-main version.
|
||||
---
|
||||
|
||||
## 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 |
|
||||
@@ -1,178 +0,0 @@
|
||||
---
|
||||
name: lancedb-column-metadata
|
||||
description: Column metadata authoring for LanceDB tables via the REST API. This skill is required for tasks like writing field descriptions, setting tags on columns (field_type, model, project_id, version), classifying columns as embeddings vs labels vs eval metrics, or grouping versioned columns into logical families — because it has the API integration needed to read the schema and persist metadata back. Invoke whenever someone wants to document, annotate, tag, or classify what their table columns ARE. Trigger even without an explicit "LanceDB" mention, as long as the context is column-level documentation or tagging for an ML or vector database table.
|
||||
metadata:
|
||||
short-description: Write column descriptions, tags, and logical groupings to a LanceDB table
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
This skill authors column-level metadata for a LanceDB table. It connects to a LanceDB deployment over its REST API, inspects the table schema, generates appropriate metadata, and writes it back.
|
||||
|
||||
## Step 0: Establish the connection
|
||||
|
||||
Use the `lancedb-connect` skill (invoke it via the Skill tool) to resolve the base URL and auth headers (`x-api-key`, `x-lancedb-database`) for whichever deployment the user is working against — enterprise/self-hosted or a local dev server. Skip it only if the connection details are already established in the conversation.
|
||||
|
||||
All examples below use `{base_url}` — substitute the resolved endpoint and include the resolved headers on every request.
|
||||
|
||||
## Metadata keys
|
||||
|
||||
All metadata uses namespaced keys:
|
||||
|
||||
| Key | Purpose | Example value |
|
||||
|-----|---------|---------------|
|
||||
| `lancedb:description` | Human-readable explanation of what the column contains | `"CLIP ViT-L/14 image embedding, L2-normalized (768-dim)"` |
|
||||
| `lancedb:tag:<name>` | Flexible key-value tag; the suffix names the tag category | `lancedb:tag:field_type: "embedding"`, `lancedb:tag:model: "clip"`, `lancedb:tag:project_id: "foo"` |
|
||||
| `lancedb:logical-column` | Logical group/family this column belongs to | `"clip_features"` |
|
||||
|
||||
Tags are open-ended — use whatever key suffix and value make sense given the user's intent. The tag suffix should describe *what is being classified* (e.g., `field_type`, `model`, `project_id`) and the value describes *how*.
|
||||
|
||||
## Step 1: Resolve the table identifier
|
||||
|
||||
You need:
|
||||
- **Table name** (required) — e.g., `my_table` or `my_namespace.my_table`
|
||||
- **Database name** — ask if not provided and not inferable from context; it goes in the `x-lancedb-database` header, never in the URL path
|
||||
|
||||
The table identifier in the URL path is typically `table_name` for a top-level table, or `namespace$table_name` if the table lives in a namespace. The API accepts a `delimiter` query parameter to parse compound identifiers (default `$`).
|
||||
|
||||
## Step 2: Describe the table
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/describe
|
||||
Content-Type: application/json
|
||||
|
||||
{}
|
||||
```
|
||||
|
||||
The response contains `schema.fields` — an array of field objects:
|
||||
|
||||
```json
|
||||
{
|
||||
"schema": {
|
||||
"fields": [
|
||||
{
|
||||
"name": "clip_embedding_v3",
|
||||
"type": { "type": "FixedSizeList", "fields": [...], "listSize": 768 },
|
||||
"nullable": true,
|
||||
"metadata": { "lancedb:description": "..." }
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Each field has:
|
||||
- `name` — field name
|
||||
- `type` — Arrow data type (check `type.type` for the type string)
|
||||
- `nullable` — boolean
|
||||
- `metadata` — existing key-value metadata (read this before writing to avoid redundant updates)
|
||||
|
||||
For struct/nested fields, recurse into `type.fields` and represent them as dot-notation paths (e.g., `parent.child`).
|
||||
|
||||
If the user hasn't specified which columns to update, work with all columns.
|
||||
|
||||
## Step 3: Generate metadata
|
||||
|
||||
Decide what to generate based on the user's request.
|
||||
|
||||
### Writing descriptions (`lancedb:description`)
|
||||
|
||||
Base descriptions on:
|
||||
- The column name and Arrow type (e.g., `FixedSizeList` of floats → likely an embedding)
|
||||
- User-supplied context (upstream pipeline, sample values, domain knowledge)
|
||||
- Name patterns: `_embedding`/`_vec`/`_embed` → vector; `_label`/`_class` → label; `_score`/`_eval`/`_metric` → evaluation metric
|
||||
|
||||
Be specific and concise. Good: `"Sentence-BERT embedding of the query text (768-dim)."` Not: `"An embedding column."`
|
||||
|
||||
### Tagging columns (`lancedb:tag:<name>`)
|
||||
|
||||
Choose tag key names that match what the user asked to annotate. Common patterns:
|
||||
|
||||
- Semantic field type → `lancedb:tag:field_type: "embedding"` / `"text"` / `"image"` / `"label"` / `"eval"` / `"id"` / `"metadata"`
|
||||
- Model or source → `lancedb:tag:model: "clip"` / `"bert"` / `"vit"`
|
||||
- Project affiliation → `lancedb:tag:project_id: "<name>"`
|
||||
- Version → `lancedb:tag:version: "v3"` (and `lancedb:tag:latest: "true"` for the newest)
|
||||
|
||||
Use Arrow type as a hint: `FixedSizeList` + float → embedding; `Utf8`/`LargeUtf8` → text; `Binary` → image or blob.
|
||||
|
||||
Multiple tags on the same column are fine — each is a separate key.
|
||||
|
||||
### Grouping into logical columns (`lancedb:logical-column`)
|
||||
|
||||
Look for naming patterns across columns:
|
||||
- `clip_v1`, `clip_v2`, `clip_v3` → logical column `"clip"`, latest is `v3`
|
||||
- `text_embed_20240101`, `text_embed_20240601` → logical column `"text_embed"`, latest is the most recent date suffix
|
||||
|
||||
Write `lancedb:logical-column` on all members of a group. Mark the newest with `lancedb:tag:latest: "true"` (in addition to its version tag).
|
||||
|
||||
## Step 4: Write the metadata
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/update_field_metadata
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"updates": [
|
||||
{
|
||||
"path": "clip_v3",
|
||||
"metadata": {
|
||||
"lancedb:description": "CLIP ViT-L/14 image embedding, L2-normalized (1024-dim).",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v3",
|
||||
"lancedb:tag:latest": "true",
|
||||
"lancedb:logical-column": "clip"
|
||||
},
|
||||
"replace": false
|
||||
},
|
||||
{
|
||||
"path": "clip_v2",
|
||||
"metadata": {
|
||||
"lancedb:description": "CLIP ViT-B/32 image embedding (768-dim), superseded by v3.",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v2",
|
||||
"lancedb:logical-column": "clip"
|
||||
},
|
||||
"replace": false
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Rules:
|
||||
- **Use `"replace": false`** (merge) by default — this preserves existing metadata the user didn't ask to change
|
||||
- Use `"replace": true` only if the user explicitly asks to overwrite all existing metadata on a column
|
||||
- Set a value to `null` to delete a specific key
|
||||
- Batch all updates in a single request when possible
|
||||
|
||||
The response includes `version` (new table version) and `fields` (the updated metadata per field).
|
||||
|
||||
## Step 5: Confirm
|
||||
|
||||
Report back:
|
||||
- Which columns were updated and what was written
|
||||
- The new table version number
|
||||
- Any columns skipped (e.g., already had up-to-date metadata)
|
||||
|
||||
---
|
||||
|
||||
## Quick examples
|
||||
|
||||
**"Write descriptions for all columns in the `product_embeddings` table"**
|
||||
1. POST `/v1/table/product_embeddings/describe` → get all fields
|
||||
2. Generate a `lancedb:description` for each column based on name + type
|
||||
3. POST `update_field_metadata` with descriptions
|
||||
4. Report
|
||||
|
||||
**"Tag the columns in `model_outputs` with their field type and model"**
|
||||
1. Describe `model_outputs`
|
||||
2. For each field, classify by name + Arrow type → set `lancedb:tag:field_type` and `lancedb:tag:model` where applicable
|
||||
3. POST `update_field_metadata`
|
||||
4. Report
|
||||
|
||||
**"Group the feature columns in `training_features` into logical families and mark the latest version"**
|
||||
1. Describe the table
|
||||
2. Find version patterns → assign `lancedb:logical-column` and `lancedb:tag:version`; mark newest with `lancedb:tag:latest: "true"`
|
||||
3. POST `update_field_metadata`
|
||||
4. Show the grouping
|
||||
@@ -1,42 +0,0 @@
|
||||
---
|
||||
name: lancedb-connect
|
||||
description: Resolve how to connect to a LanceDB deployment over the REST API — figure out the base URL, API key, and database header. Use this before making any REST requests to a LanceDB table, whenever the endpoint or auth setup is not already known. Also useful on its own when someone asks how to connect, authenticate, or curl their LanceDB instance.
|
||||
metadata:
|
||||
short-description: Resolve the base URL and auth headers for a LanceDB deployment
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
Produce two things every REST request needs:
|
||||
|
||||
1. **Base URL** — the endpoint
|
||||
2. **Headers** — `x-api-key`, and usually `x-lancedb-database`
|
||||
|
||||
## Resolution steps
|
||||
|
||||
1. If the user already gave a URL and API key (or said which environment they're working against), use that.
|
||||
2. Otherwise, look for credentials already available in the environment:
|
||||
- Env vars like `LANCEDB_URI` / `LANCEDB_HOST` / `LANCEDB_API_KEY`
|
||||
- A LanceDB endpoint already running or port-forwarded locally (the REST default port is 2333, i.e. `http://localhost:2333`)
|
||||
3. If you didn't find both pieces, ask the user directly: **"What's your LanceDB endpoint's URL, and what's your API key?"** Also ask which database to use if it isn't obvious. Don't guess or probe further — the user knows their deployment.
|
||||
|
||||
## Validating the connection
|
||||
|
||||
Make a cheap authenticated request and check the status:
|
||||
|
||||
```bash
|
||||
curl -s -w "\n%{http_code}" "{base_url}/v1/table/?limit=1" \
|
||||
-H "x-api-key: <key>" \
|
||||
-H "x-lancedb-database: <database>"
|
||||
```
|
||||
|
||||
- `200` — connection, key, and database header all good
|
||||
- `401` — API key missing or wrong
|
||||
- `400` mentioning a database header — this deployment expects `x-lancedb-database`
|
||||
|
||||
## Non-REST equivalents
|
||||
|
||||
If the caller would rather use the SDK or CLI than raw REST, the same credentials work:
|
||||
|
||||
- Python SDK: `lancedb.connect("db://<database>", api_key="<key>", host_override="<base_url>")`
|
||||
- `lancedb` CLI: a `[profiles.<name>]` entry in `~/.lancedb/config.toml` with `http_server_url`, `api_key`, `database`
|
||||
@@ -0,0 +1,81 @@
|
||||
---
|
||||
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`
|
||||
- Column metadata authoring (both SDKs): `references/column_metadata.md`
|
||||
- Branch operations (both SDKs): `references/branch_ops.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. Read `column_metadata.md` when the task is documenting, tagging, classifying, or grouping table columns (field descriptions, `lancedb:tag:*` tags, logical column families). Read `branch_ops.md` when the task involves branch lifecycle (list/create/delete), writing to a non-main branch, or verifying a change stayed off main.
|
||||
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,117 @@
|
||||
# Branch Operations
|
||||
|
||||
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. Use for branch lifecycle tasks, experimental/isolated table versions, targeting an operation at a non-main branch, or confirming a mutation did not affect main.
|
||||
|
||||
Works on local/OSS and remote Enterprise/Cloud tables.
|
||||
|
||||
## The branch model (important)
|
||||
|
||||
Branches are isolated, writable lines of history forked from another branch (or a specific version). Writes on a branch never affect `main`.
|
||||
|
||||
There is **no global "switch branch" state** — you never repoint the whole table at a branch. Instead, **operations are scoped by which table handle you use**:
|
||||
|
||||
- The handle you got from `open_table(name)` / `openTable(name)` targets `main`.
|
||||
- `branches.create(...)` and `branches.checkout(...)` return a **new table handle scoped to that branch**. Every read/write on that handle (add, update, `update_field_metadata`, `create_index`, search, …) lands on the branch.
|
||||
- The original main handle is unaffected — keep it around to verify isolation.
|
||||
|
||||
`branches.list()` returns only non-main branches. Main always exists and is not listed.
|
||||
|
||||
## Python
|
||||
|
||||
`table.branches` is a property returning the branch manager; `table.current_branch()` tells you what a handle is scoped to (`None` = main).
|
||||
|
||||
```python
|
||||
table = db.open_table("products") # scoped to main
|
||||
|
||||
# list — dict of name -> metadata (parent_branch, parent_version, ...); {} = only main
|
||||
table.branches.list()
|
||||
|
||||
# create: forks from main by default and returns a handle scoped to the new branch
|
||||
exp = table.branches.create("experiment-reindex")
|
||||
exp = table.branches.create("exp2", from_ref="main", from_version=None) # optional fork point
|
||||
|
||||
# checkout an existing branch -> branch-scoped handle
|
||||
wip = table.branches.checkout("wip-branch")
|
||||
# with version= it pins to that version (read-only detached view); omit to track latest, writable
|
||||
|
||||
# operate on the branch simply by using its handle
|
||||
wip.update_field_metadata(
|
||||
{"path": "category", "metadata": {"lancedb:description": "Product category label."}}
|
||||
)
|
||||
wip.create_scalar_index("category")
|
||||
|
||||
# delete: removes only the branch pointer; main and row data remain intact
|
||||
table.branches.delete("stale-2024")
|
||||
|
||||
# alternatively, open a branch handle directly from the connection
|
||||
wip = db.open_table("products", branch="wip-branch")
|
||||
|
||||
exp.current_branch() # "experiment-reindex"
|
||||
table.current_branch() # None (main)
|
||||
```
|
||||
|
||||
Async: same shape — `table.branches` returns `AsyncBranches`; `await table.branches.create(...)` etc.
|
||||
|
||||
## TypeScript
|
||||
|
||||
`table.branches()` is an **async method** returning the `Branches` manager; `table.currentBranch()` returns the scoped branch or `null` for main.
|
||||
|
||||
```typescript
|
||||
const table = await db.openTable("products"); // scoped to main
|
||||
const branches = await table.branches();
|
||||
|
||||
// list — Record<string, BranchContents>; {} = only main
|
||||
await branches.list();
|
||||
|
||||
// create: forks from main by default, returns a Table scoped to the new branch
|
||||
const exp = await branches.create("experiment-reindex");
|
||||
const exp2 = await branches.create("exp2", "main" /* fromRef */, undefined /* fromVersion */);
|
||||
|
||||
// checkout an existing branch -> branch-scoped Table
|
||||
const wip = await branches.checkout("wip-branch");
|
||||
// with a version arg it pins (read-only detached view); omit to track latest, writable
|
||||
|
||||
// operate on the branch simply by using its handle
|
||||
await wip.updateFieldMetadata([
|
||||
{ path: "category", metadata: { "lancedb:description": "Product category label." } },
|
||||
]);
|
||||
await wip.createIndex("category");
|
||||
|
||||
// delete: removes only the branch pointer; main and row data remain intact
|
||||
await branches.delete("stale-2024");
|
||||
|
||||
// alternatively, open a branch handle directly from the connection
|
||||
const wip2 = await db.openTable("products", { branch: "wip-branch" });
|
||||
|
||||
exp.currentBranch(); // "experiment-reindex"
|
||||
table.currentBranch(); // null (main)
|
||||
```
|
||||
|
||||
## Verifying isolation
|
||||
|
||||
After writing to a branch, confirm the change did NOT land on main by reading through both handles:
|
||||
|
||||
```python
|
||||
wip = table.branches.checkout("wip-branch")
|
||||
wip.update_field_metadata({"path": "category", "metadata": {"lancedb:description": "..."}})
|
||||
|
||||
assert b"lancedb:description" in (wip.schema.field("category").metadata or {})
|
||||
assert b"lancedb:description" not in (table.schema.field("category").metadata or {}) # main untouched
|
||||
```
|
||||
|
||||
Two handles on the same branch see each other's writes (e.g. `table.branches.create("exp")` and `db.open_table(name, branch="exp")`); main stays isolated.
|
||||
|
||||
## Quick reference
|
||||
|
||||
| Goal | Python | TypeScript |
|
||||
|------|--------|------------|
|
||||
| List branches (non-main) | `table.branches.list()` | `await (await table.branches()).list()` |
|
||||
| Create branch (off main) | `table.branches.create(name)` → branch handle | `await branches.create(name)` → branch `Table` |
|
||||
| Create from a fork point | `table.branches.create(name, from_ref=..., from_version=...)` | `await branches.create(name, fromRef, fromVersion)` |
|
||||
| Get a branch handle | `table.branches.checkout(name)` or `db.open_table(t, branch=name)` | `await branches.checkout(name)` or `await db.openTable(t, { branch: name })` |
|
||||
| Pin to a branch version (read-only) | `table.branches.checkout(name, version=v)` | `await branches.checkout(name, v)` |
|
||||
| Delete branch | `table.branches.delete(name)` | `await branches.delete(name)` |
|
||||
| Which branch is this handle on? | `table.current_branch()` (`None` = main) | `table.currentBranch()` (`null` = main) |
|
||||
| Target main | use the original (non-branch) handle | use the original (non-branch) handle |
|
||||
|
||||
Branch names must be non-empty; empty names raise a validation error.
|
||||
@@ -0,0 +1,183 @@
|
||||
# Column Metadata Authoring
|
||||
|
||||
Write column-level descriptions, tags, and logical groupings onto a LanceDB table's schema. Use this when the user wants to document, annotate, tag, or classify what their table columns ARE (embeddings vs labels vs eval metrics, model provenance, version families, etc.).
|
||||
|
||||
Works on local/OSS and remote Enterprise/Cloud tables alike — read the schema through the table handle, write through `update_field_metadata` (Python) / `updateFieldMetadata` (TypeScript).
|
||||
|
||||
## Metadata key conventions
|
||||
|
||||
All metadata uses namespaced keys:
|
||||
|
||||
| Key | Purpose | Example value |
|
||||
|-----|---------|---------------|
|
||||
| `lancedb:description` | Human-readable explanation of what the column contains | `"CLIP ViT-L/14 image embedding, L2-normalized (768-dim)"` |
|
||||
| `lancedb:tag:<name>` | Flexible key-value tag; the suffix names the tag category | `lancedb:tag:field_type: "embedding"`, `lancedb:tag:model: "clip"`, `lancedb:tag:project_id: "foo"` |
|
||||
| `lancedb:logical-column` | Logical group/family this column belongs to | `"clip_features"` |
|
||||
|
||||
Tags are open-ended — use whatever key suffix and value make sense given the user's intent. The tag suffix should describe *what is being classified* (e.g., `field_type`, `model`, `project_id`) and the value describes *how*. Multiple tags on the same column are fine — each is a separate key. All values are strings.
|
||||
|
||||
## Step 1: Read the schema and existing metadata
|
||||
|
||||
Read existing metadata before writing, to avoid redundant updates.
|
||||
|
||||
Python — `table.schema` (sync property; async: `await table.schema()`) returns a `pyarrow.Schema`. **Arrow field metadata is bytes-keyed in Python**:
|
||||
|
||||
```python
|
||||
schema = table.schema
|
||||
for field in schema:
|
||||
meta = field.metadata or {} # dict[bytes, bytes], e.g. {b"lancedb:description": b"..."}
|
||||
print(field.name, field.type, field.nullable, meta)
|
||||
```
|
||||
|
||||
TypeScript — `await table.schema()` returns an Arrow `Schema`; field metadata is a `Map<string, string>`:
|
||||
|
||||
```typescript
|
||||
const schema = await table.schema();
|
||||
for (const field of schema.fields) {
|
||||
console.log(field.name, field.type, field.nullable, field.metadata); // Map
|
||||
// field.metadata.get("lancedb:description")
|
||||
}
|
||||
```
|
||||
|
||||
For struct/nested fields, recurse into the field's children and address them as dot-paths (e.g., `parent.child`).
|
||||
|
||||
If the user hasn't specified which columns to update, work with all columns.
|
||||
|
||||
## Step 2: Generate metadata
|
||||
|
||||
Decide what to generate based on the user's request.
|
||||
|
||||
### Descriptions (`lancedb:description`)
|
||||
|
||||
Base descriptions on:
|
||||
- The column name and Arrow type (e.g., `FixedSizeList` of floats → likely an embedding)
|
||||
- User-supplied context (upstream pipeline, sample values, domain knowledge)
|
||||
- Name patterns: `_embedding`/`_vec`/`_embed` → vector; `_label`/`_class` → label; `_score`/`_eval`/`_metric` → evaluation metric
|
||||
|
||||
Be specific and concise. Good: `"Sentence-BERT embedding of the query text (768-dim)."` Not: `"An embedding column."`
|
||||
|
||||
### Tags (`lancedb:tag:<name>`)
|
||||
|
||||
Choose tag key names that match what the user asked to annotate. Common patterns:
|
||||
|
||||
- Semantic field type → `lancedb:tag:field_type: "embedding"` / `"text"` / `"image"` / `"label"` / `"eval"` / `"id"` / `"metadata"`
|
||||
- Model or source → `lancedb:tag:model: "clip"` / `"bert"` / `"vit"`
|
||||
- Project affiliation → `lancedb:tag:project_id: "<name>"`
|
||||
- Version → `lancedb:tag:version: "v3"` (and `lancedb:tag:latest: "true"` for the newest)
|
||||
|
||||
Use Arrow type as a hint: `FixedSizeList` + float → embedding; `Utf8`/`LargeUtf8` → text; `Binary` → image or blob.
|
||||
|
||||
### Logical groupings (`lancedb:logical-column`)
|
||||
|
||||
Look for naming patterns across columns:
|
||||
- `clip_v1`, `clip_v2`, `clip_v3` → logical column `"clip"`, latest is `v3`
|
||||
- `text_embed_20240101`, `text_embed_20240601` → logical column `"text_embed"`, latest is the most recent date suffix
|
||||
|
||||
Write `lancedb:logical-column` on all members of a group. Mark the newest with `lancedb:tag:latest: "true"` (in addition to its version tag).
|
||||
|
||||
## Step 3: Write the metadata
|
||||
|
||||
Each update names a field by dot-path and carries a metadata map. Semantics (identical in both SDKs):
|
||||
|
||||
- **Merge by default** (`replace` omitted/false) — preserves existing metadata the user didn't ask to change
|
||||
- `replace: true` swaps the field's entire metadata map — only if the user explicitly asks to overwrite
|
||||
- A value of `None`/`null` deletes that specific key
|
||||
- Batch all field updates into a single call when possible
|
||||
- Returns the new table version
|
||||
|
||||
Python (sync and async take one dict per field, as varargs):
|
||||
|
||||
```python
|
||||
res = table.update_field_metadata(
|
||||
{
|
||||
"path": "clip_v3",
|
||||
"metadata": {
|
||||
"lancedb:description": "CLIP ViT-L/14 image embedding, L2-normalized (1024-dim).",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v3",
|
||||
"lancedb:tag:latest": "true",
|
||||
"lancedb:logical-column": "clip",
|
||||
},
|
||||
},
|
||||
{
|
||||
"path": "clip_v2",
|
||||
"metadata": {
|
||||
"lancedb:description": "CLIP ViT-B/32 image embedding (768-dim), superseded by v3.",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v2",
|
||||
"lancedb:logical-column": "clip",
|
||||
},
|
||||
},
|
||||
)
|
||||
print(res.version) # new table version
|
||||
|
||||
# merge semantics: add a key, delete one via None, keep the rest
|
||||
table.update_field_metadata(
|
||||
{"path": "clip_v2", "metadata": {"lancedb:tag:archived": "true", "lancedb:tag:latest": None}}
|
||||
)
|
||||
```
|
||||
|
||||
(`replace_field_metadata` is deprecated — use `update_field_metadata`.)
|
||||
|
||||
TypeScript (takes an array of `FieldMetadataUpdate`):
|
||||
|
||||
```typescript
|
||||
const res = await table.updateFieldMetadata([
|
||||
{
|
||||
path: "clip_v3",
|
||||
metadata: {
|
||||
"lancedb:description": "CLIP ViT-L/14 image embedding, L2-normalized (1024-dim).",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v3",
|
||||
"lancedb:tag:latest": "true",
|
||||
"lancedb:logical-column": "clip",
|
||||
},
|
||||
},
|
||||
{
|
||||
path: "clip_v2",
|
||||
metadata: {
|
||||
"lancedb:description": "CLIP ViT-B/32 image embedding (768-dim), superseded by v3.",
|
||||
"lancedb:tag:field_type": "embedding",
|
||||
"lancedb:tag:model": "clip",
|
||||
"lancedb:tag:version": "v2",
|
||||
"lancedb:logical-column": "clip",
|
||||
},
|
||||
},
|
||||
]);
|
||||
console.log(res.version); // new table version
|
||||
|
||||
// merge semantics: add a key, delete one via null, keep the rest
|
||||
await table.updateFieldMetadata([
|
||||
{ path: "clip_v2", metadata: { "lancedb:tag:archived": "true", "lancedb:tag:latest": null } },
|
||||
]);
|
||||
```
|
||||
|
||||
## Step 4: Confirm
|
||||
|
||||
Report back:
|
||||
- Which columns were updated and what was written
|
||||
- The new table version number (from the result)
|
||||
- Any columns skipped (e.g., already had up-to-date metadata)
|
||||
|
||||
## Quick examples
|
||||
|
||||
**"Write descriptions for all columns in the `product_embeddings` table"**
|
||||
1. Read `table.schema` → all fields + existing metadata
|
||||
2. Generate a `lancedb:description` for each column based on name + type
|
||||
3. One `update_field_metadata` call with all descriptions
|
||||
4. Report
|
||||
|
||||
**"Tag the columns in `model_outputs` with their field type and model"**
|
||||
1. Read the schema
|
||||
2. For each field, classify by name + Arrow type → set `lancedb:tag:field_type` and `lancedb:tag:model` where applicable
|
||||
3. Write in one batched call
|
||||
4. Report
|
||||
|
||||
**"Group the feature columns in `training_features` into logical families and mark the latest version"**
|
||||
1. Read the schema
|
||||
2. Find version patterns → assign `lancedb:logical-column` and `lancedb:tag:version`; mark newest with `lancedb:tag:latest: "true"`
|
||||
3. Write in one batched call
|
||||
4. Show the grouping
|
||||
@@ -0,0 +1,131 @@
|
||||
# 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.
|
||||
|
||||
## Column (Field) Metadata
|
||||
|
||||
```python
|
||||
schema = table.schema # sync property; async: await table.schema()
|
||||
meta = schema.field("category").metadata # dict[bytes, bytes] — Arrow metadata is bytes-keyed
|
||||
res = table.update_field_metadata( # varargs: one dict per field; works local + remote
|
||||
{"path": "category", "metadata": {"lancedb:description": "...", "lancedb:tag:field_type": "label"}}
|
||||
)
|
||||
res.version # new table version
|
||||
```
|
||||
|
||||
Merges by default; a `None` value deletes that key; `"replace": True` swaps the whole map. Nested fields use dot-paths (`"a.b.c"`). `replace_field_metadata` is deprecated. See `references/column_metadata.md` for key conventions (`lancedb:description`, `lancedb:tag:<name>`, `lancedb:logical-column`) and the authoring workflow.
|
||||
|
||||
## Branches
|
||||
|
||||
```python
|
||||
table.branches.list() # non-main branches; {} = only main
|
||||
exp = table.branches.create("exp") # fork off main -> handle scoped to the branch
|
||||
wip = table.branches.checkout("wip") # existing branch -> scoped handle (version= pins read-only)
|
||||
wip = db.open_table("t", branch="wip") # or open scoped directly
|
||||
table.branches.delete("stale") # removes only the branch pointer
|
||||
table.current_branch() # None = main
|
||||
```
|
||||
|
||||
There is no global switch — scoping is per table handle: any read/write on a branch handle lands on that branch; the original handle keeps targeting main. See `references/branch_ops.md` for the model and isolation checks.
|
||||
|
||||
## 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,105 @@
|
||||
# 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.
|
||||
|
||||
## Column (Field) Metadata
|
||||
|
||||
```typescript
|
||||
const schema = await table.schema();
|
||||
const meta = schema.fields.find((f) => f.name === "category")?.metadata; // Map<string, string>
|
||||
const res = await table.updateFieldMetadata([
|
||||
{ path: "category", metadata: { "lancedb:description": "...", "lancedb:tag:field_type": "label" } },
|
||||
]);
|
||||
res.version; // new table version
|
||||
```
|
||||
|
||||
Merges by default; a `null` value deletes that key; `replace: true` swaps the whole map. Nested fields use dot-paths (`"a.b.c"`). See `references/column_metadata.md` for key conventions (`lancedb:description`, `lancedb:tag:<name>`, `lancedb:logical-column`) and the authoring workflow.
|
||||
|
||||
## Branches
|
||||
|
||||
```typescript
|
||||
const branches = await table.branches(); // async manager
|
||||
await branches.list(); // non-main branches; {} = only main
|
||||
const exp = await branches.create("exp"); // fork off main -> Table scoped to the branch
|
||||
const wip = await branches.checkout("wip"); // existing branch -> scoped Table (version arg pins read-only)
|
||||
const wip2 = await db.openTable("t", { branch: "wip" }); // or open scoped directly
|
||||
await branches.delete("stale"); // removes only the branch pointer
|
||||
table.currentBranch(); // null = main
|
||||
```
|
||||
|
||||
There is no global switch — scoping is per table handle: any read/write on a branch handle lands on that branch; the original handle keeps targeting main. See `references/branch_ops.md` for the model and isolation checks.
|
||||
|
||||
## 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.32.0-beta.0"
|
||||
current_version = "0.32.0-beta.2"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
# CODEOWNERS
|
||||
#
|
||||
# These owners will be the default owners for everything in the repo.
|
||||
# They will be requested for review when someone opens a pull request.
|
||||
#
|
||||
# See https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
|
||||
|
||||
# Default owners for everything
|
||||
* @jackye1995 @wjones127
|
||||
|
||||
# Release and publish workflows — changes here can affect supply chain security
|
||||
/.github/workflows/ @jackye1995 @wjones127 @Xuanwo
|
||||
|
||||
# Remote client and auth — sensitive networking and auth code
|
||||
/rust/lancedb/src/remote/ @jackye1995 @wjones127
|
||||
|
||||
# Python FFI boundary
|
||||
/python/src/ @jackye1995 @wjones127 @AyushExel
|
||||
|
||||
# NodeJS FFI boundary
|
||||
/nodejs/src/ @jackye1995 @wjones127
|
||||
@@ -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
|
||||
|
||||
Generated
+250
-227
File diff suppressed because it is too large
Load Diff
+23
-23
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.1.0-beta.2", default-features = false, "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.1.0-beta.2", "tag" = "v9.1.0-beta.2", "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 }
|
||||
@@ -39,15 +39,15 @@ arrow-schema = "58.0.0"
|
||||
arrow-select = "58.0.0"
|
||||
arrow-cast = "58.0.0"
|
||||
async-trait = "0"
|
||||
datafusion = { version = "53.0.0", default-features = false }
|
||||
datafusion-catalog = "53.0.0"
|
||||
datafusion-common = { version = "53.0.0", default-features = false }
|
||||
datafusion-execution = "53.0.0"
|
||||
datafusion-expr = "53.0.0"
|
||||
datafusion-functions = "53.0.0"
|
||||
datafusion-physical-plan = "53.0.0"
|
||||
datafusion-physical-expr = "53.0.0"
|
||||
datafusion-sql = "53.0.0"
|
||||
datafusion = { version = "54.0.0", default-features = false }
|
||||
datafusion-catalog = "54.0.0"
|
||||
datafusion-common = { version = "54.0.0", default-features = false }
|
||||
datafusion-execution = "54.0.0"
|
||||
datafusion-expr = "54.0.0"
|
||||
datafusion-functions = "54.0.0"
|
||||
datafusion-physical-plan = "54.0.0"
|
||||
datafusion-physical-expr = "54.0.0"
|
||||
datafusion-sql = "54.0.0"
|
||||
env_logger = "0.11"
|
||||
half = { "version" = "2.7.1", default-features = false, features = [
|
||||
"num-traits",
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.32.0-beta.0</version>
|
||||
<version>0.32.0-beta.2</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -33,7 +33,7 @@ protected inner: Query | Promise<Query>;
|
||||
### analyzePlan()
|
||||
|
||||
```ts
|
||||
analyzePlan(): Promise<string>
|
||||
analyzePlan(distributedMetrics?): Promise<string>
|
||||
```
|
||||
|
||||
Executes the query and returns the physical query plan annotated with runtime metrics.
|
||||
@@ -41,6 +41,12 @@ Executes the query and returns the physical query plan annotated with runtime me
|
||||
This is useful for debugging and performance analysis, as it shows how the query was executed
|
||||
and includes metrics such as elapsed time, rows processed, and I/O statistics.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
Defaults to `"aggregate"`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
@@ -38,7 +38,7 @@ protected inner: NativeQueryType | Promise<NativeQueryType>;
|
||||
### analyzePlan()
|
||||
|
||||
```ts
|
||||
analyzePlan(): Promise<string>
|
||||
analyzePlan(distributedMetrics?): Promise<string>
|
||||
```
|
||||
|
||||
Executes the query and returns the physical query plan annotated with runtime metrics.
|
||||
@@ -46,6 +46,12 @@ Executes the query and returns the physical query plan annotated with runtime me
|
||||
This is useful for debugging and performance analysis, as it shows how the query was executed
|
||||
and includes metrics such as elapsed time, rows processed, and I/O statistics.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
Defaults to `"aggregate"`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
@@ -934,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
|
||||
|
||||
@@ -29,7 +29,7 @@ protected inner: TakeQuery | Promise<TakeQuery>;
|
||||
### analyzePlan()
|
||||
|
||||
```ts
|
||||
analyzePlan(): Promise<string>
|
||||
analyzePlan(distributedMetrics?): Promise<string>
|
||||
```
|
||||
|
||||
Executes the query and returns the physical query plan annotated with runtime metrics.
|
||||
@@ -37,6 +37,12 @@ Executes the query and returns the physical query plan annotated with runtime me
|
||||
This is useful for debugging and performance analysis, as it shows how the query was executed
|
||||
and includes metrics such as elapsed time, rows processed, and I/O statistics.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
Defaults to `"aggregate"`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
@@ -51,7 +51,7 @@ addQueryVector(vector): VectorQuery
|
||||
### analyzePlan()
|
||||
|
||||
```ts
|
||||
analyzePlan(): Promise<string>
|
||||
analyzePlan(distributedMetrics?): Promise<string>
|
||||
```
|
||||
|
||||
Executes the query and returns the physical query plan annotated with runtime metrics.
|
||||
@@ -59,6 +59,12 @@ Executes the query and returns the physical query plan annotated with runtime me
|
||||
This is useful for debugging and performance analysis, as it shows how the query was executed
|
||||
and includes metrics such as elapsed time, rows processed, and I/O statistics.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
Defaults to `"aggregate"`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
@@ -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)[]>
|
||||
@@ -72,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)
|
||||
@@ -107,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)
|
||||
@@ -116,6 +118,8 @@
|
||||
|
||||
## Type Aliases
|
||||
|
||||
- [AnalyzePlanDistributedMetrics](type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
- [BaseTokenizer](type-aliases/BaseTokenizer.md)
|
||||
- [Data](type-aliases/Data.md)
|
||||
- [DataLike](type-aliases/DataLike.md)
|
||||
- [FieldLike](type-aliases/FieldLike.md)
|
||||
@@ -125,6 +129,7 @@
|
||||
- [RecordBatchLike](type-aliases/RecordBatchLike.md)
|
||||
- [SchemaLike](type-aliases/SchemaLike.md)
|
||||
- [TableLike](type-aliases/TableLike.md)
|
||||
- [TokenizeTableOptions](type-aliases/TokenizeTableOptions.md)
|
||||
|
||||
## Functions
|
||||
|
||||
@@ -135,3 +140,4 @@
|
||||
- [makeArrowTable](functions/makeArrowTable.md)
|
||||
- [packBits](functions/packBits.md)
|
||||
- [permutationBuilder](functions/permutationBuilder.md)
|
||||
- [tokenize](functions/tokenize.md)
|
||||
|
||||
@@ -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,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,11 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / AnalyzePlanDistributedMetrics
|
||||
|
||||
# Type Alias: AnalyzePlanDistributedMetrics
|
||||
|
||||
```ts
|
||||
type AnalyzePlanDistributedMetrics: "aggregate" | "per_worker" | "full";
|
||||
```
|
||||
@@ -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.32.0-beta.0</version>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.0</version>
|
||||
<version>0.32.0-beta.2</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.19</lance-core.version>
|
||||
<lance-core.version>9.1.0-beta.2</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.32.0-beta.0"
|
||||
version = "0.32.0-beta.2"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -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 }];
|
||||
@@ -2705,8 +2775,13 @@ describe("when calling analyzePlan", () => {
|
||||
.fill(1)
|
||||
.map(() => Math.random());
|
||||
const plan = await table.query().nearestTo(queryVec).analyzePlan();
|
||||
console.log("Query Plan:\n", plan); // <--- Print the plan
|
||||
expect(plan).toMatch("AnalyzeExec");
|
||||
|
||||
const fullPlan = await table
|
||||
.query()
|
||||
.nearestTo(queryVec)
|
||||
.analyzePlan("full");
|
||||
expect(fullPlan).toMatch("AnalyzeExec");
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
@@ -13,9 +13,12 @@ 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";
|
||||
@@ -90,6 +93,7 @@ export {
|
||||
QueryBase,
|
||||
VectorQuery,
|
||||
TakeQuery,
|
||||
AnalyzePlanDistributedMetrics,
|
||||
QueryExecutionOptions,
|
||||
ColumnOrdering,
|
||||
FullTextSearchOptions,
|
||||
@@ -114,6 +118,7 @@ export {
|
||||
HnswPqOptions,
|
||||
HnswSqOptions,
|
||||
FtsOptions,
|
||||
BaseTokenizer,
|
||||
} from "./indices";
|
||||
|
||||
export {
|
||||
@@ -124,6 +129,8 @@ export {
|
||||
OptimizeOptions,
|
||||
Version,
|
||||
WriteProgress,
|
||||
FtsToken,
|
||||
TokenizeTableOptions,
|
||||
LsmWriteSpec,
|
||||
ColumnAlteration,
|
||||
FieldMetadataUpdate,
|
||||
@@ -155,6 +162,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
|
||||
|
||||
+12
-3
@@ -79,6 +79,8 @@ export interface QueryExecutionOptions {
|
||||
timeoutMs?: number;
|
||||
}
|
||||
|
||||
export type AnalyzePlanDistributedMetrics = "aggregate" | "per_worker" | "full";
|
||||
|
||||
export interface ColumnOrdering {
|
||||
columnName: string;
|
||||
ascending?: boolean;
|
||||
@@ -311,13 +313,20 @@ export class QueryBase<
|
||||
* KNNVectorDistance: metric=l2, metrics=[output_rows=1, elapsed_compute=114.333µs, output_batches=1]
|
||||
* LanceScan: uri=/path/to/data, projection=[vector], row_id=true, row_addr=false, ordered=false, metrics=[output_rows=1, elapsed_compute=103.626µs, bytes_read=549, iops=2, requests=2]
|
||||
*
|
||||
* @param distributedMetrics - How distributed worker metrics are displayed for remote query plans.
|
||||
* Defaults to `"aggregate"`.
|
||||
* @returns A query execution plan with runtime metrics for each step.
|
||||
*/
|
||||
async analyzePlan(): Promise<string> {
|
||||
async analyzePlan(
|
||||
distributedMetrics?: AnalyzePlanDistributedMetrics,
|
||||
): Promise<string> {
|
||||
const distributedMetricsMode = distributedMetrics ?? "aggregate";
|
||||
if (this.inner instanceof Promise) {
|
||||
return this.inner.then((inner) => inner.analyzePlan());
|
||||
return this.inner.then((inner) =>
|
||||
inner.analyzePlan(distributedMetricsMode),
|
||||
);
|
||||
} else {
|
||||
return this.inner.analyzePlan();
|
||||
return this.inner.analyzePlan(distributedMetricsMode);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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`.
|
||||
@@ -716,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>;
|
||||
|
||||
@@ -1173,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.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.32.0-beta.0",
|
||||
"version": "0.32.0-beta.2",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
|
||||
@@ -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)]
|
||||
|
||||
+56
-21
@@ -19,6 +19,7 @@ use lancedb::index::scalar::{
|
||||
BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
|
||||
Operator, PhraseQuery,
|
||||
};
|
||||
use lancedb::query::AnalyzePlanDistributedMetrics;
|
||||
use lancedb::query::ExecutableQuery;
|
||||
use lancedb::query::Query as LanceDbQuery;
|
||||
use lancedb::query::QueryBase;
|
||||
@@ -47,6 +48,28 @@ impl From<ColumnOrdering> for LanceDbColumnOrdering {
|
||||
}
|
||||
}
|
||||
|
||||
fn analyze_plan_options(
|
||||
distributed_metrics: Option<String>,
|
||||
) -> napi::Result<QueryExecutionOptions> {
|
||||
let analyze_plan_distributed_metrics =
|
||||
match distributed_metrics.as_deref().unwrap_or("aggregate") {
|
||||
"aggregate" => AnalyzePlanDistributedMetrics::Aggregate,
|
||||
"per_worker" => AnalyzePlanDistributedMetrics::PerWorker,
|
||||
"full" => AnalyzePlanDistributedMetrics::Full,
|
||||
mode => {
|
||||
return Err(napi::Error::from_reason(format!(
|
||||
"Invalid distributedMetrics value '{}'. Expected one of: \
|
||||
'aggregate', 'per_worker', 'full'",
|
||||
mode
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let mut options = QueryExecutionOptions::default();
|
||||
options.analyze_plan_distributed_metrics = analyze_plan_distributed_metrics;
|
||||
Ok(options)
|
||||
}
|
||||
|
||||
fn bytes_to_arrow_array(data: Uint8Array, dtype: String) -> napi::Result<Arc<dyn Array>> {
|
||||
let buf = arrow_buffer::Buffer::from(data.to_vec());
|
||||
let num_bytes = buf.len();
|
||||
@@ -200,13 +223,17 @@ impl Query {
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn analyze_plan(&self) -> napi::Result<String> {
|
||||
self.inner.analyze_plan().await.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
pub async fn analyze_plan(&self, distributed_metrics: Option<String>) -> napi::Result<String> {
|
||||
let options = analyze_plan_options(distributed_metrics)?;
|
||||
self.inner
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -412,13 +439,17 @@ impl VectorQuery {
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn analyze_plan(&self) -> napi::Result<String> {
|
||||
self.inner.analyze_plan().await.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
pub async fn analyze_plan(&self, distributed_metrics: Option<String>) -> napi::Result<String> {
|
||||
let options = analyze_plan_options(distributed_metrics)?;
|
||||
self.inner
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -491,13 +522,17 @@ impl TakeQuery {
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn analyze_plan(&self) -> napi::Result<String> {
|
||||
self.inner.analyze_plan().await.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
pub async fn analyze_plan(&self, distributed_metrics: Option<String>) -> napi::Result<String> {
|
||||
let options = analyze_plan_options(distributed_metrics)?;
|
||||
self.inner
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to execute analyze plan: {}",
|
||||
convert_error(&e)
|
||||
))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+42
-3
@@ -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};
|
||||
@@ -165,7 +165,7 @@ impl Table {
|
||||
if let Some(train) = train {
|
||||
builder = builder.train(train);
|
||||
}
|
||||
builder.execute().await.default_error()
|
||||
builder.execute().await.default_error().map(|_| ())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
@@ -574,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()?;
|
||||
@@ -681,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`.
|
||||
///
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.35.0-beta.1"
|
||||
current_version = "0.35.0-beta.2"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.35.0-beta.1"
|
||||
version = "0.35.0-beta.2"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
|
||||
@@ -61,10 +61,11 @@ tests = [
|
||||
"duckdb>=0.9.0",
|
||||
"pytz>=2023.3",
|
||||
"polars>=0.19, <=1.3.0",
|
||||
"pyarrow<25",
|
||||
"pyarrow-stubs>=16.0",
|
||||
"pylance>=5.0.0b5",
|
||||
"pylance==9.0.0rc1",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=52,<53",
|
||||
"datafusion>=54,<55",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
]
|
||||
dev = [
|
||||
|
||||
@@ -6,19 +6,33 @@ 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 .udf import (
|
||||
udf,
|
||||
table_udf,
|
||||
Udf,
|
||||
Job,
|
||||
JobFailedError,
|
||||
MaterializedView,
|
||||
AsyncJob,
|
||||
AsyncMaterializedView,
|
||||
)
|
||||
from .lineage import Lineage, Node, Edge, FunctionRef
|
||||
from .schema import blob, vector, BlobType
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
from .namespace import (
|
||||
connect_namespace,
|
||||
@@ -246,6 +260,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_"
|
||||
|
||||
|
||||
@@ -454,19 +502,35 @@ async def connect_async(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"udf",
|
||||
"table_udf",
|
||||
"Udf",
|
||||
"Job",
|
||||
"JobFailedError",
|
||||
"MaterializedView",
|
||||
"AsyncJob",
|
||||
"AsyncMaterializedView",
|
||||
"Lineage",
|
||||
"Node",
|
||||
"Edge",
|
||||
"FunctionRef",
|
||||
"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]
|
||||
@@ -25,10 +25,12 @@ 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
|
||||
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
|
||||
|
||||
class MetricPoint:
|
||||
name: str
|
||||
@@ -48,6 +50,20 @@ class MetricDescription:
|
||||
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)."""
|
||||
@@ -181,6 +197,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: ...
|
||||
@@ -227,6 +254,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: ...
|
||||
@@ -258,6 +292,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:
|
||||
@@ -353,7 +394,9 @@ class Query:
|
||||
self, max_batch_length: Optional[int], timeout: Optional[timedelta]
|
||||
) -> RecordBatchStream: ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
async def analyze_plan(self) -> str: ...
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
|
||||
) -> str: ...
|
||||
def to_query_request(self) -> PyQueryRequest: ...
|
||||
|
||||
class TakeQuery:
|
||||
@@ -361,6 +404,10 @@ class TakeQuery:
|
||||
def with_row_id(self): ...
|
||||
async def output_schema(self) -> pa.Schema: ...
|
||||
async def execute(self) -> RecordBatchStream: ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
|
||||
) -> str: ...
|
||||
def to_query_request(self) -> PyQueryRequest: ...
|
||||
|
||||
class FTSQuery:
|
||||
@@ -381,6 +428,10 @@ class FTSQuery:
|
||||
async def execute(
|
||||
self, max_batch_length: Optional[int], timeout: Optional[timedelta]
|
||||
) -> RecordBatchStream: ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
|
||||
) -> str: ...
|
||||
def to_query_request(self) -> PyQueryRequest: ...
|
||||
|
||||
class VectorQuery:
|
||||
@@ -403,6 +454,10 @@ class VectorQuery:
|
||||
def bypass_vector_index(self): ...
|
||||
def nearest_to_text(self, query: dict) -> HybridQuery: ...
|
||||
def order_by(self, ordering: Optional[List[ColumnOrdering]]): ...
|
||||
async def explain_plan(self, verbose: Optional[bool]) -> str: ...
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
|
||||
) -> str: ...
|
||||
def to_query_request(self) -> PyQueryRequest: ...
|
||||
|
||||
class HybridQuery:
|
||||
@@ -493,6 +548,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`."""
|
||||
|
||||
+486
-3
@@ -65,6 +65,7 @@ if TYPE_CHECKING:
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
from .udf import MaterializedView, AsyncMaterializedView
|
||||
|
||||
from .namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
@@ -562,6 +563,277 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("serialize is not supported for this connection type")
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION).
|
||||
|
||||
Pass a ``@udf`` / ``@table_udf``-decorated function (preferred):
|
||||
|
||||
db.create_function(embed)
|
||||
|
||||
or the explicit fields:
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name: str or Udf
|
||||
A decorated UDF object, or the function name.
|
||||
language: str
|
||||
Implementation language (currently "python").
|
||||
return_type: str
|
||||
SQL return type, e.g. "FLOAT", "FLOAT[1536]",
|
||||
"STRUCT(a FLOAT, b VARCHAR)", "TABLE(chunk VARCHAR, idx INT)".
|
||||
body: str
|
||||
Function body: source text, or base64 cloudpickle bytes when
|
||||
options["body_format"] == "cloudpickle".
|
||||
options: dict, optional
|
||||
input_columns, pip, num_gpus, batch_size, timeout,
|
||||
error_policy, docker_image, body_format, ...
|
||||
replace: bool
|
||||
Drop an existing function of the same name first.
|
||||
"""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
LOOP.run(self._conn.create_function(name, language, return_type, body, options))
|
||||
|
||||
def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
|
||||
def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
LOOP.run(self._conn.drop_function(name))
|
||||
|
||||
def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "MaterializedView":
|
||||
"""Create a materialized view (CREATE MATERIALIZED VIEW); returns a
|
||||
`MaterializedView` handle (``.wait()`` blocks until it is populated).
|
||||
|
||||
Two ways to specify the view body:
|
||||
|
||||
- ergonomic: pass ``source`` (a table name or table) and ``select``
|
||||
items -- column names, expression strings ("embed(body)"),
|
||||
(alias, expression) tuples, or ``@udf`` / ``@table_udf`` objects.
|
||||
The SELECT is assembled and parsed server-side (one parser, shared
|
||||
with SQL).
|
||||
- raw: pass ``query=`` with a full SELECT, e.g.
|
||||
"SELECT id, embed(body) AS vec FROM articles WHERE id > 1".
|
||||
|
||||
`partition_by` partitions the view's (single) table function on a source
|
||||
column. If that column has an IVF vector index the server partitions by
|
||||
its index clusters (image-dedup style); otherwise it groups by distinct
|
||||
value. (Geneva's `partition_by` and `partition_by_indexed_column` unify
|
||||
here -- the engine picks the strategy from the column.)
|
||||
"""
|
||||
from .udf import build_view_query, MaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
self._drop_view_if_exists(name)
|
||||
job_id = LOOP.run(
|
||||
self._conn.create_materialized_view(
|
||||
name,
|
||||
query=query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
)
|
||||
return MaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def _drop_view_if_exists(self, name: str) -> None:
|
||||
# `replace=True` is "drop if present"; only a not-found error is
|
||||
# benign here. Anything else (perms, server fault) must surface rather
|
||||
# than be masked by a later create failure.
|
||||
try:
|
||||
self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""A `Job` for reconnecting to an inflight job by id -- e.g. an
|
||||
id you stored, or one returned from the SQL / REST surface. Submit
|
||||
methods (`refresh_column`, `MaterializedView.refresh`) already return a
|
||||
handle directly, so you do not need this to wait on a fresh submission."""
|
||||
from .udf import Job
|
||||
|
||||
return Job(self, job_id)
|
||||
|
||||
def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. `direction` is "upstream" | "downstream" | "both" (server
|
||||
default both); `depth` limits column-hops (transitive when omitted)."""
|
||||
# `self._conn` is the AsyncConnection; drive its async `lineage`
|
||||
# (which parses the JSON) on the loop, mirroring create_materialized_view.
|
||||
return LOOP.run(
|
||||
self._conn.lineage(table, column, direction=direction, depth=depth)
|
||||
)
|
||||
|
||||
def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a materialized-view refresh, return the job id.
|
||||
The public surface is ``MaterializedView.refresh()`` (which returns a
|
||||
`Job`); this stays private so refresh is only reached through the
|
||||
handle.
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._conn._refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
)
|
||||
|
||||
def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH). Returns a
|
||||
plan with .has_work / .source_version / .last_refreshed_version /
|
||||
.full_refresh / .rebuild / .units_total. `full=True` plans a full
|
||||
rebuild (incremental planning needs stable row IDs on the source)."""
|
||||
return LOOP.run(
|
||||
self._conn.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
)
|
||||
|
||||
def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options (ALTER MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.alter_materialized_view(name, auto_refresh=auto_refresh))
|
||||
|
||||
def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition (DROP MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.drop_materialized_view(name))
|
||||
|
||||
def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
|
||||
def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
|
||||
def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
|
||||
Passing ``table`` (the job's table) lets the server answer with an O(1)
|
||||
single-node read instead of scanning the database's active jobs.
|
||||
Returns the job's status, or None if it's unknown or no longer active.
|
||||
"""
|
||||
return LOOP.run(self._conn.get_job(job_id, table))
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False if none was inflight (already finished or
|
||||
unknown id) -- cancellation is best-effort.
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
def describe_platform_job(self, platform_job_id: str):
|
||||
"""Describe a platform job (POST /v1/jobs/describe): registry-backed
|
||||
lifecycle state plus the owner-written status payload. None when the
|
||||
registry has no such job."""
|
||||
return LOOP.run(self._conn.describe_platform_job(platform_job_id))
|
||||
|
||||
def resolve_platform_job_id(
|
||||
self, manifest_job_id: str, table: "str | None" = None
|
||||
):
|
||||
"""Resolve a submission (manifest) job id to its platform job id.
|
||||
None until the job has registered (dispatch is async)."""
|
||||
return LOOP.run(self._conn.resolve_platform_job_id(manifest_job_id, table))
|
||||
|
||||
def cancel_platform_job(self, platform_job_id: str) -> None:
|
||||
"""Cancel a platform job (POST /v1/jobs/cancel). Idempotent on
|
||||
already-terminal jobs."""
|
||||
return LOOP.run(self._conn.cancel_platform_job(platform_job_id))
|
||||
|
||||
def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Pass ``job_id`` to narrow to a single job. Unlike :meth:`list_jobs`
|
||||
(live, inflight) these are the terminal records.
|
||||
"""
|
||||
return LOOP.run(self._conn.job_history(job_id))
|
||||
|
||||
def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS),
|
||||
optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return LOOP.run(self._conn.errors(job_id, table))
|
||||
|
||||
|
||||
class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
@@ -1655,7 +1927,7 @@ class AsyncConnection(object):
|
||||
namespace_client=namespace_client,
|
||||
)
|
||||
|
||||
return AsyncTable(new_table)
|
||||
return AsyncTable(new_table, conn=self)
|
||||
|
||||
async def open_table(
|
||||
self,
|
||||
@@ -1728,7 +2000,7 @@ class AsyncConnection(object):
|
||||
namespace_client=namespace_client,
|
||||
managed_versioning=managed_versioning,
|
||||
)
|
||||
tbl = AsyncTable(table)
|
||||
tbl = AsyncTable(table, conn=self)
|
||||
# "main" is the default branch, so treat it as no branch: remote rejects
|
||||
# every branch checkout (even "main"), and the version still applies.
|
||||
if branch is not None and branch != "main":
|
||||
@@ -1785,7 +2057,218 @@ class AsyncConnection(object):
|
||||
source_tag=source_tag,
|
||||
is_shallow=is_shallow,
|
||||
)
|
||||
return AsyncTable(table)
|
||||
return AsyncTable(table, conn=self)
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
async def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION). Accepts a ``@udf``/``@table_udf``
|
||||
object (preferred) or the explicit (name, language, return_type, body,
|
||||
options)."""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
await self._inner.create_function(name, language, return_type, body, options)
|
||||
|
||||
async def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return await self._inner.list_functions()
|
||||
|
||||
async def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
await self._inner.drop_function(name)
|
||||
|
||||
async def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "AsyncMaterializedView":
|
||||
"""Create a materialized view; returns an `AsyncMaterializedView`
|
||||
handle (``.wait()`` blocks until populated). Pass either ``query=`` (a
|
||||
full SELECT) or ``source`` + ``select`` items; `partition_by`
|
||||
partitions the view's table function on a source column (index-cluster
|
||||
if the column is IVF-indexed, else distinct-value). See the sync
|
||||
method for the select grammar."""
|
||||
from .udf import build_view_query, AsyncMaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
job_id = await self._inner.create_materialized_view(
|
||||
name,
|
||||
query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
return AsyncMaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""An `AsyncJob` for reconnecting to an inflight job by id (a
|
||||
stored id, or one from the SQL / REST surface). Submit methods already
|
||||
return a handle, so this is only needed to re-attach to an existing
|
||||
job."""
|
||||
from .udf import AsyncJob
|
||||
|
||||
return AsyncJob(self, job_id)
|
||||
|
||||
async def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view (or column). See the sync
|
||||
`Connection.lineage`. Returns a `Lineage`."""
|
||||
from .lineage import Lineage
|
||||
|
||||
raw = await self._inner.table_lineage(table, column, direction, depth)
|
||||
return Lineage.from_json(raw)
|
||||
|
||||
async def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a refresh, return the job id. The public surface is
|
||||
``AsyncMaterializedView.refresh()`` (returns an `AsyncJob`).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return await self._inner.refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
|
||||
async def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return await self._inner.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
|
||||
async def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options."""
|
||||
await self._inner.alter_materialized_view(name, auto_refresh)
|
||||
|
||||
async def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition."""
|
||||
await self._inner.drop_materialized_view(name)
|
||||
|
||||
async def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return await self._inner.list_materialized_views()
|
||||
|
||||
async def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
|
||||
async def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
``table`` (the job's table) enables an O(1) server-side lookup.
|
||||
Returns the job's status, or None if unknown / no longer active."""
|
||||
return await self._inner.get_job(job_id, table)
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False otherwise (best-effort).
|
||||
"""
|
||||
return await self._inner.cancel_job(job_id)
|
||||
|
||||
async def describe_platform_job(self, platform_job_id: str):
|
||||
"""Describe a platform job: registry-backed lifecycle state plus the
|
||||
owner-written status payload. None when the registry has no such
|
||||
job."""
|
||||
return await self._inner.describe_platform_job(platform_job_id)
|
||||
|
||||
async def resolve_platform_job_id(
|
||||
self, manifest_job_id: str, table: "str | None" = None
|
||||
):
|
||||
"""Resolve a submission (manifest) job id to its platform job id.
|
||||
None until the job has registered (dispatch is async)."""
|
||||
return await self._inner.resolve_platform_job_id(manifest_job_id, table)
|
||||
|
||||
async def cancel_platform_job(self, platform_job_id: str) -> None:
|
||||
"""Cancel a platform job. Idempotent on already-terminal jobs."""
|
||||
return await self._inner.cancel_platform_job(platform_job_id)
|
||||
|
||||
async def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Reads each table's durable job-history store. Pass ``job_id`` to narrow
|
||||
to a single job. Unlike :meth:`list_jobs` (live, inflight) these are the
|
||||
terminal records, with created/updated/completed timestamps.
|
||||
"""
|
||||
return await self._inner.job_history(job_id)
|
||||
|
||||
async def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS).
|
||||
|
||||
Optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return await self._inner.errors(job_id, table)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
|
||||
@@ -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__}",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""Client-side model of derived-compute lineage.
|
||||
|
||||
`Connection.lineage()` / `Table.lineage()` / `MaterializedView.lineage()` return
|
||||
a `Lineage`: the graph of what a column or materialized view derives from
|
||||
(upstream), what derives from it (downstream), and -- for each derived column --
|
||||
the function that produced it, the version it was produced with, and whether
|
||||
that is stale relative to the function the registry now holds.
|
||||
|
||||
The server returns this as JSON (the wire contract); these classes deserialize
|
||||
it. Nothing here talks to the server.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionRef:
|
||||
"""The function that produced a derived column, with version + location."""
|
||||
|
||||
name: str
|
||||
#: Version that produced the data (stamped at compute time), if known.
|
||||
as_computed_version: Optional[str] = None
|
||||
#: Version the registry currently holds for this function name.
|
||||
current_version: Optional[str] = None
|
||||
#: True when the column was produced by an older function than the registry
|
||||
#: now holds -- i.e. silently stale; re-refresh to catch up.
|
||||
stale_vs_current: bool = False
|
||||
language: Optional[str] = None
|
||||
docker_image: Optional[str] = None
|
||||
env_digest: Optional[str] = None
|
||||
code_uri: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "FunctionRef":
|
||||
return cls(
|
||||
name=d["name"],
|
||||
as_computed_version=d.get("as_computed_version"),
|
||||
current_version=d.get("current_version"),
|
||||
stale_vs_current=d.get("stale_vs_current", False),
|
||||
language=d.get("language"),
|
||||
docker_image=d.get("docker_image"),
|
||||
env_digest=d.get("env_digest"),
|
||||
code_uri=d.get("code_uri"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
"""A lineage node: a table, view, column, or function."""
|
||||
|
||||
kind: str # "table" | "view" | "column" | "function"
|
||||
id: str # "table", "table.column", or "fn:name@version"
|
||||
table: Optional[str] = None
|
||||
function: Optional[FunctionRef] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Node":
|
||||
fn = d.get("function")
|
||||
return cls(
|
||||
kind=d["kind"],
|
||||
id=d["id"],
|
||||
table=d.get("table"),
|
||||
function=FunctionRef._from(fn) if fn else None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Edge:
|
||||
"""`downstream` depends on `upstream`, produced by `via` (a function name,
|
||||
or None for a passthrough)."""
|
||||
|
||||
downstream: str
|
||||
upstream: str
|
||||
via: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Edge":
|
||||
return cls(downstream=d["downstream"], upstream=d["upstream"], via=d.get("via"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Lineage:
|
||||
"""A derived-compute lineage graph (nodes + labeled edges)."""
|
||||
|
||||
target: str
|
||||
nodes: List[Node] = field(default_factory=list)
|
||||
edges: List[Edge] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, raw: Union[str, bytes, dict]) -> "Lineage":
|
||||
d = json.loads(raw) if isinstance(raw, (str, bytes)) else raw
|
||||
return cls(
|
||||
target=d.get("target", ""),
|
||||
nodes=[Node._from(n) for n in d.get("nodes", [])],
|
||||
edges=[Edge._from(e) for e in d.get("edges", [])],
|
||||
)
|
||||
|
||||
def functions(self) -> List[FunctionRef]:
|
||||
"""The function nodes in the graph."""
|
||||
return [n.function for n in self.nodes if n.function is not None]
|
||||
|
||||
def stale(self) -> List[FunctionRef]:
|
||||
"""Functions whose as-computed version is behind the current registry
|
||||
version -- the columns they produced are silently out of date."""
|
||||
return [f for f in self.functions() if f.stale_vs_current]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
def prune(d: dict) -> dict:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
return {
|
||||
"target": self.target,
|
||||
"nodes": [
|
||||
prune(
|
||||
{
|
||||
"kind": n.kind,
|
||||
"id": n.id,
|
||||
"table": n.table,
|
||||
"function": prune(vars(n.function)) if n.function else None,
|
||||
}
|
||||
)
|
||||
for n in self.nodes
|
||||
],
|
||||
"edges": [prune(vars(e)) for e in self.edges],
|
||||
}
|
||||
|
||||
def to_graphviz(self) -> str:
|
||||
"""Graphviz DOT for the lineage DAG: columns/tables as nodes, function
|
||||
names on edges, drift edges dashed + red."""
|
||||
stale_names = {f.name for f in self.stale()}
|
||||
out = [
|
||||
"digraph lineage {",
|
||||
" rankdir=LR;",
|
||||
' node [fontname="monospace"];',
|
||||
]
|
||||
for n in self.nodes:
|
||||
if n.kind == "function":
|
||||
continue
|
||||
shape = "ellipse" if n.kind in ("table", "view") else "box"
|
||||
out.append(f' "{n.id}" [shape={shape}];')
|
||||
for e in self.edges:
|
||||
attrs = ""
|
||||
if e.via:
|
||||
if e.via in stale_names:
|
||||
attrs = f' [label="{e.via}" color=red style=dashed]'
|
||||
else:
|
||||
attrs = f' [label="{e.via}"]'
|
||||
out.append(f' "{e.upstream}" -> "{e.downstream}"{attrs};')
|
||||
out.append("}")
|
||||
return "\n".join(out)
|
||||
|
||||
def _repr_html_(self) -> str:
|
||||
warn = ""
|
||||
drift = self.stale()
|
||||
if drift:
|
||||
names = ", ".join(sorted({f.name for f in drift}))
|
||||
warn = (
|
||||
f'<p style="color:#b00000"><b>stale vs current:</b> {names} '
|
||||
"(re-refresh to catch up)</p>"
|
||||
)
|
||||
rows = "".join(
|
||||
f"<tr><td><code>{e.downstream}</code></td>"
|
||||
f"<td>← {e.via or ''}</td>"
|
||||
f"<td><code>{e.upstream}</code></td></tr>"
|
||||
for e in self.edges
|
||||
)
|
||||
return (
|
||||
f"<b>lineage: <code>{self.target}</code></b>{warn}"
|
||||
"<table><tr><th>derived</th><th>via</th><th>from</th></tr>"
|
||||
f"{rows}</table>"
|
||||
)
|
||||
+382
-102
@@ -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
|
||||
@@ -71,27 +79,25 @@ if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
T = TypeVar("T", bound="LanceModel")
|
||||
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
|
||||
|
||||
|
||||
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:
|
||||
@@ -140,13 +146,7 @@ def _combine_where(
|
||||
return f"({existing_sql}) AND ({new_sql})"
|
||||
|
||||
|
||||
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 _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
if columns is None:
|
||||
return {}
|
||||
if isinstance(columns, list):
|
||||
@@ -171,7 +171,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 = {
|
||||
@@ -179,10 +183,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}
|
||||
@@ -215,11 +219,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"):
|
||||
@@ -229,7 +233,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()
|
||||
@@ -239,7 +243,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)
|
||||
@@ -260,13 +266,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,
|
||||
@@ -674,7 +738,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
|
||||
#
|
||||
@@ -958,7 +1022,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()
|
||||
@@ -1017,6 +1081,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
|
||||
@@ -1195,6 +1264,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.
|
||||
|
||||
@@ -1268,7 +1373,9 @@ class LanceQueryBuilder(ABC):
|
||||
self._order_by = ordering
|
||||
return self
|
||||
|
||||
def analyze_plan(self) -> str:
|
||||
def analyze_plan(
|
||||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||||
) -> str:
|
||||
"""
|
||||
Run the query and return its execution plan with runtime metrics.
|
||||
|
||||
@@ -1306,12 +1413,22 @@ class LanceQueryBuilder(ABC):
|
||||
fragments_scanned=..., ranges_scanned=1, rows_scanned=1,
|
||||
bytes_read=..., iops=..., requests=..., task_wait_time=...]
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||||
Defaults to "aggregate".
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
"aggregate" preserves the legacy summary, "per_worker" shows each
|
||||
worker separately, and "full" includes both.
|
||||
|
||||
Returns
|
||||
-------
|
||||
plan : str
|
||||
The physical query execution plan with runtime metrics.
|
||||
"""
|
||||
return self._table._analyze_plan(self.to_query_object())
|
||||
return self._table._analyze_plan(
|
||||
self.to_query_object(), distributed_metrics=distributed_metrics
|
||||
)
|
||||
|
||||
def vector(self, vector: Union[np.ndarray, list]) -> Self:
|
||||
"""Set the vector to search for.
|
||||
@@ -1371,13 +1488,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:
|
||||
@@ -1625,7 +1758,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:
|
||||
"""
|
||||
@@ -1685,7 +1820,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
|
||||
)
|
||||
@@ -1829,8 +1964,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
|
||||
-------
|
||||
@@ -1840,6 +1974,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
|
||||
@@ -1864,7 +2013,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,
|
||||
@@ -1882,22 +2031,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
|
||||
@@ -1925,7 +2065,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(
|
||||
@@ -1947,7 +2089,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:
|
||||
@@ -2019,14 +2161,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
|
||||
-------
|
||||
@@ -2051,15 +2192,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(
|
||||
@@ -2443,9 +2594,17 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
indented_fts = "\n".join(" " + line for line in fts_plan.splitlines())
|
||||
return f"{reranker_label}\n {indented_vector}\n {indented_fts}"
|
||||
|
||||
def analyze_plan(self):
|
||||
def analyze_plan(
|
||||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||||
) -> str:
|
||||
"""Execute the query and display with runtime metrics.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||||
Defaults to "aggregate".
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
|
||||
Returns
|
||||
-------
|
||||
plan : str
|
||||
@@ -2453,9 +2612,19 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._create_query_builders()
|
||||
|
||||
results = ["Vector Search Plan:"]
|
||||
results.append(self._table._analyze_plan(self._vector_query.to_query_object()))
|
||||
results.append(
|
||||
self._table._analyze_plan(
|
||||
self._vector_query.to_query_object(),
|
||||
distributed_metrics=distributed_metrics,
|
||||
)
|
||||
)
|
||||
results.append("FTS Search Plan:")
|
||||
results.append(self._table._analyze_plan(self._fts_query.to_query_object()))
|
||||
results.append(
|
||||
self._table._analyze_plan(
|
||||
self._fts_query.to_query_object(),
|
||||
distributed_metrics=distributed_metrics,
|
||||
)
|
||||
)
|
||||
return "\n".join(results)
|
||||
|
||||
def _create_query_builders(self):
|
||||
@@ -2500,7 +2669,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
|
||||
)
|
||||
@@ -2530,6 +2699,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:
|
||||
"""
|
||||
@@ -2539,11 +2711,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.
|
||||
@@ -2596,6 +2803,7 @@ class AsyncQueryBase(object):
|
||||
"""
|
||||
Include the _rowid column in the results.
|
||||
"""
|
||||
self._with_row_id = True
|
||||
self._inner.with_row_id()
|
||||
return self
|
||||
|
||||
@@ -2642,6 +2850,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
|
||||
@@ -2672,8 +2881,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]:
|
||||
@@ -2740,7 +2949,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):
|
||||
@@ -2781,14 +2990,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,
|
||||
@@ -2880,14 +3111,22 @@ class AsyncQueryBase(object):
|
||||
""" # noqa: E501
|
||||
return await self._inner.explain_plan(verbose)
|
||||
|
||||
async def analyze_plan(self):
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||||
) -> str:
|
||||
"""Execute the query and display with runtime metrics.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||||
Defaults to "aggregate".
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
|
||||
Returns
|
||||
-------
|
||||
plan : str
|
||||
"""
|
||||
return await self._inner.analyze_plan()
|
||||
return await self._inner.analyze_plan(distributed_metrics)
|
||||
|
||||
|
||||
class AsyncStandardQuery(AsyncQueryBase):
|
||||
@@ -3573,9 +3812,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()
|
||||
|
||||
@@ -3591,8 +3845,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)
|
||||
|
||||
@@ -3615,18 +3875,18 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||||
RRFReranker(K=60)
|
||||
ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance]
|
||||
Take: columns="vector, _rowid, _distance, (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
FilterExec: _distance@2 IS NOT NULL
|
||||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||||
KNNVectorDistance: metric=l2
|
||||
LanceRead: uri=..., projection=[vector], ...
|
||||
Take: columns="vector, _rowid, _distance, (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
FilterExec: _distance@2 IS NOT NULL
|
||||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||||
KNNVectorDistance: metric=l2
|
||||
LanceRead: uri=..., projection=[vector], ...
|
||||
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score]
|
||||
Take: columns="_rowid, _score, (vector), (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
MatchQuery: column=text, query=hello
|
||||
Take: columns="_rowid, _score, (vector), (text)"
|
||||
CoalesceBatchesExec: target_batch_size=1024
|
||||
GlobalLimitExec: skip=0, fetch=10
|
||||
MatchQuery: column=text, query=[hello]
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -3645,7 +3905,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
indented_fts = "\n".join(" " + line for line in fts_plan.splitlines())
|
||||
return f"{self._reranker}\n {indented_vector}\n {indented_fts}"
|
||||
|
||||
async def analyze_plan(self):
|
||||
async def analyze_plan(
|
||||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||||
) -> str:
|
||||
"""
|
||||
Execute the query and return the physical execution plan with runtime metrics.
|
||||
|
||||
@@ -3654,14 +3916,24 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
elapsed time, I/O stats, and more. It’s useful for debugging and
|
||||
performance analysis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||||
Defaults to "aggregate".
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
|
||||
Returns
|
||||
-------
|
||||
plan : str
|
||||
"""
|
||||
results = ["Vector Search Query:"]
|
||||
results.append(await self._inner.to_vector_query().analyze_plan())
|
||||
results.append(
|
||||
await self._inner.to_vector_query().analyze_plan(distributed_metrics)
|
||||
)
|
||||
results.append("FTS Search Query:")
|
||||
results.append(await self._inner.to_fts_query().analyze_plan())
|
||||
results.append(
|
||||
await self._inner.to_fts_query().analyze_plan(distributed_metrics)
|
||||
)
|
||||
|
||||
return "\n".join(results)
|
||||
|
||||
@@ -3945,14 +4217,22 @@ class BaseQueryBuilder(object):
|
||||
""" # noqa: E501
|
||||
return LOOP.run(self._inner.explain_plan(verbose))
|
||||
|
||||
def analyze_plan(self):
|
||||
def analyze_plan(
|
||||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||||
) -> str:
|
||||
"""Execute the query and display with runtime metrics.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||||
Defaults to "aggregate".
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
|
||||
Returns
|
||||
-------
|
||||
plan : str
|
||||
"""
|
||||
return LOOP.run(self._inner.analyze_plan())
|
||||
return LOOP.run(self._inner.analyze_plan(distributed_metrics))
|
||||
|
||||
|
||||
class LanceTakeQueryBuilder(BaseQueryBuilder):
|
||||
|
||||
@@ -13,10 +13,14 @@ from typing import (
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
TYPE_CHECKING,
|
||||
Union,
|
||||
Literal,
|
||||
overload,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..udf import Job
|
||||
import warnings
|
||||
|
||||
from lancedb import __version__
|
||||
@@ -28,6 +32,7 @@ from lancedb._lancedb import (
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
IndexConfig,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
@@ -55,7 +60,12 @@ from lancedb.merge import LanceMergeInsertBuilder
|
||||
from lancedb.embeddings import EmbeddingFunctionRegistry
|
||||
from lancedb.table import _normalize_progress
|
||||
|
||||
from ..query import LanceVectorQueryBuilder, LanceQueryBuilder, LanceTakeQueryBuilder
|
||||
from ..query import (
|
||||
AnalyzePlanDistributedMetrics,
|
||||
LanceQueryBuilder,
|
||||
LanceTakeQueryBuilder,
|
||||
LanceVectorQueryBuilder,
|
||||
)
|
||||
from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
|
||||
from ..types import BaseTokenizerType
|
||||
|
||||
@@ -244,6 +254,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))
|
||||
@@ -700,8 +727,15 @@ class RemoteTable(Table):
|
||||
def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str:
|
||||
return LOOP.run(self._table._explain_plan(query, verbose))
|
||||
|
||||
def _analyze_plan(self, query: Query) -> str:
|
||||
return LOOP.run(self._table._analyze_plan(query))
|
||||
def _analyze_plan(
|
||||
self,
|
||||
query: Query,
|
||||
*,
|
||||
distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
|
||||
) -> str:
|
||||
return LOOP.run(
|
||||
self._table._analyze_plan(query, distributed_metrics=distributed_metrics)
|
||||
)
|
||||
|
||||
def _output_schema(self, query: Query) -> pa.Schema:
|
||||
return LOOP.run(self._table._output_schema(query))
|
||||
@@ -884,8 +918,142 @@ class RemoteTable(Table):
|
||||
def count_rows(self, filter: Optional[str] = None) -> int:
|
||||
return LOOP.run(self._table.count_rows(filter))
|
||||
|
||||
def add_columns(self, transforms: Dict[str, str]) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
def add_columns(
|
||||
self,
|
||||
transforms: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
computed: Optional[Dict[str, tuple]] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
LOOP.run(self._table.add_columns(computed=computed))
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "Job":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `Job` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("c").wait()``). Server-backed feature
|
||||
(LanceDB Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from ..udf import Job
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return Job(self._job_conn(), job_id)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Derived-compute lineage of this table, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. See `Connection.lineage`."""
|
||||
return self._job_conn().lineage(
|
||||
self._name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
def _job_conn(self):
|
||||
"""A client connection for polling jobs this table spawns. Built lazily
|
||||
from the table's serialized connection state and cached (not pickled --
|
||||
a forked/unpickled table rebuilds it on next use)."""
|
||||
from lancedb import deserialize_conn
|
||||
|
||||
conn = getattr(self, "_job_conn_cache", None)
|
||||
if conn is None:
|
||||
conn = deserialize_conn(self._serialized_connection_state())
|
||||
self._job_conn_cache = conn
|
||||
return conn
|
||||
|
||||
def load_columns(
|
||||
self,
|
||||
source: Union[str, Iterable[str]],
|
||||
pk: str,
|
||||
columns: Union[Iterable[str], Dict[str, str]],
|
||||
*,
|
||||
source_format: str = "parquet",
|
||||
source_pk: Optional[str] = None,
|
||||
on_missing: str = "carry",
|
||||
source_storage_options: Optional[Dict[str, str]] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
commit_granularity: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Fill existing columns from an external source by primary-key join.
|
||||
|
||||
The distributed-job equivalent of Geneva's ``Table.load_columns()``:
|
||||
imports precomputed values (e.g. embeddings) from Parquet/Lance/IPC into
|
||||
this table, matching on a primary key. Returns the load job id.
|
||||
Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source: str | list[str]
|
||||
One source URI or a list of URIs.
|
||||
pk: str
|
||||
Destination primary-key column. Also the source key unless
|
||||
``source_pk`` is given.
|
||||
columns: list[str] | dict[str, str]
|
||||
Value columns to load. A list loads same-named columns; a dict maps
|
||||
``{target: source}``.
|
||||
source_format: str
|
||||
``"parquet"`` (default), ``"lance"``, or ``"ipc"``.
|
||||
source_pk: str, optional
|
||||
Source primary-key column when it differs from ``pk``.
|
||||
on_missing: str
|
||||
Behavior for destination rows with no source match:
|
||||
``"carry"`` (default, keep existing), ``"null"``, or ``"error"``.
|
||||
"""
|
||||
if isinstance(source, str):
|
||||
source = [source]
|
||||
if isinstance(columns, dict):
|
||||
mappings = [(target, src) for target, src in columns.items()]
|
||||
else:
|
||||
mappings = [(c, None) for c in columns]
|
||||
return LOOP.run(
|
||||
self._table.load_columns(
|
||||
list(source),
|
||||
source_format,
|
||||
pk,
|
||||
mappings,
|
||||
source_key=source_pk,
|
||||
source_storage_options=source_storage_options,
|
||||
on_missing=on_missing,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
commit_granularity=commit_granularity,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
@@ -994,6 +1162,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",
|
||||
]
|
||||
|
||||
@@ -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.
|
||||
|
||||
+521
-67
@@ -29,6 +29,14 @@ from urllib.parse import urlparse
|
||||
from lancedb.scannable import _register_optional_converters, to_scannable
|
||||
|
||||
from . import __version__
|
||||
from ._blob import (
|
||||
BlobFile,
|
||||
_normalize_blob_row_ids,
|
||||
_wrap_blob_files,
|
||||
strip_auto_row_ids,
|
||||
validate_blob_mode,
|
||||
)
|
||||
from .types import BlobMode
|
||||
from lancedb.arrow import peek_reader
|
||||
from lancedb.background_loop import LOOP, embedding_executor
|
||||
from .dependencies import (
|
||||
@@ -65,6 +73,7 @@ from .expr import Expr
|
||||
from .merge import LanceMergeInsertBuilder
|
||||
from .pydantic import LanceModel, model_to_dict
|
||||
from .query import (
|
||||
AnalyzePlanDistributedMetrics,
|
||||
AsyncFTSQuery,
|
||||
AsyncHybridQuery,
|
||||
AsyncQuery,
|
||||
@@ -88,10 +97,7 @@ from .util import (
|
||||
value_to_sql,
|
||||
)
|
||||
from .index import lang_mapping
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
_VALID_BLOB_MODES = ("lazy", "bytes", "descriptions")
|
||||
from .schema import blob_v2_column_paths, schema_has_blob_field
|
||||
|
||||
|
||||
def _should_push_down_query_table(
|
||||
@@ -100,23 +106,6 @@ def _should_push_down_query_table(
|
||||
return namespace_client is not None and "QueryTable" in pushdown_operations
|
||||
|
||||
|
||||
def _validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in _VALID_BLOB_MODES:
|
||||
modes = ", ".join(repr(mode) for mode in _VALID_BLOB_MODES)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
|
||||
|
||||
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 _schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return any(_field_is_blob(field) for field in schema)
|
||||
|
||||
|
||||
_MODEL_BACKED_TOKENIZER_PREFIXES = ("jieba", "lindera")
|
||||
_MODEL_BACKED_TOKENIZER_ERRORS = (
|
||||
"unknown base tokenizer",
|
||||
@@ -173,6 +162,7 @@ def _maybe_add_fts_error_note(
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db import LanceDBConnection
|
||||
from .udf import Job
|
||||
from ._lancedb import (
|
||||
Table as LanceDBTable,
|
||||
OptimizeStats,
|
||||
@@ -185,6 +175,7 @@ if TYPE_CHECKING:
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
UpdateResult,
|
||||
@@ -712,6 +703,24 @@ def _normalize_progress(progress):
|
||||
return progress, False
|
||||
|
||||
|
||||
def _computed_groups(computed):
|
||||
"""Group computed columns by expression, preserving declaration order
|
||||
(struct-returning functions need their columns adjacent so schema order
|
||||
matches field order). Accepts the ergonomic forms -- `fn("col")` values
|
||||
and tuple keys for struct fan-out -- via `_normalize_computed`."""
|
||||
from .udf import _normalize_computed
|
||||
|
||||
groups = []
|
||||
for name, (sql_type, expression) in _normalize_computed(computed).items():
|
||||
for expr, cols in groups:
|
||||
if expr == expression:
|
||||
cols.append((name, sql_type))
|
||||
break
|
||||
else:
|
||||
groups.append((expression, [(name, sql_type)]))
|
||||
return groups
|
||||
|
||||
|
||||
class Table(ABC):
|
||||
"""
|
||||
A Table is a collection of Records in a LanceDB Database.
|
||||
@@ -817,6 +826,59 @@ class Table(ABC):
|
||||
"""The number of rows in this Table"""
|
||||
return self.count_rows(None)
|
||||
|
||||
def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF -- no compute happens
|
||||
here (the agent fills them lazily, or refresh_column() triggers a run).
|
||||
|
||||
.. deprecated::
|
||||
A computed column is an expression over a registered function, so
|
||||
bind it as one: ``add_columns(computed={"vec": embed("data")})``.
|
||||
``embed("data")`` applies the function to the `data` column and
|
||||
infers the type from the function's return signature -- the
|
||||
function never couples to a particular column. Prefer that form.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"vec": embed("data")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
else:
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
self.add_columns(computed=computed)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def embedding_functions(self) -> Dict[str, EmbeddingFunctionConfig]:
|
||||
@@ -893,7 +955,7 @@ class Table(ABC):
|
||||
wait_timeout: Optional[timedelta] = ...,
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
) -> None: ...
|
||||
) -> "Job": ...
|
||||
|
||||
# Legacy API overload (deprecated)
|
||||
@overload
|
||||
@@ -917,7 +979,7 @@ class Table(ABC):
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
target_partition_size: Optional[int] = ...,
|
||||
) -> None: ...
|
||||
) -> "Job": ...
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
@@ -968,6 +1030,14 @@ class Table(ABC):
|
||||
train : bool, default True
|
||||
Whether to train the index with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Job
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``job.wait()`` blocks until it completes; when the
|
||||
build finished within this call, the job is already ``finished``.
|
||||
Prefer ``job.wait()`` over the deprecated ``wait_timeout``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
New API (recommended):
|
||||
@@ -1159,6 +1229,8 @@ class Table(ABC):
|
||||
- "whitespace": Split text by whitespace, but not punctuation.
|
||||
- "raw": No tokenization. The entire text is treated as a single token.
|
||||
- "ngram": N-Gram tokenizer.
|
||||
- "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"
|
||||
@@ -1523,6 +1595,31 @@ class Table(ABC):
|
||||
A query object that can be executed to get the rows.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def blob_columns(self) -> list[str]:
|
||||
"""Names of the blob v2 columns declared on this table."""
|
||||
|
||||
@abstractmethod
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> pa.LargeBinaryArray:
|
||||
"""Materialize full blob bytes for ``column`` at the given rows.
|
||||
|
||||
Convenience for small payloads. For large values use
|
||||
:meth:`fetch_blob_files`.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
"""Open lazy, seekable :class:`~lancedb._blob.BlobFile` handles.
|
||||
|
||||
Prefer this over :meth:`fetch_blobs` for large payloads. ``row_ids`` is
|
||||
a ``list[int]`` or query ``pyarrow.Table`` with ``_rowid`` (or stashed
|
||||
row-id metadata). Null rows are ``None``. Local tables only.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def _execute_query(
|
||||
self,
|
||||
@@ -1536,7 +1633,12 @@ class Table(ABC):
|
||||
def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def _analyze_plan(self, query: Query) -> str: ...
|
||||
def _analyze_plan(
|
||||
self,
|
||||
query: Query,
|
||||
*,
|
||||
distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
|
||||
) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def _output_schema(self, query: Query) -> pa.Schema: ...
|
||||
@@ -1786,6 +1888,24 @@ class Table(ABC):
|
||||
[Table.create_index][lancedb.table.Table.create_index]
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
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.
|
||||
|
||||
Specify exactly one of ``column`` or ``index_name``.
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
||||
"""
|
||||
@@ -2085,8 +2205,8 @@ class LanceTable(Table):
|
||||
def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import LanceDBConnection
|
||||
|
||||
async_tbl = AsyncTable(tbl)
|
||||
conn = LanceDBConnection.from_inner(tbl.database())
|
||||
async_tbl = AsyncTable(tbl, conn=conn._conn)
|
||||
return cls(
|
||||
conn,
|
||||
async_tbl.name,
|
||||
@@ -2204,6 +2324,19 @@ class LanceTable(Table):
|
||||
def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder:
|
||||
return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
|
||||
|
||||
def blob_columns(self) -> list[str]:
|
||||
return LOOP.run(self._table.blob_columns())
|
||||
|
||||
def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> pa.LargeBinaryArray:
|
||||
return LOOP.run(self._table.fetch_blobs(column, row_ids))
|
||||
|
||||
def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
return LOOP.run(self._table.fetch_blob_files(column, row_ids))
|
||||
|
||||
@property
|
||||
def tags(self) -> Tags:
|
||||
"""Tag management for the table.
|
||||
@@ -2399,9 +2532,14 @@ class LanceTable(Table):
|
||||
-------
|
||||
pd.DataFrame
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
if blob_mode == "descriptions" or not _schema_has_blob_field(self.schema):
|
||||
return self.to_arrow().to_pandas(**kwargs)
|
||||
validate_blob_mode(blob_mode)
|
||||
if blob_mode == "descriptions" or not schema_has_blob_field(self.schema):
|
||||
arrow_tbl = self.to_arrow()
|
||||
if blob_mode == "descriptions":
|
||||
arrow_tbl = strip_auto_row_ids(
|
||||
arrow_tbl, blob_v2_column_paths(self.schema)
|
||||
)
|
||||
return arrow_tbl.to_pandas(**kwargs)
|
||||
|
||||
if (
|
||||
blob_mode == "lazy"
|
||||
@@ -2410,6 +2548,9 @@ class LanceTable(Table):
|
||||
):
|
||||
return self.to_arrow().to_pandas(**kwargs)
|
||||
|
||||
if blob_mode == "bytes" and blob_v2_column_paths(self.schema):
|
||||
return self.search().to_pandas(blob_mode=blob_mode, **kwargs)
|
||||
|
||||
return self.to_lance().to_pandas(blob_mode=blob_mode, **kwargs)
|
||||
|
||||
def to_arrow(self) -> pa.Table:
|
||||
@@ -2467,7 +2608,7 @@ class LanceTable(Table):
|
||||
wait_timeout: Optional[timedelta] = ...,
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
) -> None: ...
|
||||
) -> "Job": ...
|
||||
|
||||
# Legacy API overload (deprecated)
|
||||
@overload
|
||||
@@ -2493,7 +2634,7 @@ class LanceTable(Table):
|
||||
name: Optional[str] = ...,
|
||||
train: bool = ...,
|
||||
target_partition_size: Optional[int] = ...,
|
||||
) -> None: ...
|
||||
) -> "Job": ...
|
||||
|
||||
def create_index(
|
||||
self,
|
||||
@@ -2552,6 +2693,14 @@ class LanceTable(Table):
|
||||
train : bool, default True
|
||||
Whether to train the index with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Job
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``job.wait()`` blocks until it completes; when the
|
||||
build finished within this call, the job is already ``finished``.
|
||||
Prefer ``job.wait()`` over the deprecated ``wait_timeout``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
New API (recommended):
|
||||
@@ -2627,7 +2776,7 @@ class LanceTable(Table):
|
||||
target_partition_size=target_partition_size,
|
||||
)
|
||||
self.checkout_latest()
|
||||
return
|
||||
return self._sync_job(None)
|
||||
else:
|
||||
# New API: metric is the column name
|
||||
column = metric
|
||||
@@ -2664,19 +2813,30 @@ class LanceTable(Table):
|
||||
),
|
||||
)
|
||||
self.checkout_latest()
|
||||
return
|
||||
return self._sync_job(None)
|
||||
|
||||
return LOOP.run(
|
||||
self._table.create_index(
|
||||
column,
|
||||
replace=replace,
|
||||
config=config,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
train=train,
|
||||
return self._sync_job(
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
column,
|
||||
replace=replace,
|
||||
config=config,
|
||||
wait_timeout=wait_timeout,
|
||||
name=name,
|
||||
train=train,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def _sync_job(self, ajob) -> "Job":
|
||||
"""Convert an AsyncJob (or None for work done in-process) into a sync
|
||||
Job bound to this table's connection."""
|
||||
from .udf import Job
|
||||
|
||||
if ajob is not None and ajob.id:
|
||||
return Job(self._conn, ajob.id, table=self.name)
|
||||
return Job._completed(self._conn, table=self.name)
|
||||
|
||||
def _is_legacy_create_index_call(
|
||||
self,
|
||||
first_arg: str,
|
||||
@@ -2927,8 +3087,12 @@ class LanceTable(Table):
|
||||
config = LabelList()
|
||||
else:
|
||||
raise ValueError(f"Unknown index type {index_type}")
|
||||
return LOOP.run(
|
||||
self._table.create_index(column, replace=replace, config=config, name=name)
|
||||
return self._sync_job(
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
column, replace=replace, config=config, name=name
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@deprecation.deprecated(
|
||||
@@ -3011,7 +3175,7 @@ class LanceTable(Table):
|
||||
)
|
||||
|
||||
try:
|
||||
LOOP.run(
|
||||
ajob = LOOP.run(
|
||||
self._table.create_index(
|
||||
field_names,
|
||||
replace=replace,
|
||||
@@ -3026,6 +3190,7 @@ class LanceTable(Table):
|
||||
language=config.language,
|
||||
)
|
||||
raise e
|
||||
return self._sync_job(ajob)
|
||||
|
||||
@staticmethod
|
||||
def infer_tokenizer_configs(tokenizer_name: str) -> dict:
|
||||
@@ -3575,8 +3740,15 @@ class LanceTable(Table):
|
||||
def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str:
|
||||
return LOOP.run(self._table._explain_plan(query, verbose))
|
||||
|
||||
def _analyze_plan(self, query: Query) -> str:
|
||||
return LOOP.run(self._table._analyze_plan(query))
|
||||
def _analyze_plan(
|
||||
self,
|
||||
query: Query,
|
||||
*,
|
||||
distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
|
||||
) -> str:
|
||||
return LOOP.run(
|
||||
self._table._analyze_plan(query, distributed_metrics=distributed_metrics)
|
||||
)
|
||||
|
||||
def _output_schema(self, query: Query) -> pa.Schema:
|
||||
return LOOP.run(self._table._output_schema(query))
|
||||
@@ -3710,6 +3882,26 @@ class LanceTable(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.
|
||||
|
||||
Specify exactly one of ``column`` or ``index_name``.
|
||||
|
||||
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.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.tokenize(query, column=column, index_name=index_name)
|
||||
)
|
||||
|
||||
def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
||||
"""
|
||||
Retrieve statistics about an index
|
||||
@@ -3727,9 +3919,68 @@ class LanceTable(Table):
|
||||
return LOOP.run(self._table.index_stats(index_name))
|
||||
|
||||
def add_columns(
|
||||
self, transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
self,
|
||||
transforms: Dict[str, str]
|
||||
| pa.field
|
||||
| List[pa.field]
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
result_unused = LOOP.run(self._table.add_columns(computed=computed))
|
||||
del result_unused
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "Job":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `Job` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("col").wait()``) -- mirrors
|
||||
`MaterializedView.refresh()`. Server-backed feature (LanceDB
|
||||
Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from .udf import Job
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return Job(self._conn, job_id, table=self.name)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
@@ -4079,17 +4330,58 @@ def _handle_bad_vector_column(
|
||||
raise ValueError(
|
||||
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
||||
)
|
||||
vec_arr = pc.if_else(
|
||||
is_bad,
|
||||
pa.scalar([fill_value] * dim, type=vec_arr.type),
|
||||
vec_arr,
|
||||
)
|
||||
vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
|
||||
else:
|
||||
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
||||
|
||||
return data.set_column(position, vector_column_name, vec_arr)
|
||||
|
||||
|
||||
def _fill_bad_vector_values(
|
||||
arr: Union[pa.Array, pa.ChunkedArray],
|
||||
dim: int,
|
||||
fill_value: float,
|
||||
) -> pa.Array:
|
||||
if not isinstance(arr, pa.ChunkedArray):
|
||||
arr = pa.chunked_array([arr])
|
||||
arr = arr.combine_chunks()
|
||||
|
||||
# A fixed-size slice truncates long vectors and pads short vectors with nulls.
|
||||
# Slice an array marking the original child nulls in parallel so padding nulls
|
||||
# can be distinguished from null values that were already present.
|
||||
sliced = pc.list_slice(arr, 0, dim, return_fixed_size_list=True)
|
||||
child_nulls = pc.is_null(arr.values)
|
||||
parent_nulls = pc.is_null(arr)
|
||||
if pa.types.is_list(arr.type):
|
||||
original_child_nulls = pa.ListArray.from_arrays(
|
||||
arr.offsets, child_nulls, mask=parent_nulls
|
||||
)
|
||||
elif pa.types.is_large_list(arr.type):
|
||||
original_child_nulls = pa.LargeListArray.from_arrays(
|
||||
arr.offsets, child_nulls, mask=parent_nulls
|
||||
)
|
||||
else:
|
||||
original_child_nulls = pa.FixedSizeListArray.from_arrays(
|
||||
child_nulls, arr.type.list_size, mask=parent_nulls
|
||||
)
|
||||
sliced_child_nulls = pc.list_slice(
|
||||
original_child_nulls, 0, dim, return_fixed_size_list=True
|
||||
)
|
||||
needs_fill = pc.is_null(sliced_child_nulls.values)
|
||||
|
||||
values = sliced.values
|
||||
if pa.types.is_floating(values.type):
|
||||
values_for_nan_check = (
|
||||
values.cast(pa.float32()) if pa.types.is_float16(values.type) else values
|
||||
)
|
||||
needs_fill = pc.or_kleene(needs_fill, pc.is_nan(values_for_nan_check))
|
||||
|
||||
fill_scalar = pa.scalar(fill_value).cast(values.type)
|
||||
filled_values = pc.if_else(needs_fill, fill_scalar, values)
|
||||
filled = pa.FixedSizeListArray.from_arrays(filled_values, dim)
|
||||
return filled.cast(arr.type)
|
||||
|
||||
|
||||
def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
|
||||
if isinstance(arr, pa.ChunkedArray):
|
||||
values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
|
||||
@@ -4304,6 +4596,7 @@ class AsyncTable:
|
||||
self,
|
||||
table: LanceDBTable,
|
||||
*,
|
||||
conn: Optional[Any] = None,
|
||||
namespace_path: Optional[List[str]] = None,
|
||||
namespace_client: Optional[Any] = None,
|
||||
pushdown_operations: Optional[set] = None,
|
||||
@@ -4317,6 +4610,9 @@ class AsyncTable:
|
||||
[AsyncConnection.open_table][lancedb.AsyncConnection.open_table] to obtain
|
||||
Table objects."""
|
||||
self._inner = table
|
||||
#: The owning AsyncConnection, when known -- lets index/refresh calls
|
||||
#: hand back AsyncJob handles that can reach the platform jobs API.
|
||||
self._conn = conn
|
||||
self._namespace_path = namespace_path or []
|
||||
self._namespace_client = namespace_client
|
||||
self._pushdown_operations = pushdown_operations or set()
|
||||
@@ -4539,14 +4835,18 @@ class AsyncTable:
|
||||
-------
|
||||
pd.DataFrame
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
if blob_mode == "descriptions" or not _schema_has_blob_field(
|
||||
await self.schema()
|
||||
):
|
||||
return (await self.to_arrow()).to_pandas(**kwargs)
|
||||
validate_blob_mode(blob_mode)
|
||||
schema = await self.schema()
|
||||
if blob_mode == "descriptions" or not schema_has_blob_field(schema):
|
||||
arrow_tbl = await self.to_arrow()
|
||||
if blob_mode == "descriptions":
|
||||
arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema))
|
||||
return arrow_tbl.to_pandas(**kwargs)
|
||||
|
||||
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
|
||||
return (await self.to_arrow()).to_pandas(**kwargs)
|
||||
if blob_mode == "bytes" and blob_v2_column_paths(schema):
|
||||
return await self.query().to_pandas(blob_mode=blob_mode, **kwargs)
|
||||
return (await self._to_lance()).to_pandas(blob_mode=blob_mode, **kwargs)
|
||||
|
||||
async def to_arrow(self) -> pa.Table:
|
||||
@@ -4620,6 +4920,14 @@ class AsyncTable:
|
||||
train: bool, default True
|
||||
Whether to train the index with existing data. Vector indices always train
|
||||
with existing data.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AsyncJob
|
||||
A handle on the index build. When the server defers the build to a
|
||||
background job, ``await job.wait()`` blocks until it completes;
|
||||
when the build finished within this call, the job is already
|
||||
``finished``. Prefer ``await job.wait()`` over ``wait_timeout``.
|
||||
"""
|
||||
if config is not None:
|
||||
if not isinstance(
|
||||
@@ -4645,7 +4953,7 @@ class AsyncTable:
|
||||
+ str(type(config))
|
||||
)
|
||||
try:
|
||||
await self._inner.create_index(
|
||||
job_id = await self._inner.create_index(
|
||||
column,
|
||||
index=config,
|
||||
replace=replace,
|
||||
@@ -4662,6 +4970,12 @@ class AsyncTable:
|
||||
)
|
||||
raise e
|
||||
|
||||
from .udf import AsyncJob
|
||||
|
||||
if job_id:
|
||||
return AsyncJob(self._conn, job_id, table=self.name)
|
||||
return AsyncJob._completed(self._conn, table=self.name)
|
||||
|
||||
async def drop_index(self, name: str) -> None:
|
||||
"""
|
||||
Drop an index from the table.
|
||||
@@ -5270,10 +5584,15 @@ class AsyncTable:
|
||||
async_query = self._sync_query_to_async(query)
|
||||
return await async_query.explain_plan(verbose)
|
||||
|
||||
async def _analyze_plan(self, query: Query) -> str:
|
||||
async def _analyze_plan(
|
||||
self,
|
||||
query: Query,
|
||||
*,
|
||||
distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
|
||||
) -> str:
|
||||
# This method is used by the sync table
|
||||
async_query = self._sync_query_to_async(query)
|
||||
return await async_query.analyze_plan()
|
||||
return await async_query.analyze_plan(distributed_metrics)
|
||||
|
||||
async def _output_schema(self, query: Query) -> pa.Schema:
|
||||
async_query = self._sync_query_to_async(query)
|
||||
@@ -5432,9 +5751,44 @@ class AsyncTable:
|
||||
|
||||
return await self._inner.update(updates_sql, where)
|
||||
|
||||
async def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
Returns the refresh job id. Server-backed feature.
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh); they override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill)."""
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
return await self._inner.refresh_column(
|
||||
list(columns),
|
||||
where_clause=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
|
||||
async def add_columns(
|
||||
self, transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
self,
|
||||
transforms: dict[str, str]
|
||||
| pa.field
|
||||
| List[pa.field]
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
"""
|
||||
Add new columns with defined values.
|
||||
|
||||
@@ -5453,6 +5807,7 @@ class AsyncTable:
|
||||
version: the new version number of the table after adding columns.
|
||||
|
||||
"""
|
||||
result = None
|
||||
if isinstance(transforms, pa.Field):
|
||||
transforms = [transforms]
|
||||
if isinstance(transforms, list) and all(
|
||||
@@ -5460,9 +5815,69 @@ class AsyncTable:
|
||||
):
|
||||
transforms = pa.schema(transforms)
|
||||
if isinstance(transforms, pa.Schema):
|
||||
return await self._inner.add_columns_with_schema(transforms)
|
||||
result = await self._inner.add_columns_with_schema(transforms)
|
||||
elif transforms is not None:
|
||||
result = await self._inner.add_columns(list(transforms.items()))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
for expression, cols in _computed_groups(computed):
|
||||
await self._inner.add_computed_columns(cols, expression)
|
||||
return result
|
||||
|
||||
async def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF (async).
|
||||
|
||||
.. deprecated::
|
||||
Use ``add_columns(computed={"col": fn("input_col")})`` -- a computed
|
||||
column is an expression over a registered function, so bind it that
|
||||
way instead of coupling the UDF to the column here.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"col": fn("input_col")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
else:
|
||||
return await self._inner.add_columns(list(transforms.items()))
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
await self.add_columns(computed=computed)
|
||||
|
||||
async def alter_columns(
|
||||
self, *alterations: Iterable[dict[str, Any]]
|
||||
@@ -5647,6 +6062,24 @@ class AsyncTable:
|
||||
"""
|
||||
return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self)
|
||||
|
||||
async def blob_columns(self) -> list[str]:
|
||||
return await self._inner.blob_columns()
|
||||
|
||||
async def fetch_blobs(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> pa.LargeBinaryArray:
|
||||
return await self._inner.fetch_blobs(
|
||||
column, _normalize_blob_row_ids(row_ids, column)
|
||||
)
|
||||
|
||||
async def fetch_blob_files(
|
||||
self, column: str, row_ids: Union[list[int], pa.Table]
|
||||
) -> "list[Optional[BlobFile]]":
|
||||
handles = await self._inner.fetch_blob_files(
|
||||
column, _normalize_blob_row_ids(row_ids, column)
|
||||
)
|
||||
return _wrap_blob_files(handles)
|
||||
|
||||
@property
|
||||
def tags(self) -> AsyncTags:
|
||||
"""Tag management for the dataset.
|
||||
@@ -5748,6 +6181,24 @@ class AsyncTable:
|
||||
"""
|
||||
return await self._inner.list_indices()
|
||||
|
||||
async 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.
|
||||
|
||||
Specify exactly one of ``column`` or ``index_name``.
|
||||
|
||||
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.
|
||||
"""
|
||||
return await self._inner.tokenize(query, column=column, index_name=index_name)
|
||||
|
||||
async def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
||||
"""
|
||||
Retrieve statistics about an index
|
||||
@@ -6175,7 +6626,7 @@ class AsyncBranches:
|
||||
if from_ref == "main":
|
||||
from_ref = None
|
||||
inner = await self._table.branches.create(name, from_ref, from_version)
|
||||
return AsyncTable(inner)
|
||||
return AsyncTable(inner, conn=self._table._conn)
|
||||
|
||||
async def checkout(self, name: str, version: Optional[int] = None) -> "AsyncTable":
|
||||
"""Check out an existing branch and return a handle scoped to it.
|
||||
@@ -6189,7 +6640,10 @@ class AsyncBranches:
|
||||
handle is a read-only view of that version; when omitted it tracks
|
||||
the branch's latest and stays writable.
|
||||
"""
|
||||
return AsyncTable(await self._table.branches.checkout(name, version))
|
||||
return AsyncTable(
|
||||
await self._table.branches.checkout(name, version),
|
||||
conn=self._table._conn,
|
||||
)
|
||||
|
||||
async def delete(self, name: str) -> None:
|
||||
"""Delete a branch."""
|
||||
|
||||
@@ -1,11 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from typing import Literal
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
from .expr import Expr
|
||||
|
||||
# Query type literals
|
||||
QueryType = Literal["vector", "fts", "hybrid", "auto"]
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
QueryProjectionSpec = Union[
|
||||
List[str],
|
||||
List[Tuple[str, Union[str, Expr]]],
|
||||
Dict[str, Union[str, Expr]],
|
||||
]
|
||||
QueryProjection = Optional[QueryProjectionSpec]
|
||||
|
||||
# Distance type literals
|
||||
DistanceType = Literal["l2", "cosine", "dot"]
|
||||
DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"]
|
||||
@@ -42,5 +55,7 @@ IndexType = Literal[
|
||||
]
|
||||
|
||||
# Tokenizer literals
|
||||
BuiltinTokenizerType = Literal["simple", "raw", "whitespace", "ngram"]
|
||||
BuiltinTokenizerType = Literal[
|
||||
"simple", "raw", "whitespace", "ngram", "icu", "icu/split"
|
||||
]
|
||||
BaseTokenizerType = BuiltinTokenizerType | str
|
||||
|
||||
@@ -0,0 +1,847 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""UDF authoring for LanceDB derived compute (server-backed).
|
||||
|
||||
`@udf` / `@table_udf` turn a plain Python function into a registrable
|
||||
server-side UDF: a cloudpickled (or source) body, a SQL signature inferred
|
||||
from type hints, and the runtime options (pip deps, GPUs, batching, ...).
|
||||
Register and use them through the existing connection/table API:
|
||||
|
||||
import lancedb
|
||||
from lancedb import udf, table_udf
|
||||
|
||||
db = lancedb.connect("db://my_db", api_key="...", host_override="...")
|
||||
|
||||
@udf(pip=["torch>=2.0"], num_gpus=1)
|
||||
def embed(text: str) -> list[float]:
|
||||
return model.encode(text).tolist()
|
||||
|
||||
db.create_function(embed) # CREATE FUNCTION (once)
|
||||
tbl = db.open_table("docs")
|
||||
tbl.add_columns(computed={"vec": embed("text")}) # bind embed(text) -> vec
|
||||
tbl.refresh_column("vec").wait() # materialize (returns a Job)
|
||||
view = db.create_materialized_view("chunks", tbl, ["id", chunk_fn])
|
||||
|
||||
`embed("text")` applies the registered function to the `text` column and yields
|
||||
the expression `embed(text)`; the function itself stays decoupled from any
|
||||
column, so the same `embed` works on any column or table.
|
||||
|
||||
These operations are server-backed (LanceDB Enterprise / Cloud); the
|
||||
decorator itself works locally (define + call), only registration needs a
|
||||
remote connection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import dataclasses
|
||||
import functools
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
import json
|
||||
import time
|
||||
import typing
|
||||
|
||||
# -- type hints -> SQL type strings -------------------------------------
|
||||
|
||||
_SCALARS = {
|
||||
int: "BIGINT",
|
||||
# Pragmatic default for ML workloads: python float maps to FLOAT
|
||||
# (Float32). Use an explicit `returns=` for DOUBLE.
|
||||
float: "FLOAT",
|
||||
str: "VARCHAR",
|
||||
bool: "BOOLEAN",
|
||||
bytes: "BLOB",
|
||||
}
|
||||
|
||||
|
||||
class TypeInferenceError(TypeError):
|
||||
pass
|
||||
|
||||
|
||||
def sql_type(hint) -> str:
|
||||
"""SQL type string for a python type hint."""
|
||||
if hint in _SCALARS:
|
||||
return _SCALARS[hint]
|
||||
origin = typing.get_origin(hint)
|
||||
if origin in (list, typing.List):
|
||||
(item,) = typing.get_args(hint) or (None,)
|
||||
if item in _SCALARS:
|
||||
return f"{_SCALARS[item]}[]"
|
||||
raise TypeInferenceError(
|
||||
f"unsupported list item type {item!r}; use an explicit returns="
|
||||
)
|
||||
fields = _struct_fields(hint)
|
||||
if fields is not None:
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"STRUCT({inner})"
|
||||
raise TypeInferenceError(
|
||||
f"cannot infer a SQL type for {hint!r}; pass an explicit type string"
|
||||
)
|
||||
|
||||
|
||||
def _struct_fields(hint):
|
||||
"""(name, hint) pairs for a TypedDict or dataclass, else None."""
|
||||
if dataclasses.is_dataclass(hint):
|
||||
return [(f.name, f.type) for f in dataclasses.fields(hint)]
|
||||
# TypedDict detection: a dict subclass with __annotations__.
|
||||
if (
|
||||
isinstance(hint, type)
|
||||
and issubclass(hint, dict)
|
||||
and typing.get_type_hints(hint)
|
||||
):
|
||||
return list(typing.get_type_hints(hint).items())
|
||||
return None
|
||||
|
||||
|
||||
def return_type(fn, override: "str | None", table: bool) -> str:
|
||||
"""SQL return type for a function: explicit override wins, else the
|
||||
return annotation. Table functions render as TABLE(...) and accept
|
||||
struct-shaped hints (TypedDict/dataclass, optionally list-wrapped)."""
|
||||
if override is not None:
|
||||
s = override.strip()
|
||||
if table and not s.upper().startswith("TABLE"):
|
||||
if s.upper().startswith("STRUCT"):
|
||||
return "TABLE" + s[len("STRUCT") :]
|
||||
raise TypeInferenceError(
|
||||
"a table function's returns= must be TABLE(...) or STRUCT(...)"
|
||||
)
|
||||
return s
|
||||
|
||||
hints = typing.get_type_hints(fn)
|
||||
ret = hints.get("return")
|
||||
if ret is None:
|
||||
raise TypeInferenceError(
|
||||
f"function {fn.__name__!r} needs a return annotation or returns="
|
||||
)
|
||||
if table:
|
||||
# Accept list[Row] / Row where Row is a TypedDict or dataclass.
|
||||
if typing.get_origin(ret) in (list, typing.List):
|
||||
(ret,) = typing.get_args(ret)
|
||||
fields = _struct_fields(ret)
|
||||
if fields is None:
|
||||
raise TypeInferenceError(
|
||||
"a table function must return rows shaped as a TypedDict or "
|
||||
"dataclass (optionally list-wrapped); or pass returns=..."
|
||||
)
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"TABLE({inner})"
|
||||
return sql_type(ret)
|
||||
|
||||
|
||||
def param_types(fn) -> "list[tuple[str, str]]":
|
||||
"""(name, sql type) per parameter, from annotations. Each UDF
|
||||
parameter binds to a source column of the same name by default."""
|
||||
hints = typing.get_type_hints(fn)
|
||||
out = []
|
||||
for name, p in inspect.signature(fn).parameters.items():
|
||||
if p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD):
|
||||
raise TypeInferenceError("*args/**kwargs are not supported in UDFs")
|
||||
hint = hints.get(name)
|
||||
if hint is None:
|
||||
raise TypeInferenceError(
|
||||
f"parameter {name!r} of {fn.__name__!r} needs a type annotation"
|
||||
)
|
||||
out.append((name, sql_type(hint)))
|
||||
return out
|
||||
|
||||
|
||||
# -- column expressions -------------------------------------------------
|
||||
|
||||
|
||||
class ColumnExpr(str):
|
||||
"""A computed-column expression produced by applying a registered
|
||||
function to column names, e.g. ``embed("data") -> "embed(data)"``.
|
||||
|
||||
It IS the expression string everywhere a string is expected (views, SQL,
|
||||
logging), and additionally carries the function's declared return type so
|
||||
``add_columns(computed=...)`` can declare the column without a hand-written
|
||||
type. ``field_types`` holds the per-field SQL types of a STRUCT return, for
|
||||
fanning one expression out to several columns.
|
||||
"""
|
||||
|
||||
data_type: "str | None"
|
||||
field_types: "list[str] | None"
|
||||
|
||||
def __new__(cls, expr: str, data_type=None, field_types=None):
|
||||
obj = super().__new__(cls, expr)
|
||||
obj.data_type = data_type
|
||||
obj.field_types = field_types
|
||||
return obj
|
||||
|
||||
|
||||
def _normalize_computed(computed: dict) -> dict:
|
||||
"""Normalize the user-facing ``computed=`` mapping to the canonical
|
||||
``{name: (sql_type, expression)}`` form.
|
||||
|
||||
Accepts, per entry:
|
||||
- value is a `ColumnExpr` (from ``fn("col")``): the column's SQL type
|
||||
comes from the function's return type -- no hand-written type needed. A
|
||||
tuple key (``("chunk", "idx")``) fans a STRUCT return out to one
|
||||
(type, expression) entry per field, in declared order.
|
||||
- value is a legacy ``(sql_type, expression)`` tuple: passed through (the
|
||||
escape hatch, e.g. bare-name function strings).
|
||||
"""
|
||||
out: dict = {}
|
||||
for key, val in computed.items():
|
||||
if isinstance(val, ColumnExpr):
|
||||
expr = str(val)
|
||||
if isinstance(key, (tuple, list)):
|
||||
if not val.field_types:
|
||||
raise ValueError(
|
||||
f"columns {tuple(key)} need a STRUCT-returning function; "
|
||||
f"{expr} returns a single value"
|
||||
)
|
||||
if len(val.field_types) != len(key):
|
||||
raise ValueError(
|
||||
f"{len(key)} columns but {len(val.field_types)} struct fields "
|
||||
f"in {expr}"
|
||||
)
|
||||
for name, t in zip(key, val.field_types):
|
||||
out[name] = (t, expr)
|
||||
else:
|
||||
if val.data_type is None:
|
||||
raise ValueError(f"cannot infer a type for {expr}; pass types=")
|
||||
out[key] = (val.data_type, expr)
|
||||
else:
|
||||
out[key] = val
|
||||
return out
|
||||
|
||||
|
||||
# -- the @udf / @table_udf decorators -----------------------------------
|
||||
|
||||
|
||||
class Udf:
|
||||
def __init__(
|
||||
self,
|
||||
fn,
|
||||
*,
|
||||
returns: "str | None" = None,
|
||||
table: bool = False,
|
||||
name: "str | None" = None,
|
||||
pip: "list[str] | None" = None,
|
||||
pip_index_url: "str | None" = None,
|
||||
pip_extra_index_urls: "list[str] | None" = None,
|
||||
find_links: "list[str] | None" = None,
|
||||
requirements: "str | list[str] | None" = None,
|
||||
conda: "list[str] | None" = None,
|
||||
conda_channels: "list[str] | None" = None,
|
||||
env: "dict[str, str] | list[str] | None" = None,
|
||||
num_cpus: "int | None" = None,
|
||||
num_gpus: "int | None" = None,
|
||||
batch_size: "int | None" = None,
|
||||
timeout: "float | None" = None,
|
||||
error_policy: "str | None" = None,
|
||||
max_skip_ratio: "float | None" = None,
|
||||
retries: "int | None" = None,
|
||||
docker_image: "str | None" = None,
|
||||
description: "str | None" = None,
|
||||
prefer_source: bool = False,
|
||||
):
|
||||
functools.update_wrapper(self, fn)
|
||||
self.fn = fn
|
||||
self.name = name or fn.__name__
|
||||
self.table = table
|
||||
self.params = param_types(fn)
|
||||
self.returns = return_type(fn, returns, table)
|
||||
self.prefer_source = prefer_source
|
||||
self.options: "dict[str, str]" = {}
|
||||
if conda and (pip or requirements):
|
||||
raise ValueError("pass conda or pip/requirements, not both")
|
||||
if conda_channels and not conda:
|
||||
raise ValueError("conda_channels requires conda")
|
||||
if pip:
|
||||
self.options["pip"] = ",".join(pip)
|
||||
if pip_extra_index_urls:
|
||||
self.options["pip_extra_index_urls"] = ",".join(pip_extra_index_urls)
|
||||
if find_links:
|
||||
self.options["find_links"] = ",".join(find_links)
|
||||
if requirements:
|
||||
self.options["requirements"] = _format_requirements(requirements)
|
||||
if conda:
|
||||
self.options["conda"] = ",".join(conda)
|
||||
if conda_channels:
|
||||
self.options["conda_channels"] = ",".join(conda_channels)
|
||||
if env:
|
||||
self.options["env"] = _format_env(env)
|
||||
for key, val in [
|
||||
("pip_index_url", pip_index_url),
|
||||
("num_cpus", num_cpus),
|
||||
("num_gpus", num_gpus),
|
||||
("batch_size", batch_size),
|
||||
("timeout", timeout),
|
||||
("error_policy", error_policy),
|
||||
("max_skip_ratio", max_skip_ratio),
|
||||
("retries", retries),
|
||||
("docker_image", docker_image),
|
||||
]:
|
||||
if val is not None:
|
||||
self.options[key] = str(val)
|
||||
# Keep the source in the description (when available) so the
|
||||
# catalog stays inspectable even for pickled bodies.
|
||||
if description is not None:
|
||||
self.options["description"] = description
|
||||
else:
|
||||
try:
|
||||
self.options["description"] = textwrap.dedent(inspect.getsource(fn))
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""Call with real values to run locally; call with column-name
|
||||
strings to build an expression for backfills and views, e.g.
|
||||
``embed("data")`` -> the expression ``embed(data)`` (a `ColumnExpr`
|
||||
carrying the function's return type for `add_columns(computed=...)`)."""
|
||||
if args and all(isinstance(a, str) for a in args) and not kwargs:
|
||||
return self.expression(*args)
|
||||
return self.fn(*args, **kwargs)
|
||||
|
||||
def expression(self, *columns: str) -> ColumnExpr:
|
||||
"""The expression applying this function to `columns` (default: the
|
||||
function's own parameter names). Returns a `ColumnExpr` -- a string
|
||||
that also carries the declared return type (and struct field types)."""
|
||||
cols = columns or [p for p, _ in self.params]
|
||||
expr = f"{self.name}({', '.join(cols)})"
|
||||
field_types = None
|
||||
if self.returns.upper().startswith("STRUCT"):
|
||||
field_types = struct_field_types(self.returns)
|
||||
return ColumnExpr(expr, data_type=self.returns, field_types=field_types)
|
||||
|
||||
def _body(self) -> "tuple[str, str]":
|
||||
"""(body literal, body_format). Source when requested and
|
||||
retrievable; cloudpickle otherwise (handles closures)."""
|
||||
if self.prefer_source:
|
||||
try:
|
||||
src = textwrap.dedent(inspect.getsource(self.fn))
|
||||
# Strip the decorator line(s) so the stored body is a
|
||||
# plain function definition.
|
||||
lines = src.splitlines(keepends=True)
|
||||
while lines and lines[0].lstrip().startswith("@"):
|
||||
lines.pop(0)
|
||||
return "".join(lines), "source"
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
import cloudpickle
|
||||
|
||||
raw = cloudpickle.dumps(self.fn)
|
||||
return base64.b64encode(raw).decode("ascii"), "cloudpickle"
|
||||
|
||||
def _body_and_options(self) -> "tuple[str, dict[str, str]]":
|
||||
"""The body literal plus the finalized options (body_format /
|
||||
python_version / cloudpickle-pip bookkeeping for a non-source
|
||||
body)."""
|
||||
body, body_format = self._body()
|
||||
options = dict(self.options)
|
||||
if body_format != "source":
|
||||
options["body_format"] = body_format
|
||||
# Pickled code objects only load under the same interpreter
|
||||
# minor version; record ours so the worker can fail with a
|
||||
# clear message instead of a bytecode error.
|
||||
options["python_version"] = self.pickle_environment()
|
||||
# The worker deserializes the body with cloudpickle; make sure
|
||||
# the job's pip environment provides it. Conda bakes inject
|
||||
# cloudpickle server-side, so do not create an invalid pip+conda
|
||||
# declaration here.
|
||||
if "conda" not in options:
|
||||
pip = [d for d in options.get("pip", "").split(",") if d]
|
||||
if not any(d.startswith("cloudpickle") for d in pip):
|
||||
pip.append("cloudpickle")
|
||||
options["pip"] = ",".join(pip)
|
||||
return body, options
|
||||
|
||||
def create_request(self) -> dict:
|
||||
"""Keyword arguments for `connection.create_function`."""
|
||||
body, options = self._body_and_options()
|
||||
return {
|
||||
"name": self.name,
|
||||
"language": "python",
|
||||
"return_type": self.returns,
|
||||
"body": body,
|
||||
"options": options,
|
||||
}
|
||||
|
||||
def create_statement(self) -> str:
|
||||
"""The equivalent `CREATE FUNCTION` SQL (for SQL-surface callers)."""
|
||||
params = ", ".join(f"{n} {t}" for n, t in self.params)
|
||||
body, options = self._body_and_options()
|
||||
with_clause = ""
|
||||
if options:
|
||||
rendered = ", ".join(
|
||||
f"{k} = '{_escape(v)}'" for k, v in sorted(options.items())
|
||||
)
|
||||
with_clause = f" WITH ({rendered})"
|
||||
return (
|
||||
f"CREATE FUNCTION {self.name}({params}) RETURNS {self.returns} "
|
||||
f"LANGUAGE python AS '{_escape_body(body)}'{with_clause}"
|
||||
)
|
||||
|
||||
def pickle_environment(self) -> str:
|
||||
"""Python version the body pickles under -- workers should match
|
||||
the minor version for cloudpickle compatibility."""
|
||||
return f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
|
||||
|
||||
def _escape(s: str) -> str:
|
||||
return str(s).replace("'", "''")
|
||||
|
||||
|
||||
def _format_requirements(requirements: "str | list[str]") -> str:
|
||||
if isinstance(requirements, str):
|
||||
return requirements
|
||||
return "\n".join(str(req) for req in requirements)
|
||||
|
||||
|
||||
def _format_env(env: "dict[str, str] | list[str]") -> str:
|
||||
if isinstance(env, dict):
|
||||
return "; ".join(f"{key}={value}" for key, value in env.items())
|
||||
return "; ".join(str(entry) for entry in env)
|
||||
|
||||
|
||||
def _escape_body(body: str) -> str:
|
||||
# The server unescapes \n / \t in single-quoted bodies; encode real
|
||||
# newlines accordingly and escape quotes.
|
||||
return (
|
||||
body.replace("\\", "\\\\")
|
||||
.replace("'", "''")
|
||||
.replace("\n", "\\n")
|
||||
.replace("\t", "\\t")
|
||||
)
|
||||
|
||||
|
||||
def udf(fn=None, **kwargs):
|
||||
"""Decorate a function as a scalar (or struct-returning) UDF.
|
||||
|
||||
@udf
|
||||
def doubled(val: int) -> float: ...
|
||||
|
||||
@udf(pip=["torch>=2"], num_gpus=1)
|
||||
def embed(body: str) -> list[float]: ...
|
||||
"""
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
def table_udf(fn=None, **kwargs):
|
||||
"""Decorate a table function (UDTF): each input row may emit zero or
|
||||
more output rows. Only usable in materialized views.
|
||||
|
||||
class Chunk(TypedDict):
|
||||
chunk: str
|
||||
chunk_idx: int
|
||||
|
||||
@table_udf
|
||||
def chunker(body: str) -> list[Chunk]: ...
|
||||
"""
|
||||
kwargs["table"] = True
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
# -- view / job handles (thin references over a connection) -------------
|
||||
|
||||
|
||||
def struct_field_types(returns: str) -> "list[str]":
|
||||
"""Field type strings of a STRUCT(...) SQL type, in declared order."""
|
||||
inner = returns.strip()[len("STRUCT(") : -1]
|
||||
fields, depth, start = [], 0, 0
|
||||
for i, c in enumerate(inner):
|
||||
if c in "([":
|
||||
depth += 1
|
||||
elif c in ")]":
|
||||
depth -= 1
|
||||
elif c == "," and depth == 0:
|
||||
fields.append(inner[start:i].strip())
|
||||
start = i + 1
|
||||
fields.append(inner[start:].strip())
|
||||
# Each field is "name TYPE"; drop the name.
|
||||
return [f.split(None, 1)[1] for f in fields]
|
||||
|
||||
|
||||
def build_view_query(source, select) -> str:
|
||||
"""Assemble a view SELECT from a source (name or table) and select
|
||||
items: a column name, an expression string, a (alias, expression)
|
||||
tuple, or a @udf/@table_udf object."""
|
||||
src = source.name if hasattr(source, "name") else source
|
||||
items = []
|
||||
for item in select:
|
||||
if isinstance(item, Udf):
|
||||
items.append(item.expression())
|
||||
elif isinstance(item, tuple):
|
||||
alias, expr = item
|
||||
expr = expr.expression() if isinstance(expr, Udf) else expr
|
||||
items.append(f"{expr} AS {alias}")
|
||||
else:
|
||||
items.append(item)
|
||||
return f"SELECT {', '.join(items)} FROM {src}"
|
||||
|
||||
|
||||
def _job_id_matches(handle_id: str, listed_id: str) -> bool:
|
||||
# The refresh/backfill endpoints return the submission id (a uuid), but
|
||||
# the agent names the manifest job "<table>-<type>-<first 8 of the
|
||||
# submission id>" -- which is what list_jobs and cancel report. Match the
|
||||
# canonical id directly, or by that submission prefix.
|
||||
if listed_id == handle_id:
|
||||
return True
|
||||
prefix = handle_id[:8]
|
||||
return len(prefix) >= 4 and prefix in listed_id
|
||||
|
||||
|
||||
class MaterializedView:
|
||||
"""A reference to a materialized view (name + connection). Operations are
|
||||
server-backed connection calls bound to the name.
|
||||
|
||||
``create_materialized_view`` returns one of these; ``job_id`` is the
|
||||
initial-population job (None when the view was created with no data), so
|
||||
``db.create_materialized_view(...).wait()`` blocks until it is populated.
|
||||
"""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return Job(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
def refresh(self, full: bool = False) -> "Job":
|
||||
"""Refresh the materialized view; returns a `Job` to wait on,
|
||||
poll, or cancel (``view.refresh().wait()``).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh. A full rebuild preserves
|
||||
the view's indexes -- they are reindexed by the distributed indexer.
|
||||
"""
|
||||
job_id = self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return Job(self.conn, job_id, table=self.name)
|
||||
|
||||
def explain_refresh(self, full: bool = False):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
def alter(self, auto_refresh: bool) -> None:
|
||||
self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
def drop(self) -> None:
|
||||
self.conn.drop_materialized_view(self.name)
|
||||
|
||||
# A materialized view is a first-class table: it can be indexed and
|
||||
# searched like any other. These open the materialized dataset by name and
|
||||
# delegate. Indexes declared this way are recorded against the view, so the
|
||||
# engine re-applies them after a full refresh rebuilds the dataset (a full
|
||||
# refresh overwrites the dataset, which would otherwise drop its indices).
|
||||
def _table(self):
|
||||
return self.conn.open_table(self.name)
|
||||
|
||||
def create_index(self, *args, **kwargs):
|
||||
"""Build an index on the materialized view (see Table.create_index)."""
|
||||
return self._table().create_index(*args, **kwargs)
|
||||
|
||||
def create_scalar_index(self, *args, **kwargs):
|
||||
"""Build a scalar index on the materialized view."""
|
||||
return self._table().create_scalar_index(*args, **kwargs)
|
||||
|
||||
def create_fts_index(self, *args, **kwargs):
|
||||
"""Build a full-text-search index on the materialized view."""
|
||||
return self._table().create_fts_index(*args, **kwargs)
|
||||
|
||||
def search(self, *args, **kwargs):
|
||||
"""Search the materialized view (vector / FTS / hybrid)."""
|
||||
return self._table().search(*args, **kwargs)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or one of its columns). Delegates
|
||||
to the backing table; the server already includes the view's sources
|
||||
and downstream dependents. Returns a `Lineage`."""
|
||||
return self._table().lineage(column, direction=direction, depth=depth)
|
||||
|
||||
|
||||
_PROGRESS = re.compile(r"(\d+)/(\d+)")
|
||||
|
||||
|
||||
class JobFailedError(RuntimeError):
|
||||
"""Raised by ``Job.wait()`` when the server reports the job ``failed``.
|
||||
|
||||
Carries the server-side error so a doomed backfill (e.g. a multi-column
|
||||
``REFRESH COLUMN`` of a scalar UDF) surfaces its real cause promptly,
|
||||
instead of the caller blocking until ``wait()``'s timeout.
|
||||
"""
|
||||
|
||||
def __init__(self, job_id: str, error: "str | None"):
|
||||
self.job_id = job_id
|
||||
self.error = error
|
||||
super().__init__(f"job {job_id} failed: {error or 'unknown error'}")
|
||||
|
||||
|
||||
class Job:
|
||||
"""A reference to a server-side job, backed by the platform jobs API.
|
||||
|
||||
Holds the submission (manifest) id and resolves the platform job id
|
||||
lazily; ``status``/``progress``/``wait`` read the registry-backed
|
||||
describe endpoint, so terminal states and errors are first-class.
|
||||
"""
|
||||
|
||||
#: How long an unresolved job is treated as still materializing
|
||||
#: (submission -> dispatch -> registry record is async).
|
||||
GRACE_SECONDS = 20.0
|
||||
|
||||
#: Platform lifecycle state -> the user-facing vocabulary.
|
||||
_STATES = {
|
||||
"IN_PROGRESS": "running",
|
||||
"DONE": "finished",
|
||||
"FAILED": "failed",
|
||||
"CANCELLED": "cancelled",
|
||||
}
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
#: The submission (manifest) id the launching call handed out.
|
||||
self.id = job_id
|
||||
#: The job's table, when known -- narrows platform-id resolution.
|
||||
self.table = table
|
||||
self._platform_id: "str | None" = None
|
||||
self._created = time.monotonic()
|
||||
self._finished = False
|
||||
|
||||
@classmethod
|
||||
def _completed(cls, conn=None, table: "str | None" = None) -> "Job":
|
||||
"""A job for work that completed synchronously within the call that
|
||||
returned it (native tables, scalar/FTS builds). ``status``/``wait``
|
||||
report ``finished`` immediately and ``cancel`` is a no-op."""
|
||||
job = cls(conn, "", table)
|
||||
job._finished = True
|
||||
return job
|
||||
|
||||
def _resolve(self) -> "str | None":
|
||||
if self._platform_id is None:
|
||||
self._platform_id = self.conn.resolve_platform_job_id(self.id, self.table)
|
||||
return self._platform_id
|
||||
|
||||
def _describe(self):
|
||||
platform_id = self._resolve()
|
||||
if platform_id is None:
|
||||
return None
|
||||
return self.conn.describe_platform_job(platform_id)
|
||||
|
||||
@staticmethod
|
||||
def _payload(described) -> dict:
|
||||
# Older records carry the status-store URI string instead of a
|
||||
# payload; anything non-dict means "no structured status".
|
||||
try:
|
||||
payload = json.loads(described.status_json)
|
||||
except (TypeError, ValueError):
|
||||
return {}
|
||||
return payload if isinstance(payload, dict) else {}
|
||||
|
||||
def status(self) -> str:
|
||||
"""pending / running / finished / failed / cancelled (or unknown
|
||||
when the job never appeared in the registry)."""
|
||||
if self._finished:
|
||||
return "finished"
|
||||
described = self._describe()
|
||||
if described is not None:
|
||||
return self._STATES.get(described.job_state, described.job_state)
|
||||
if time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "unknown"
|
||||
|
||||
def progress(self) -> "tuple[int, int] | None":
|
||||
"""(units_done, units_total) once workers have published progress."""
|
||||
if self._finished:
|
||||
return None
|
||||
described = self._describe()
|
||||
if described is None:
|
||||
return None
|
||||
payload = self._payload(described)
|
||||
if payload.get("units_total") is not None:
|
||||
return payload.get("units_done") or 0, payload["units_total"]
|
||||
return None
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
described = self._describe()
|
||||
if described is None:
|
||||
if time.monotonic() - self._created > self.GRACE_SECONDS:
|
||||
raise JobFailedError(
|
||||
self.id,
|
||||
"job did not appear in the job registry within the "
|
||||
"grace period",
|
||||
)
|
||||
time.sleep(min(poll, 0.5))
|
||||
continue
|
||||
state = self._STATES.get(described.job_state, described.job_state)
|
||||
if state == "finished":
|
||||
return state
|
||||
if state == "cancelled":
|
||||
return state
|
||||
if state == "failed":
|
||||
raise JobFailedError(self.id, self._payload(described).get("error"))
|
||||
time.sleep(poll)
|
||||
raise TimeoutError(f"job {self.id} still {self.status()} after {timeout}s")
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Request cancellation. Workers drain cooperatively; poll ``status``
|
||||
for the terminal ``cancelled``."""
|
||||
if self._finished:
|
||||
return
|
||||
deadline = time.monotonic() + 5.0
|
||||
while (platform_id := self._resolve()) is None:
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError(
|
||||
f"job {self.id} has not registered yet; retry cancel shortly"
|
||||
)
|
||||
time.sleep(0.5)
|
||||
self.conn.cancel_platform_job(platform_id)
|
||||
|
||||
|
||||
class AsyncMaterializedView:
|
||||
"""Async reference to a materialized view (name + async connection)."""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return await AsyncJob(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
async def refresh(self, full: bool = False) -> "AsyncJob":
|
||||
"""Refresh the materialized view; returns an `AsyncJob` to wait
|
||||
on, poll, or cancel.
|
||||
|
||||
``full=True`` forces a full rebuild instead of an incremental refresh
|
||||
(indexes are preserved and reindexed by the distributed indexer).
|
||||
"""
|
||||
job_id = await self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return AsyncJob(self.conn, job_id, table=self.name)
|
||||
|
||||
async def explain_refresh(self, full: bool = False):
|
||||
return await self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
async def alter(self, auto_refresh: bool) -> None:
|
||||
await self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
async def drop(self) -> None:
|
||||
await self.conn.drop_materialized_view(self.name)
|
||||
|
||||
async def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or column). Returns a `Lineage`."""
|
||||
return await self.conn.lineage(
|
||||
self.name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
|
||||
class AsyncJob:
|
||||
"""Async reference to a server-side job, backed by the platform jobs API.
|
||||
|
||||
Same contract as `Job` with awaitable methods.
|
||||
"""
|
||||
|
||||
GRACE_SECONDS = 20.0
|
||||
_STATES = Job._STATES
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.id = job_id
|
||||
self.table = table
|
||||
self._platform_id: "str | None" = None
|
||||
self._created = time.monotonic()
|
||||
self._finished = False
|
||||
|
||||
@classmethod
|
||||
def _completed(cls, conn=None, table: "str | None" = None) -> "AsyncJob":
|
||||
"""See ``Job._completed``."""
|
||||
job = cls(conn, "", table)
|
||||
job._finished = True
|
||||
return job
|
||||
|
||||
async def _resolve(self) -> "str | None":
|
||||
if self._platform_id is None:
|
||||
self._platform_id = await self.conn.resolve_platform_job_id(
|
||||
self.id, self.table
|
||||
)
|
||||
return self._platform_id
|
||||
|
||||
async def _describe(self):
|
||||
platform_id = await self._resolve()
|
||||
if platform_id is None:
|
||||
return None
|
||||
return await self.conn.describe_platform_job(platform_id)
|
||||
|
||||
async def status(self) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
described = await self._describe()
|
||||
if described is not None:
|
||||
return self._STATES.get(described.job_state, described.job_state)
|
||||
if time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "unknown"
|
||||
|
||||
async def progress(self) -> "tuple[int, int] | None":
|
||||
if self._finished:
|
||||
return None
|
||||
described = await self._describe()
|
||||
if described is None:
|
||||
return None
|
||||
payload = Job._payload(described)
|
||||
if payload.get("units_total") is not None:
|
||||
return payload.get("units_done") or 0, payload["units_total"]
|
||||
return None
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
if self._finished:
|
||||
return "finished"
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
described = await self._describe()
|
||||
if described is None:
|
||||
if time.monotonic() - self._created > self.GRACE_SECONDS:
|
||||
raise JobFailedError(
|
||||
self.id,
|
||||
"job did not appear in the job registry within the "
|
||||
"grace period",
|
||||
)
|
||||
await asyncio.sleep(min(poll, 0.5))
|
||||
continue
|
||||
state = self._STATES.get(described.job_state, described.job_state)
|
||||
if state in ("finished", "cancelled"):
|
||||
return state
|
||||
if state == "failed":
|
||||
raise JobFailedError(self.id, Job._payload(described).get("error"))
|
||||
await asyncio.sleep(poll)
|
||||
raise TimeoutError(f"job {self.id} still {await self.status()} after {timeout}s")
|
||||
|
||||
async def cancel(self) -> None:
|
||||
if self._finished:
|
||||
return
|
||||
deadline = time.monotonic() + 5.0
|
||||
while (platform_id := await self._resolve()) is None:
|
||||
if time.monotonic() > deadline:
|
||||
raise RuntimeError(
|
||||
f"job {self.id} has not registered yet; retry cancel shortly"
|
||||
)
|
||||
await asyncio.sleep(0.5)
|
||||
await self.conn.cancel_platform_job(platform_id)
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import io
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
|
||||
from lancedb.index import FTS
|
||||
from lancedb.schema import blob_column_paths, blob_v2_column_paths
|
||||
|
||||
|
||||
def _blob_table(name, rows):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add(rows)
|
||||
return table
|
||||
|
||||
|
||||
def _blob_array(name, values):
|
||||
blob_type = lancedb.blob(name).type
|
||||
storage_type = blob_type.storage_type
|
||||
storage = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(values, type=pa.large_binary()),
|
||||
pa.array([None] * len(values), type=pa.string()),
|
||||
pa.array([None] * len(values), type=pa.uint64()),
|
||||
pa.array([None] * len(values), type=pa.uint64()),
|
||||
],
|
||||
fields=list(storage_type),
|
||||
)
|
||||
return pa.ExtensionArray.from_storage(blob_type, storage)
|
||||
|
||||
|
||||
def _row_ids_by_id(table):
|
||||
hits = table.search().with_row_id(True).limit(1000).to_arrow()
|
||||
assert "_rowid" in hits.column_names
|
||||
return dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
|
||||
|
||||
def test_blob_factory_declares_v2_field():
|
||||
field = lancedb.blob("image")
|
||||
assert isinstance(field.type, pa.ExtensionType)
|
||||
assert field.type.extension_name == "lance.blob.v2"
|
||||
|
||||
|
||||
def test_blob_v2_column_paths_include_list_children():
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("info", pa.struct([lancedb.blob("blob")])),
|
||||
pa.field("images", pa.list_(lancedb.blob("image"))),
|
||||
pa.field("large_images", pa.large_list(lancedb.blob("large_image"))),
|
||||
pa.field(
|
||||
"fixed_images",
|
||||
pa.list_(lancedb.blob("fixed_image"), list_size=2),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert blob_v2_column_paths(schema) == [
|
||||
"info.blob",
|
||||
"images.image",
|
||||
"large_images.large_image",
|
||||
"fixed_images.fixed_image",
|
||||
]
|
||||
|
||||
|
||||
def _legacy_v1_table(name):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field(
|
||||
"legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
)
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add([{"id": 1, "legacy": b"old"}])
|
||||
return table
|
||||
|
||||
|
||||
def test_blob_v2_column_paths_exclude_legacy_metadata():
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
lancedb.blob("image"),
|
||||
pa.field(
|
||||
"legacy", pa.large_binary(), metadata={"lance-encoding:blob": "true"}
|
||||
),
|
||||
]
|
||||
)
|
||||
assert blob_v2_column_paths(schema) == ["image"]
|
||||
assert blob_column_paths(schema) == ["image", "legacy"]
|
||||
|
||||
|
||||
def test_blob_v2_paths_match_blob_columns():
|
||||
table = _blob_table("paths_match", [{"id": 1, "image": b"x"}])
|
||||
assert blob_v2_column_paths(table.schema) == table.blob_columns()
|
||||
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
nested = db.create_table("nested_paths", data=data)
|
||||
assert blob_v2_column_paths(nested.schema) == nested.blob_columns()
|
||||
|
||||
|
||||
def test_auto_row_id_stash_round_trip():
|
||||
table = _blob_table(
|
||||
"stash_round_trip",
|
||||
[{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
|
||||
)
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
row_ids = hits["_rowid"].to_pylist()
|
||||
|
||||
stashed = stash_auto_row_ids(hits, ["image"])
|
||||
|
||||
assert "_rowid" not in stashed.column_names
|
||||
assert stashed.schema.field("image").metadata == hits.schema.field("image").metadata
|
||||
assert read_row_ids_from_hits(stashed, "image") == row_ids
|
||||
|
||||
|
||||
def test_blob_query_omits_auto_row_id():
|
||||
table = _blob_table("rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
|
||||
def test_blob_query_explicit_row_id_opt_in():
|
||||
table = _blob_table("explicit_rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
assert "_rowid" in hits.column_names
|
||||
|
||||
|
||||
def test_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
table = _blob_table("descriptions_no_leak", [{"id": 1, "image": b"x"}])
|
||||
df = table.to_pandas(blob_mode="descriptions")
|
||||
descriptor = df["image"].iloc[0]
|
||||
assert "_lance_row_id" not in descriptor
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("descriptions_no_leak_async", schema=schema)
|
||||
await table.add([{"id": 1, "image": b"x"}])
|
||||
df = await table.to_pandas(blob_mode="descriptions")
|
||||
descriptor = df["image"].iloc[0]
|
||||
assert "_lance_row_id" not in descriptor
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
def test_fetch_blobs_round_trip():
|
||||
table = _blob_table(
|
||||
"round_trip",
|
||||
[{"id": 1, "image": b"alpha"}, {"id": 2, "image": b"beta"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
|
||||
|
||||
|
||||
def test_fetch_blobs_accepts_query_result():
|
||||
table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"gamma"}
|
||||
|
||||
|
||||
def test_fetch_blobs_null_alignment():
|
||||
table = _blob_table(
|
||||
"nulls",
|
||||
[{"id": 1, "image": b"present"}, {"id": 2, "image": None}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
request = [by_id[1], by_id[2], by_id[1]]
|
||||
blobs = table.fetch_blobs("image", request)
|
||||
assert len(blobs) == len(request)
|
||||
assert blobs[0].as_py() == b"present"
|
||||
assert blobs[1].as_py() is None
|
||||
assert blobs[2].as_py() == b"present"
|
||||
|
||||
|
||||
def test_fetch_blobs_nested_path():
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first", "second"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested-alpha", b"nested-beta"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested", data=data)
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("info.blob", [by_id[1], by_id[2]])
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"nested-alpha", b"nested-beta"]
|
||||
|
||||
|
||||
def test_fetch_blob_files_lazy_read():
|
||||
payload = b"lazy-read" * 100
|
||||
table = _blob_table("lazy", [{"id": 1, "image": payload}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
handles = table.fetch_blob_files("image", [by_id[1]])
|
||||
assert len(handles) == 1
|
||||
assert handles[0].read() == payload
|
||||
|
||||
|
||||
def test_fetch_blob_files_null_alignment():
|
||||
table = _blob_table(
|
||||
"lazy_nulls",
|
||||
[{"id": 1, "image": b"here"}, {"id": 2, "image": None}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
handles = table.fetch_blob_files("image", [by_id[2], by_id[1]])
|
||||
assert len(handles) == 2
|
||||
assert handles[0] is None
|
||||
assert handles[1].read() == b"here"
|
||||
|
||||
|
||||
def test_fetch_blobs_rejects_non_blob_column():
|
||||
table = _blob_table("reject", [{"id": 1, "image": b"x"}])
|
||||
with pytest.raises(ValueError, match="not a blob column"):
|
||||
table.fetch_blobs("id", [0])
|
||||
|
||||
|
||||
def test_legacy_v1_query_omits_auto_row_id():
|
||||
table = _legacy_v1_table("legacy_v1")
|
||||
hits = table.search().select(["legacy"]).limit(10).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
|
||||
def test_fetch_blobs_rejects_legacy_v1_column():
|
||||
table = _legacy_v1_table("legacy_fetch")
|
||||
with pytest.raises(ValueError, match="legacy blob column.*blob v2"):
|
||||
table.fetch_blobs("legacy", [0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_fetch_blob_files_lazy_read():
|
||||
db = await lancedb.connect_async("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = await db.create_table("async_lazy", schema=schema)
|
||||
payload = b"async-lazy" * 100
|
||||
await table.add([{"id": 1, "image": payload}])
|
||||
hits = (
|
||||
await table.query().select({"image_alias": "image"}).limit(10).to_arrow()
|
||||
).combine_chunks()
|
||||
assert "_rowid" not in hits.column_names
|
||||
handles = await table.fetch_blob_files("image", hits)
|
||||
assert len(handles) == 1
|
||||
assert await handles[0].aread() == payload
|
||||
|
||||
|
||||
def test_fetch_blobs_from_query_result_without_row_id_raises():
|
||||
table = _blob_table("no_rowid", [{"id": 1, "image": b"x"}])
|
||||
hits = table.search().select(["id"]).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
with pytest.raises(ValueError, match="_rowid"):
|
||||
table.fetch_blobs("image", hits)
|
||||
|
||||
|
||||
_HYBRID_BLOB_SCHEMA = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("text", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("image"),
|
||||
]
|
||||
)
|
||||
_HYBRID_BLOB_ROWS = [
|
||||
{"id": 1, "text": "hello alpha", "vector": [1.0, 0.0], "image": b"alpha"},
|
||||
{"id": 2, "text": "hello beta", "vector": [0.9, 0.1], "image": b"beta"},
|
||||
{"id": 3, "text": "other", "vector": [0.0, 1.0], "image": b"other"},
|
||||
]
|
||||
|
||||
|
||||
def _hybrid_blob_table(db):
|
||||
table = db.create_table("hybrid_blob_fetch", schema=_HYBRID_BLOB_SCHEMA)
|
||||
table.add(_HYBRID_BLOB_ROWS)
|
||||
table.create_index("text", config=FTS(with_position=False))
|
||||
return table
|
||||
|
||||
|
||||
async def _hybrid_blob_table_async(db):
|
||||
table = await db.create_table("hybrid_blob_fetch_async", schema=_HYBRID_BLOB_SCHEMA)
|
||||
await table.add(_HYBRID_BLOB_ROWS)
|
||||
await table.create_index("text", config=FTS(with_position=False))
|
||||
return table
|
||||
|
||||
|
||||
def test_blob_v2_hybrid_fetch_blobs():
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
hits = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("image").type.names
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_blob_v2_hybrid_fetch_blobs_async():
|
||||
db = await lancedb.connect_async("memory:///hybrid_blob_fetch_async")
|
||||
table = await _hybrid_blob_table_async(db)
|
||||
hits = await (
|
||||
table.query()
|
||||
.nearest_to([1.0, 0.0])
|
||||
.nearest_to_text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("image").type.names
|
||||
blobs = await table.fetch_blobs("image", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
def test_blob_file_seek_read_and_read_range():
|
||||
payload = _identifiable_payload(1024)
|
||||
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
handle = table.fetch_blob_files("image", [by_id[1]])[0]
|
||||
|
||||
assert handle.seek(100) == 100
|
||||
assert handle.read(16) == payload[100:116]
|
||||
handle.seek(100)
|
||||
assert handle.read_range(500, 8) == payload[500:508]
|
||||
assert handle.tell() == 100
|
||||
|
||||
with pytest.raises(ValueError, match="whence"):
|
||||
handle.seek(0, 99)
|
||||
|
||||
|
||||
def test_fetch_blob_files_from_query_partial_read():
|
||||
payload = _identifiable_payload(65536)
|
||||
table = _blob_table("query_partial", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
assert handle.size() == 65536
|
||||
assert handle.read_range(0, 128) == payload[:128]
|
||||
assert handle.tell() == 0
|
||||
assert handle.seek(40000) == 40000
|
||||
assert handle.read(16) == payload[40000:40016]
|
||||
|
||||
|
||||
def test_blob_file_buffered_reader():
|
||||
payload = _identifiable_payload(4096)
|
||||
table = _blob_table("buffered_reader", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
reader = io.BufferedReader(handle)
|
||||
assert reader.read(8) == payload[:8]
|
||||
assert reader.read(8) == payload[8:16]
|
||||
assert reader.read() == payload[16:]
|
||||
|
||||
|
||||
def test_fetch_blob_files_cross_fragment_nulls_and_dups():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("cross_fragment", schema=schema)
|
||||
table.add([{"id": 1, "image": b"alpha"}])
|
||||
table.add([{"id": 2, "image": None}, {"id": 3, "image": b"beta"}])
|
||||
|
||||
by_id = _row_ids_by_id(table)
|
||||
request = [by_id[3], by_id[2], by_id[1], by_id[3]]
|
||||
handles = table.fetch_blob_files("image", request)
|
||||
assert len(handles) == 4
|
||||
assert handles[1] is None
|
||||
assert handles[0].read() == b"beta"
|
||||
assert handles[2].read() == b"alpha"
|
||||
assert handles[3].seek(1) == 1
|
||||
assert handles[3].read() == b"eta"
|
||||
|
||||
|
||||
def test_blob_file_pyav_decode_seek(tmp_path):
|
||||
av = pytest.importorskip("av")
|
||||
import fractions
|
||||
|
||||
clip = tmp_path / "clip.mp4"
|
||||
with av.open(str(clip), mode="w") as container:
|
||||
stream = container.add_stream("mpeg4", rate=5)
|
||||
stream.width, stream.height, stream.pix_fmt = 32, 32, "yuv420p"
|
||||
stream.time_base = fractions.Fraction(1, 5)
|
||||
for pts in range(5):
|
||||
frame = av.VideoFrame(32, 32, "yuv420p")
|
||||
frame.pts = pts
|
||||
container.mux(stream.encode(frame))
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
table = _blob_table("pyav", [{"id": 1, "image": clip.read_bytes()}])
|
||||
hits = table.search().select(["image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
|
||||
with av.open(handle) as container:
|
||||
stream = container.streams.video[0]
|
||||
container.seek(0)
|
||||
assert next(container.decode(stream)) is not None
|
||||
|
||||
|
||||
def test_blob_v2_hybrid_fetch_blob_files_seek():
|
||||
table = _hybrid_blob_table(lancedb.connect("memory:///"))
|
||||
hits = (
|
||||
table.search(query_type="hybrid")
|
||||
.vector([1.0, 0.0])
|
||||
.text("hello")
|
||||
.select(["id", "image"])
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
assert "_rowid" not in hits.column_names
|
||||
|
||||
handles = table.fetch_blob_files("image", hits)
|
||||
assert len(handles) == 2
|
||||
assert {handle.read_range(0, 2) for handle in handles} == {b"al", b"be"}
|
||||
first = handles[0]
|
||||
assert first.seek(1) == 1
|
||||
assert first.read(2) in {b"lp", b"et"}
|
||||
|
||||
|
||||
def test_blob_file_header_sniff_from_search():
|
||||
payload = b"%PDF-1.7\n" + bytes(4096)
|
||||
table = _blob_table("header_sniff", [{"id": 1, "image": payload}])
|
||||
hits = table.search().select(["id", "image"]).limit(1).to_arrow()
|
||||
handle = table.fetch_blob_files("image", hits)[0]
|
||||
assert handle.read_range(0, 4) == b"%PDF"
|
||||
assert handle.tell() == 0
|
||||
|
||||
|
||||
def test_blob_file_multiple_handles_independent_cursors():
|
||||
table = _blob_table(
|
||||
"multi_handle",
|
||||
[{"id": 1, "image": b"first-payload"}, {"id": 2, "image": b"second-payload"}],
|
||||
)
|
||||
by_id = _row_ids_by_id(table)
|
||||
first, second = table.fetch_blob_files("image", [by_id[1], by_id[2]])
|
||||
assert first.seek(6) == 6
|
||||
assert second.tell() == 0
|
||||
assert first.read(7) == b"payload"
|
||||
assert second.read(6) == b"second"
|
||||
|
||||
|
||||
def test_fetch_blob_files_nested_path_seek():
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first", "second"], type=pa.string()),
|
||||
_blob_array("blob", [b"nested-alpha", b"nested-beta"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested_seek", data=data)
|
||||
by_id = _row_ids_by_id(table)
|
||||
handle = table.fetch_blob_files("info.blob", [by_id[2]])[0]
|
||||
assert handle.seek(7) == 7
|
||||
assert handle.read() == b"beta"
|
||||
|
||||
|
||||
def test_fetch_blobs_survives_sort_after_query():
|
||||
table = _blob_table(
|
||||
"sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("image", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def test_fetch_blobs_survives_filter_and_sort_after_query():
|
||||
table = _blob_table(
|
||||
"filter_sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
filtered = hits.filter(pc.field("id") >= 2)
|
||||
sort_idx = pc.sort_indices(filtered["id"], sort_keys=[("id", "descending")])
|
||||
filtered_sorted = filtered.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("image", filtered_sorted)
|
||||
expected = [f"payload-{i}".encode() for i in filtered_sorted["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def test_fetch_blob_files_survives_sort_after_query():
|
||||
table = _blob_table(
|
||||
"lazy_sort_survives",
|
||||
[{"id": i, "image": f"payload-{i}".encode()} for i in range(5)],
|
||||
)
|
||||
hits = table.search().select(["id", "image"]).to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
handles = table.fetch_blob_files("image", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [handle.read() for handle in handles] == expected
|
||||
|
||||
|
||||
def test_fetch_blobs_nested_path_survives_sort_after_query():
|
||||
db = lancedb.connect("memory:///")
|
||||
values = [f"payload-{i}".encode() for i in range(4)]
|
||||
info = pa.StructArray.from_arrays(
|
||||
[pa.array(["row"] * 4, type=pa.string()), _blob_array("blob", values)],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array(range(4), type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested_sort_survives", data=data)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
sort_idx = pc.sort_indices(hits["id"], sort_keys=[("id", "descending")])
|
||||
sorted_hits = hits.take(sort_idx)
|
||||
|
||||
blobs = table.fetch_blobs("info.blob", sorted_hits)
|
||||
expected = [f"payload-{i}".encode() for i in sorted_hits["id"].to_pylist()]
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == expected
|
||||
|
||||
|
||||
def _identifiable_payload(size: int) -> bytes:
|
||||
block = 256
|
||||
return b"".join(bytes([i % 256]) * block for i in range(size // block))
|
||||
@@ -786,6 +786,97 @@ def test_language(mem_db: DBConnection):
|
||||
assert len(results) == 0
|
||||
|
||||
|
||||
def test_tokenize_uses_simple_index_tokenizer(mem_db: DBConnection):
|
||||
data = pa.table({"text": ["Running in cafés"], "other": ["Running in cafés"]})
|
||||
table = mem_db.create_table("test_tokenize", data=data)
|
||||
table.create_index("text", config=FTS(base_tokenizer="simple"))
|
||||
|
||||
tokens = table.tokenize("Running in cafés", column="text")
|
||||
|
||||
assert [(token.text, token.position) for token in tokens] == [
|
||||
("run", 0),
|
||||
("cafe", 2),
|
||||
]
|
||||
|
||||
|
||||
def test_tokenize_uses_icu_index_tokenizer_by_name(mem_db: DBConnection):
|
||||
data = pa.table({"text": ["Hello, こんにちは世界!"]})
|
||||
table = mem_db.create_table("test_tokenize_icu", data=data)
|
||||
table.create_index(
|
||||
"text",
|
||||
config=FTS(
|
||||
base_tokenizer="icu",
|
||||
stem=False,
|
||||
remove_stop_words=False,
|
||||
),
|
||||
name="text_icu_idx",
|
||||
)
|
||||
|
||||
tokens = table.tokenize("Hello, こんにちは世界!", index_name="text_icu_idx")
|
||||
|
||||
assert [(token.text, token.position) for token in tokens] == [
|
||||
("hello", 0),
|
||||
("こんにちは", 1),
|
||||
("世界", 2),
|
||||
]
|
||||
|
||||
|
||||
def test_tokenize_requires_one_selector(mem_db: DBConnection):
|
||||
data = pa.table({"text": ["hello world"]})
|
||||
table = mem_db.create_table("test_tokenize_selector", data=data)
|
||||
table.create_index("text", config=FTS(), name="text_idx")
|
||||
|
||||
with pytest.raises(ValueError, match="Specify exactly one"):
|
||||
table.tokenize("hello")
|
||||
|
||||
with pytest.raises(ValueError, match="Specify exactly one"):
|
||||
table.tokenize("hello", column="text", index_name="text_idx")
|
||||
|
||||
|
||||
def test_tokenize_requires_fts_index(mem_db: DBConnection):
|
||||
data = pa.table({"text": ["hello world"]})
|
||||
table = mem_db.create_table("test_tokenize_no_index", data=data)
|
||||
|
||||
with pytest.raises(ValueError, match="does not have a full text search index"):
|
||||
table.tokenize("hello", column="text")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tokenize_async(async_table):
|
||||
await async_table.create_index("text", config=FTS())
|
||||
|
||||
tokens = await async_table.tokenize("Running in cafés", column="text")
|
||||
|
||||
assert [(token.text, token.position) for token in tokens] == [
|
||||
("run", 0),
|
||||
("cafe", 2),
|
||||
]
|
||||
|
||||
|
||||
def test_tokenize_uses_explicit_simple_tokenizer():
|
||||
tokens = ldb.tokenize("Running in cafés", base_tokenizer="simple")
|
||||
|
||||
assert [(token.text, token.position) for token in tokens] == [
|
||||
("run", 0),
|
||||
("cafe", 2),
|
||||
]
|
||||
|
||||
|
||||
def test_tokenize_uses_explicit_icu_tokenizer():
|
||||
tokens = ldb.tokenize(
|
||||
"Hello, こんにちは世界!",
|
||||
base_tokenizer="icu",
|
||||
stem=False,
|
||||
remove_stop_words=False,
|
||||
)
|
||||
|
||||
assert [(token.text, token.position) for token in tokens] == [
|
||||
("hello", 0),
|
||||
("こんにちは", 1),
|
||||
("世界", 2),
|
||||
]
|
||||
|
||||
|
||||
def test_fts_on_list(mem_db: DBConnection):
|
||||
data = pa.table(
|
||||
{
|
||||
@@ -1084,6 +1175,84 @@ def test_fts_query_to_json():
|
||||
assert json_str == expected
|
||||
|
||||
|
||||
def test_fts_phrase_query_is_preserved_in_query_object():
|
||||
query = LanceFtsQueryBuilder(mock.Mock(), "puppy runs").phrase_query()
|
||||
|
||||
query_object = query.to_query_object()
|
||||
|
||||
assert query_object.full_text_query.query == '"puppy runs"'
|
||||
|
||||
|
||||
def test_fts_phrase_query_execution_preserves_user_text():
|
||||
table = mock.Mock()
|
||||
table.schema = pa.schema([])
|
||||
table._execute_query.return_value = pa.table({"text": ["result"]}).to_reader()
|
||||
|
||||
class CapturingReranker:
|
||||
score = "relevance"
|
||||
|
||||
def __init__(self):
|
||||
self.queries = []
|
||||
|
||||
def rerank_fts(self, query, results):
|
||||
self.queries.append(query)
|
||||
return results.append_column("_relevance_score", [[1.0]])
|
||||
|
||||
reranker = CapturingReranker()
|
||||
query = (
|
||||
LanceFtsQueryBuilder(table, "puppy runs")
|
||||
.phrase_query()
|
||||
.with_row_id(False)
|
||||
.rerank(reranker)
|
||||
)
|
||||
|
||||
query.to_arrow()
|
||||
|
||||
backend_query = table._execute_query.call_args.args[0]
|
||||
assert (
|
||||
backend_query.full_text_query.query,
|
||||
reranker.queries,
|
||||
query._query,
|
||||
) == ('"puppy runs"', ["puppy runs"], "puppy runs")
|
||||
|
||||
|
||||
def test_fts_phrase_query_false_preserves_string():
|
||||
query = LanceFtsQueryBuilder(mock.Mock(), "puppy runs").phrase_query(False)
|
||||
|
||||
query_object = query.to_query_object()
|
||||
|
||||
assert query_object.full_text_query.query == "puppy runs"
|
||||
|
||||
|
||||
def test_fts_phrase_query_preserves_fully_quoted_string():
|
||||
query = LanceFtsQueryBuilder(mock.Mock(), '"puppy runs"').phrase_query()
|
||||
|
||||
query_object = query.to_query_object()
|
||||
|
||||
assert query_object.full_text_query.query == '"puppy runs"'
|
||||
|
||||
|
||||
def test_fts_phrase_query_preserves_structured_phrase_query():
|
||||
phrase_query = PhraseQuery("puppy runs", "text")
|
||||
query = LanceFtsQueryBuilder(mock.Mock(), phrase_query).phrase_query()
|
||||
|
||||
query_object = query.to_query_object()
|
||||
|
||||
assert query_object.full_text_query.query == phrase_query
|
||||
|
||||
|
||||
def test_fts_phrase_query_rejects_other_structured_queries():
|
||||
query = LanceFtsQueryBuilder(
|
||||
mock.Mock(), MatchQuery("puppy", "text")
|
||||
).phrase_query()
|
||||
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=r"phrase_query\(\) requires a string or PhraseQuery, got MatchQuery",
|
||||
):
|
||||
query.to_query_object()
|
||||
|
||||
|
||||
def test_fts_fast_search(table):
|
||||
table.create_fts_index("text")
|
||||
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Unit tests for GeminiText embedding function."""
|
||||
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
# Mock google.genai modules before they are imported by gemini_text.py
|
||||
mock_google = MagicMock()
|
||||
mock_genai = MagicMock()
|
||||
mock_types = MagicMock()
|
||||
|
||||
mock_google.genai = mock_genai
|
||||
mock_genai.types = mock_types
|
||||
|
||||
sys.modules["google"] = mock_google
|
||||
sys.modules["google.genai"] = mock_genai
|
||||
sys.modules["google.genai.types"] = mock_types
|
||||
|
||||
import pytest # noqa: E402
|
||||
import numpy as np # noqa: E402
|
||||
from lancedb.embeddings import get_registry # noqa: E402
|
||||
from lancedb import __version__ # noqa: E402
|
||||
|
||||
|
||||
class TestGeminiText:
|
||||
"""Tests for GeminiText model registration, configuration, and execution."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_mocks(self):
|
||||
"""Set up standard mocks for google-genai Client and Config."""
|
||||
# Reset mocks
|
||||
mock_genai.reset_mock()
|
||||
mock_types.reset_mock()
|
||||
|
||||
self.mock_client = MagicMock()
|
||||
mock_genai.Client.return_value = self.mock_client
|
||||
|
||||
# Mock response for embed_content
|
||||
self.mock_embedding_1 = MagicMock()
|
||||
self.mock_embedding_1.values = [0.1] * 768
|
||||
self.mock_embedding_2 = MagicMock()
|
||||
self.mock_embedding_2.values = [0.2] * 768
|
||||
|
||||
self.mock_response = MagicMock()
|
||||
self.mock_response.embeddings = [self.mock_embedding_1, self.mock_embedding_2]
|
||||
self.mock_client.models.embed_content.return_value = self.mock_response
|
||||
|
||||
def test_gemini_registered(self):
|
||||
"""Test that gemini-text is registered in the embedding function registry."""
|
||||
registry = get_registry()
|
||||
assert registry.get("gemini-text") is not None
|
||||
|
||||
def test_client_init_headers(self):
|
||||
"""Test that Client is initialized with the partner-attribution header."""
|
||||
with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}):
|
||||
with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"):
|
||||
registry = get_registry()
|
||||
func = registry.get("gemini-text").create()
|
||||
|
||||
# Access the client property to trigger initialization
|
||||
_ = func.client
|
||||
|
||||
mock_genai.Client.assert_called_once_with(
|
||||
api_key="test-key",
|
||||
http_options={
|
||||
"headers": {
|
||||
"x-goog-api-client": f"lancedb/{__version__}",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
def test_generate_embeddings_batched(self):
|
||||
"""Test that multiple texts are sent in a single batched API request."""
|
||||
with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}):
|
||||
with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"):
|
||||
registry = get_registry()
|
||||
func = registry.get("gemini-text").create()
|
||||
|
||||
texts = ["hello", "world"]
|
||||
embeddings = func.generate_embeddings(texts)
|
||||
|
||||
# Check embed_content was called exactly once
|
||||
self.mock_client.models.embed_content.assert_called_once()
|
||||
|
||||
# Verify call arguments
|
||||
call_kwargs = self.mock_client.models.embed_content.call_args.kwargs
|
||||
assert call_kwargs["model"] == "gemini-embedding-001"
|
||||
assert len(call_kwargs["contents"]) == 2
|
||||
assert call_kwargs["contents"][0] == {"parts": [{"text": "hello"}]}
|
||||
assert call_kwargs["contents"][1] == {"parts": [{"text": "world"}]}
|
||||
|
||||
# Verify returns are correct numpy arrays
|
||||
assert len(embeddings) == 2
|
||||
assert isinstance(embeddings[0], np.ndarray)
|
||||
assert embeddings[0].shape == (768,)
|
||||
assert np.allclose(embeddings[0], 0.1)
|
||||
assert np.allclose(embeddings[1], 0.2)
|
||||
|
||||
def test_generate_embeddings_retrieval_document(self):
|
||||
"""Test that retrieval_document task type prepends the document title part."""
|
||||
with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}):
|
||||
with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"):
|
||||
registry = get_registry()
|
||||
func = registry.get("gemini-text").create(
|
||||
source_task_type="retrieval_document"
|
||||
)
|
||||
|
||||
texts = ["doc text"]
|
||||
|
||||
# We need mock to return only 1 embedding since we only pass 1 text
|
||||
mock_embedding = MagicMock()
|
||||
mock_embedding.values = [0.3] * 768
|
||||
self.mock_response.embeddings = [mock_embedding]
|
||||
|
||||
embeddings = func.generate_embeddings(
|
||||
texts, task_type="retrieval_document"
|
||||
)
|
||||
|
||||
# Check call arguments for retrieval_document
|
||||
call_kwargs = self.mock_client.models.embed_content.call_args.kwargs
|
||||
assert call_kwargs["contents"][0] == {
|
||||
"parts": [{"text": "Embedding of a document"}, {"text": "doc text"}]
|
||||
}
|
||||
mock_types.EmbedContentConfig.assert_called_once_with(
|
||||
output_dimensionality=768, task_type="RETRIEVAL_DOCUMENT"
|
||||
)
|
||||
|
||||
assert len(embeddings) == 1
|
||||
assert np.allclose(embeddings[0], 0.3)
|
||||
|
||||
def test_custom_dimension(self):
|
||||
"""Test that custom dimension (dim) can be configured and passed to config."""
|
||||
with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}):
|
||||
with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"):
|
||||
registry = get_registry()
|
||||
func = registry.get("gemini-text").create(dim=3072)
|
||||
|
||||
assert func.ndims() == 3072
|
||||
|
||||
texts = ["hello"]
|
||||
mock_embedding = MagicMock()
|
||||
mock_embedding.values = [0.5] * 3072
|
||||
self.mock_response.embeddings = [mock_embedding]
|
||||
|
||||
_ = func.generate_embeddings(texts)
|
||||
|
||||
mock_types.EmbedContentConfig.assert_called_once_with(
|
||||
output_dimensionality=3072
|
||||
)
|
||||
|
||||
def test_generate_embeddings_chunked(self):
|
||||
"""Test that generate_embeddings chunks texts into groups of 100."""
|
||||
with patch.dict("os.environ", {"GOOGLE_API_KEY": "test-key"}):
|
||||
with patch("lancedb.embeddings.gemini_text.attempt_import_or_raise"):
|
||||
registry = get_registry()
|
||||
func = registry.get("gemini-text").create()
|
||||
|
||||
# Passing 250 texts should make 3 calls (100, 100, 50)
|
||||
texts = [f"text_{i}" for i in range(250)]
|
||||
|
||||
# Mock client response to return correct number of embeddings per chunk
|
||||
def mock_embed_side_effect(model, contents, config=None):
|
||||
mock_resp = MagicMock()
|
||||
mock_embeddings = []
|
||||
for _ in contents:
|
||||
emb = MagicMock()
|
||||
# Each embedding is length 768
|
||||
emb.values = [0.1] * 768
|
||||
mock_embeddings.append(emb)
|
||||
mock_resp.embeddings = mock_embeddings
|
||||
return mock_resp
|
||||
|
||||
self.mock_client.models.embed_content.side_effect = (
|
||||
mock_embed_side_effect
|
||||
)
|
||||
|
||||
embeddings = func.generate_embeddings(texts)
|
||||
|
||||
# embed_content should be called 3 times
|
||||
assert self.mock_client.models.embed_content.call_count == 3
|
||||
assert len(embeddings) == 250
|
||||
@@ -1,6 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import lancedb
|
||||
|
||||
from lancedb.query import LanceHybridQueryBuilder
|
||||
@@ -139,6 +141,20 @@ def test_hybrid_query_distance_range(sync_table: Table):
|
||||
assert 0.2 <= dist.as_py() <= 0.5
|
||||
|
||||
|
||||
def test_hybrid_query_applies_zero_upper_distance_bound(sync_table: Table):
|
||||
result = (
|
||||
sync_table.search(query_type="hybrid")
|
||||
.vector([0.0, 0.4])
|
||||
.text("elephant")
|
||||
.distance_range(upper_bound=0.0)
|
||||
.rerank(RRFReranker(return_score="all"))
|
||||
.limit(4)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert len(result) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hybrid_query_distance_range_async(table: AsyncTable):
|
||||
reranker = RRFReranker(return_score="all")
|
||||
@@ -177,6 +193,31 @@ async def test_analyze_plan(table: AsyncTable):
|
||||
assert "metrics=" in res
|
||||
|
||||
|
||||
def test_hybrid_phrase_query_is_preserved_in_analyze_plan():
|
||||
table = mock.Mock()
|
||||
analyzed_queries = []
|
||||
distributed_metric_modes = []
|
||||
|
||||
def capture_query(query, *, distributed_metrics="aggregate"):
|
||||
analyzed_queries.append(query)
|
||||
distributed_metric_modes.append(distributed_metrics)
|
||||
return ""
|
||||
|
||||
table._analyze_plan.side_effect = capture_query
|
||||
|
||||
(
|
||||
LanceHybridQueryBuilder(table)
|
||||
.vector([0.1, 0.2])
|
||||
.text("puppy runs")
|
||||
.phrase_query()
|
||||
.analyze_plan(distributed_metrics="full")
|
||||
)
|
||||
|
||||
assert len(analyzed_queries) == 2
|
||||
assert analyzed_queries[1].full_text_query.query == '"puppy runs"'
|
||||
assert distributed_metric_modes == ["full", "full"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def table_with_id(tmpdir_factory) -> Table:
|
||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||
|
||||
@@ -0,0 +1,209 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""Job / AsyncJob against the platform jobs API.
|
||||
|
||||
The reference resolves its submission (manifest) id to a platform job id,
|
||||
then polls describe for registry-backed state: terminal states are
|
||||
first-class (DONE / FAILED / CANCELLED), progress comes from the
|
||||
owner-written status payload, and a failed job raises ``JobFailedError``
|
||||
promptly with the server error.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from lancedb.udf import Job, AsyncJob, JobFailedError
|
||||
|
||||
|
||||
class FakeDescription:
|
||||
"""Mirror of the pyo3 PlatformJobDescription fields the Job reads."""
|
||||
|
||||
def __init__(self, job_state, status=None):
|
||||
self.job_id = "plat-1"
|
||||
self.job_type = "indexer"
|
||||
self.job_subtype = "udf"
|
||||
self.job_state = job_state
|
||||
self.creation_ms = 0
|
||||
self.status_json = json.dumps(status if status is not None else {})
|
||||
|
||||
|
||||
class FakeConn:
|
||||
"""Scripted timeline: resolve returns None until `resolve_after` calls,
|
||||
then the platform id; describe walks a list of descriptions (holding the
|
||||
last once exhausted)."""
|
||||
|
||||
def __init__(self, descriptions, resolve_after=0):
|
||||
self._descs = list(descriptions)
|
||||
self._resolve_after = resolve_after
|
||||
self.resolve_calls = 0
|
||||
self.describe_calls = 0
|
||||
self.cancelled = []
|
||||
|
||||
def resolve_platform_job_id(self, manifest_job_id, table=None):
|
||||
self.resolve_calls += 1
|
||||
if self.resolve_calls <= self._resolve_after:
|
||||
return None
|
||||
return "plat-1"
|
||||
|
||||
def describe_platform_job(self, platform_job_id):
|
||||
assert platform_job_id == "plat-1"
|
||||
snap = self._descs[min(self.describe_calls, len(self._descs) - 1)]
|
||||
self.describe_calls += 1
|
||||
return snap
|
||||
|
||||
def cancel_platform_job(self, platform_job_id):
|
||||
self.cancelled.append(platform_job_id)
|
||||
|
||||
|
||||
class AsyncFakeConn(FakeConn):
|
||||
async def resolve_platform_job_id(self, manifest_job_id, table=None):
|
||||
return FakeConn.resolve_platform_job_id(self, manifest_job_id, table)
|
||||
|
||||
async def describe_platform_job(self, platform_job_id):
|
||||
return FakeConn.describe_platform_job(self, platform_job_id)
|
||||
|
||||
async def cancel_platform_job(self, platform_job_id):
|
||||
return FakeConn.cancel_platform_job(self, platform_job_id)
|
||||
|
||||
|
||||
def test_status_maps_platform_states():
|
||||
for wire, want in [
|
||||
("IN_PROGRESS", "running"),
|
||||
("DONE", "finished"),
|
||||
("FAILED", "failed"),
|
||||
("CANCELLED", "cancelled"),
|
||||
]:
|
||||
job = Job(FakeConn([FakeDescription(wire)]), "job-1", table="t")
|
||||
assert job.status() == want
|
||||
|
||||
|
||||
def test_status_pending_before_resolution():
|
||||
job = Job(FakeConn([], resolve_after=10_000), "job-1", table="t")
|
||||
assert job.status() == "pending"
|
||||
|
||||
|
||||
def test_progress_from_status_payload():
|
||||
conn = FakeConn(
|
||||
[
|
||||
FakeDescription(
|
||||
"IN_PROGRESS",
|
||||
status={"units_done": 3, "units_total": 8, "rows_committed": 100},
|
||||
)
|
||||
]
|
||||
)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.progress() == (3, 8)
|
||||
|
||||
|
||||
def test_progress_none_for_uri_only_status():
|
||||
# Older records carry the status-store URI string, not a payload.
|
||||
desc = FakeDescription("IN_PROGRESS")
|
||||
desc.status_json = json.dumps("s3://bucket/job/job_status")
|
||||
job = Job(FakeConn([desc]), "job-1", table="t")
|
||||
assert job.progress() is None
|
||||
|
||||
|
||||
def test_wait_raises_on_failed_promptly():
|
||||
conn = FakeConn(
|
||||
[
|
||||
FakeDescription("IN_PROGRESS"),
|
||||
FakeDescription(
|
||||
"FAILED", status={"error": "multi-column backfill needs a STRUCT"}
|
||||
),
|
||||
]
|
||||
)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
job.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5 # prompt, nowhere near the 30s timeout
|
||||
assert "STRUCT" in str(exc.value)
|
||||
assert exc.value.error == "multi-column backfill needs a STRUCT"
|
||||
assert exc.value.job_id == "job-1"
|
||||
|
||||
|
||||
def test_wait_returns_finished_on_done():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")])
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
|
||||
def test_wait_returns_cancelled():
|
||||
conn = FakeConn([FakeDescription("CANCELLED")])
|
||||
job = Job(conn, "job-1", table="t")
|
||||
assert job.wait(timeout=30, poll=0.01) == "cancelled"
|
||||
|
||||
|
||||
def test_wait_raises_when_job_never_registers():
|
||||
# An unresolved job past the grace window is a lost submission, not an
|
||||
# eternal "pending" hang.
|
||||
conn = FakeConn([], resolve_after=10_000)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
job.GRACE_SECONDS = 0.05
|
||||
job._created = time.monotonic() - 1.0
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
job.wait(timeout=5, poll=0.01)
|
||||
assert "registry" in str(exc.value)
|
||||
|
||||
|
||||
def test_cancel_resolves_then_cancels():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS")], resolve_after=1)
|
||||
job = Job(conn, "job-1", table="t")
|
||||
job.cancel()
|
||||
assert conn.cancelled == ["plat-1"]
|
||||
|
||||
|
||||
def test_async_wait_raises_on_failed_promptly():
|
||||
conn = AsyncFakeConn(
|
||||
[FakeDescription("FAILED", status={"error": "boom"})],
|
||||
)
|
||||
job = AsyncJob(conn, "job-1", table="t")
|
||||
|
||||
async def run():
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
await job.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5
|
||||
assert exc.value.error == "boom"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_async_wait_returns_finished():
|
||||
conn = AsyncFakeConn([FakeDescription("IN_PROGRESS"), FakeDescription("DONE")])
|
||||
job = AsyncJob(conn, "job-1", table="t")
|
||||
|
||||
async def run():
|
||||
assert await job.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_completed_job_is_finished_without_conn():
|
||||
job = Job._completed(table="t")
|
||||
assert job.status() == "finished"
|
||||
assert job.wait(timeout=0.01) == "finished"
|
||||
assert job.progress() is None
|
||||
job.cancel() # no-op, must not touch a connection
|
||||
|
||||
|
||||
def test_completed_job_ignores_registry():
|
||||
conn = FakeConn([FakeDescription("IN_PROGRESS")])
|
||||
job = Job._completed(conn, table="t")
|
||||
assert job.wait(timeout=0.01) == "finished"
|
||||
assert conn.resolve_calls == 0
|
||||
assert conn.describe_calls == 0
|
||||
|
||||
|
||||
def test_completed_async_job_is_finished():
|
||||
async def run():
|
||||
job = AsyncJob._completed(table="t")
|
||||
assert await job.status() == "finished"
|
||||
assert await job.wait(timeout=0.01) == "finished"
|
||||
assert await job.progress() is None
|
||||
await job.cancel()
|
||||
|
||||
asyncio.run(run())
|
||||
@@ -128,21 +128,33 @@ def test_split_hash(mem_db):
|
||||
|
||||
def test_split_hash_with_discard(mem_db):
|
||||
"""Test hash-based splitting with discard weight."""
|
||||
total_rows = 1000
|
||||
tbl = mem_db.create_table(
|
||||
"test_table",
|
||||
pa.table({"id": range(100), "category": ["A", "B"] * 50, "value": range(100)}),
|
||||
pa.table(
|
||||
{
|
||||
"id": range(total_rows),
|
||||
"category": [f"category-{i}" for i in range(total_rows)],
|
||||
"value": range(total_rows),
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
permutation_tbl = (
|
||||
# Hash a high-cardinality column: "category" has only two distinct
|
||||
# values, so whether anything is discarded would hinge on where those
|
||||
# two hashes land rather than on the discard weight.
|
||||
permutation_builder(tbl)
|
||||
.split_hash(["category"], [1, 1], discard_weight=2) # Should discard ~50%
|
||||
.split_hash(["id"], [1, 1], discard_weight=2) # Should discard ~50%
|
||||
.execute()
|
||||
)
|
||||
|
||||
# Should have fewer than 100 rows due to discard
|
||||
# Should have fewer rows due to discard, but should not be empty.
|
||||
row_count = permutation_tbl.count_rows()
|
||||
assert row_count < 100
|
||||
assert row_count > 0 # But not empty
|
||||
assert 0 < row_count < total_rows
|
||||
|
||||
data = permutation_tbl.search(None).to_arrow().to_pydict()
|
||||
assert set(data["split_id"]) == {0, 1}
|
||||
|
||||
|
||||
def test_split_sequential(mem_db):
|
||||
|
||||
@@ -11,6 +11,7 @@ import lancedb
|
||||
from lancedb.db import AsyncConnection
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import get_registry, register
|
||||
from lancedb.expr import col
|
||||
from lancedb.index import FTS, IvfPq
|
||||
import lancedb.pydantic
|
||||
import numpy as np
|
||||
@@ -63,11 +64,71 @@ def _blob_query_data():
|
||||
)
|
||||
|
||||
|
||||
def _create_blob_v2_query_table(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
async def _create_blob_v2_query_table_async(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
|
||||
|
||||
def _assert_blob_bytes_projection(df):
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
|
||||
|
||||
def _blob_query_table(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, _blob_query_data())
|
||||
return _create_blob_v2_query_table(db, name)
|
||||
|
||||
|
||||
async def _blob_query_table_async(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, _blob_query_data())
|
||||
return await _create_blob_v2_query_table_async(db, name)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def table(tmpdir_factory) -> lancedb.table.Table:
|
||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||
@@ -235,10 +296,11 @@ def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
||||
pytest.importorskip("lance")
|
||||
table = tmp_db.create_table(
|
||||
"test_query_to_pandas_blob_projection", _blob_query_data()
|
||||
table = _blob_query_table(
|
||||
tmp_db, f"test_query_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
)
|
||||
|
||||
df = (
|
||||
@@ -250,9 +312,8 @@ def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
@@ -348,18 +409,6 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert lazy_df["id"].tolist() == [1]
|
||||
_assert_lazy_blob(lazy_df["blob"].iloc[0], b"one")
|
||||
|
||||
bytes_df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
assert bytes_df["id_alias"].tolist() == [3, 4]
|
||||
assert bytes_df["payload"].tolist() == [b"three", b"four"]
|
||||
assert bytes_df["double_id"].tolist() == [6, 8]
|
||||
|
||||
desc_df = await (
|
||||
table.query()
|
||||
.where("id = 1")
|
||||
@@ -371,6 +420,31 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_bytes_projection(
|
||||
tmp_db_async, blob_schema
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await _blob_query_table_async(
|
||||
tmp_db_async,
|
||||
f"test_async_query_to_pandas_blob_{blob_schema}_bytes",
|
||||
blob_schema,
|
||||
)
|
||||
|
||||
df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow(
|
||||
@@ -502,6 +576,18 @@ def test_with_row_id(table: lancedb.table.Table):
|
||||
assert rs["_rowid"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_blob_v2_query_omits_auto_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid")
|
||||
|
||||
query_obj = table.search().select(["id", "blob"]).limit(2).to_query_object()
|
||||
assert query_obj.with_row_id is None
|
||||
|
||||
rs = table.search().select(["id", "blob"]).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in rs.column_names
|
||||
assert rs["id"].to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_where_repeated_combines_with_and(table: lancedb.table.Table):
|
||||
# Calling where() more than once should AND the filters together instead of
|
||||
# silently replacing the previous one (regression test for #2649).
|
||||
@@ -1187,7 +1273,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
|
||||
query = await table_async.search("dog", query_type="fts", fts_columns="text")
|
||||
plan = await query.explain_plan()
|
||||
# Should show FTS details (issue #2465 is now fixed)
|
||||
assert "MatchQuery: column=text, query=dog" in plan
|
||||
assert "MatchQuery: column=text, query=[dog]" in plan
|
||||
assert "GlobalLimitExec" in plan # Default limit
|
||||
|
||||
# Test FTS query with limit
|
||||
@@ -1195,7 +1281,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
|
||||
"dog", query_type="fts", fts_columns="text"
|
||||
)
|
||||
plan_with_limit = await query_with_limit.limit(1).explain_plan()
|
||||
assert "MatchQuery: column=text, query=dog" in plan_with_limit
|
||||
assert "MatchQuery: column=text, query=[dog]" in plan_with_limit
|
||||
assert "GlobalLimitExec: skip=0, fetch=1" in plan_with_limit
|
||||
|
||||
# Test FTS query with offset and limit
|
||||
@@ -1203,7 +1289,7 @@ async def test_explain_plan_fts(table_async: AsyncTable):
|
||||
"dog", query_type="fts", fts_columns="text"
|
||||
)
|
||||
plan_with_offset = await query_with_offset.offset(1).limit(1).explain_plan()
|
||||
assert "MatchQuery: column=text, query=dog" in plan_with_offset
|
||||
assert "MatchQuery: column=text, query=[dog]" in plan_with_offset
|
||||
assert "GlobalLimitExec: skip=1, fetch=1" in plan_with_offset
|
||||
|
||||
|
||||
@@ -1247,7 +1333,7 @@ async def test_explain_plan_with_filters(table_async: AsyncTable):
|
||||
"dog", query_type="fts", fts_columns="text"
|
||||
)
|
||||
plan_fts_filter = await query_fts_filter.where("id = 1").explain_plan()
|
||||
assert "MatchQuery: column=text, query=dog" in plan_fts_filter
|
||||
assert "MatchQuery: column=text, query=[dog]" in plan_fts_filter
|
||||
assert "LanceRead" in plan_fts_filter
|
||||
assert "full_filter=id = Int64(1)" in plan_fts_filter # Should show filter details
|
||||
|
||||
@@ -1946,3 +2032,39 @@ def test_fast_search(tmp_path):
|
||||
# 2. Fast Search -> Should NOT include "LanceScan" (Uses Index)
|
||||
plan = table.search(q).fast_search().explain_plan(True)
|
||||
assert "LanceScan" not in plan
|
||||
|
||||
|
||||
def test_blob_v2_with_row_id_bytes_pandas(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_rowid_bytes_pandas")
|
||||
|
||||
df = (
|
||||
table.search()
|
||||
.with_row_id(True)
|
||||
.select(["id", "blob"])
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert "_rowid" in df.columns
|
||||
assert df["id"].tolist() == [1, 2, 3, 4]
|
||||
assert df["blob"].tolist() == [b"one", b"two", b"three", b"four"]
|
||||
|
||||
|
||||
def test_blob_v2_expr_projection_stash(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_expr_projection_stash")
|
||||
|
||||
hits = table.search().select({"blob_alias": col("blob")}).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
|
||||
def test_blob_v2_to_batches_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_to_batches_rowid")
|
||||
|
||||
hits = table.search().select(["id", "blob"]).limit(2).to_batches().read_all()
|
||||
|
||||
assert "_rowid" in hits.column_names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
@@ -23,6 +23,7 @@ from lancedb.rerankers import (
|
||||
AnswerdotaiRerankers,
|
||||
VoyageAIReranker,
|
||||
MRRReranker,
|
||||
WatsonxReranker,
|
||||
)
|
||||
from lancedb.table import LanceTable
|
||||
|
||||
@@ -727,3 +728,19 @@ def test_linear_combination_missing_fts_is_penalised():
|
||||
f"Document with FTS score (rowid 0, {scores[0]:.4f}) should beat "
|
||||
f"document with no FTS match (rowid 1, {scores[1]:.4f})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("WATSONX_API_KEY") is None
|
||||
or (
|
||||
os.environ.get("WATSONX_PROJECT_ID") is None
|
||||
and os.environ.get("WATSONX_SPACE_ID") is None
|
||||
),
|
||||
reason="WATSONX_API_KEY and one of WATSONX_PROJECT_ID / "
|
||||
"WATSONX_SPACE_ID must be set",
|
||||
)
|
||||
def test_watsonx_reranker(tmp_path):
|
||||
pytest.importorskip("ibm_watsonx_ai")
|
||||
table, schema = get_test_table(tmp_path)
|
||||
reranker = WatsonxReranker()
|
||||
_run_test_reranker(reranker, table, "single player experience", None, schema)
|
||||
|
||||
@@ -45,6 +45,32 @@ def _blob_test_data():
|
||||
)
|
||||
|
||||
|
||||
def _blob_v2_table(db: DBConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
async def _blob_v2_table_async(db: AsyncConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
def _blob_table(db: DBConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, data=_blob_test_data())
|
||||
return _blob_v2_table(db, name)
|
||||
|
||||
|
||||
async def _blob_table_async(db: AsyncConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, data=_blob_test_data())
|
||||
return await _blob_v2_table_async(db, name)
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
@@ -107,6 +133,18 @@ def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_table_to_pandas_blob_bytes(tmp_db: DBConnection, blob_schema):
|
||||
pytest.importorskip("lance")
|
||||
table = _blob_table(tmp_db, f"test_to_pandas_blob_{blob_schema}_bytes", blob_schema)
|
||||
|
||||
df = table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
||||
@@ -118,15 +156,20 @@ def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection):
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_table_to_pandas_blob_bytes(
|
||||
tmp_db_async: AsyncConnection, blob_schema
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await tmp_db_async.create_table(
|
||||
"test_async_to_pandas_blob_bytes", data=_blob_test_data()
|
||||
table = await _blob_table_async(
|
||||
tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
)
|
||||
|
||||
df = await table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1568,16 +1611,23 @@ def test_create_with_nans(mem_db: DBConnection):
|
||||
"fill_test",
|
||||
data=[
|
||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
||||
{"vector": [2.1, 4.1], "item": "foo", "price": 9.0},
|
||||
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
||||
{"vector": [5], "item": "bar", "price": 22.0},
|
||||
],
|
||||
on_bad_vectors="fill",
|
||||
fill_value=0.0,
|
||||
)
|
||||
assert len(table) == 3
|
||||
assert len(table) == 5
|
||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||
v = arrow_tbl["vector"].to_pylist()[0]
|
||||
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||
filled_vectors = {
|
||||
row["price"]: row["vector"]
|
||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
||||
}
|
||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||
|
||||
|
||||
def test_add_with_nans(mem_db: DBConnection):
|
||||
@@ -1620,15 +1670,21 @@ def test_add_with_nans(mem_db: DBConnection):
|
||||
data=[
|
||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
||||
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
||||
{"vector": [5], "item": "bar", "price": 22.0},
|
||||
],
|
||||
on_bad_vectors="fill",
|
||||
fill_value=0.0,
|
||||
)
|
||||
assert len(table) == 3
|
||||
assert len(table) == 4
|
||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||
v = arrow_tbl["vector"].to_pylist()[0]
|
||||
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||
filled_vectors = {
|
||||
row["price"]: row["vector"]
|
||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
||||
}
|
||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||
|
||||
|
||||
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
|
||||
@@ -1789,7 +1845,9 @@ def test_on_bad_vectors_fill_preserves_arrow_nested_vector_type(mem_db: DBConnec
|
||||
fill_value=0.0,
|
||||
)
|
||||
|
||||
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0], [0.0, 0.0]]
|
||||
vector = table.to_arrow()["vector"]
|
||||
assert vector.type == pa.list_(pa.float32())
|
||||
assert vector.to_pylist() == [[1.0, 2.0], [0.0, 3.0]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -13,6 +13,7 @@ from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
from lancedb.table import (
|
||||
_append_vector_columns,
|
||||
_cast_to_target_schema,
|
||||
_fill_bad_vector_values,
|
||||
_handle_bad_vectors,
|
||||
_into_pyarrow_reader,
|
||||
_infer_target_schema,
|
||||
@@ -287,7 +288,9 @@ def test_append_vector_columns():
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
vector = pa.array([[1.0, 2.0], [3.0], [4.0, 5.0]])
|
||||
vector = pa.array(
|
||||
[[1.0, 2.0], [3.0], [4.0, 5.0], [6.0, 7.0, 8.0], [None, 9.0], None]
|
||||
)
|
||||
schema = pa.schema({"vector": pa.list_(pa.float64())})
|
||||
data = pa.table({"vector": vector}, schema=schema)
|
||||
|
||||
@@ -313,15 +316,54 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
).read_all()
|
||||
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0]])
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array([[1.0, 2.0], [42.0, 42.0], [4.0, 5.0]])
|
||||
expected = pa.array(
|
||||
[
|
||||
[1.0, 2.0],
|
||||
[3.0, 42.0],
|
||||
[4.0, 5.0],
|
||||
[6.0, 7.0],
|
||||
[None, 9.0],
|
||||
[42.0, 42.0],
|
||||
]
|
||||
)
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0]])
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None])
|
||||
|
||||
assert output["vector"].combine_chunks() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("vector_type", "vectors", "expected"),
|
||||
[
|
||||
(
|
||||
pa.list_(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.large_list(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.list_(pa.float64(), 2),
|
||||
[[1.0, float("nan")], None, [None, 3.0]],
|
||||
[[1.0, 42.0], [42.0, 42.0], [None, 3.0]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
|
||||
arr = pa.array([[0.0, 0.0], *vectors, [9.0, 9.0]], type=vector_type)
|
||||
arr = arr.slice(1, len(vectors))
|
||||
|
||||
actual = _fill_bad_vector_values(arr, dim=2, fill_value=42.0)
|
||||
|
||||
assert actual.type == vector_type
|
||||
assert actual.to_pylist() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
|
||||
@@ -351,7 +393,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[3.0, 4.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array([[42.0, 42.0], [3.0, 4.0]])
|
||||
expected = pa.array([[1.0, 42.0], [3.0, 4.0]])
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([None, [3.0, 4.0]])
|
||||
|
||||
|
||||
+456
-1
@@ -18,7 +18,10 @@ use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
connection::NamespaceClientPushdownOperation,
|
||||
database::namespace::LanceNamespaceDatabase,
|
||||
database::{CreateTableMode, Database, ReadConsistency},
|
||||
database::{
|
||||
CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode, Database,
|
||||
ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest,
|
||||
},
|
||||
};
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
@@ -27,6 +30,107 @@ use pyo3::{
|
||||
types::{PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A registered materialized view definition.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
pub projection: Vec<String>,
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// One inflight server-side job.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// A described platform job (POST /v1/jobs/describe).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct PlatformJobDescription {
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub job_subtype: String,
|
||||
/// "IN_PROGRESS" | "CANCELLED" | "FAILED" | "DONE".
|
||||
pub job_state: String,
|
||||
pub creation_ms: i64,
|
||||
/// The owner-written status payload as a JSON string (units_done /
|
||||
/// units_total / rows_committed / error when present).
|
||||
pub status_json: String,
|
||||
}
|
||||
|
||||
/// One durable, completed/terminal server-side job record (SHOW JOB HISTORY).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobHistoryEntry {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// One per-row UDF error recorded by `error_policy=skip` (SHOW ERRORS).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobErrorEntry {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a REFRESH MATERIALIZED VIEW would execute (EXPLAIN REFRESH).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
pub rebuild: bool,
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Connection {
|
||||
inner: Option<LanceConnection>,
|
||||
@@ -310,6 +414,357 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, language, return_type, body, options=None))]
|
||||
pub fn create_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: Option<HashMap<String, String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_function(CreateFunctionRequest {
|
||||
name,
|
||||
language,
|
||||
return_type,
|
||||
body,
|
||||
options: options.unwrap_or_default(),
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let functions = inner.list_functions().await.infer_error()?;
|
||||
Ok(functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_function(self_: PyRef<'_, Self>, name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_function(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, query, auto_refresh=false, with_no_data=false, partition_by=None))]
|
||||
pub fn create_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
partition_by: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_materialized_view(CreateMaterializedViewRequest {
|
||||
name,
|
||||
query,
|
||||
auto_refresh,
|
||||
with_no_data,
|
||||
partition_by,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None, num_workers=None, max_workers=None))]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_materialized_view(RefreshMaterializedViewRequest {
|
||||
name,
|
||||
full,
|
||||
src_version,
|
||||
num_workers,
|
||||
max_workers,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), returned as the
|
||||
/// server's lineage JSON string (the Python layer parses it).
|
||||
pub fn table_lineage(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
column: Option<String>,
|
||||
direction: Option<String>,
|
||||
depth: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.table_lineage(TableLineageRequest {
|
||||
name,
|
||||
column,
|
||||
direction,
|
||||
depth,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None))]
|
||||
pub fn explain_refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let p = inner
|
||||
.explain_refresh_materialized_view(&name, full, src_version)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: p.table_name,
|
||||
has_work: p.has_work,
|
||||
source_version: p.source_version,
|
||||
last_refreshed_version: p.last_refreshed_version,
|
||||
full_refresh: p.full_refresh,
|
||||
rebuild: p.rebuild,
|
||||
units_total: p.units_total,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn alter_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
auto_refresh: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.alter_materialized_view(&name, auto_refresh)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_materialized_view(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let views = inner.list_materialized_views().await.infer_error()?;
|
||||
Ok(views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let jobs = inner.list_jobs().await.infer_error()?;
|
||||
Ok(jobs
|
||||
.into_iter()
|
||||
.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel_job(&job_id).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn describe_platform_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
platform_job_id: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let described = inner
|
||||
.describe_platform_job(&platform_job_id)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(described.map(|d| PlatformJobDescription {
|
||||
job_id: d.job_id,
|
||||
job_type: d.job_type,
|
||||
job_subtype: d.job_subtype,
|
||||
job_state: d.job_state,
|
||||
creation_ms: d.creation_ms,
|
||||
status_json: d.status.to_string(),
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (manifest_job_id, table=None))]
|
||||
pub fn resolve_platform_job_id(
|
||||
self_: PyRef<'_, Self>,
|
||||
manifest_job_id: String,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.resolve_platform_job_id(&manifest_job_id, table.as_deref())
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_platform_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
platform_job_id: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.cancel_platform_job(&platform_job_id)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id, table=None))]
|
||||
pub fn get_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: String,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let job = inner
|
||||
.get_job(&job_id, table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(job.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None))]
|
||||
pub fn job_history(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner.job_history(job_id.as_deref()).await.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| JobHistoryEntry {
|
||||
table: r.table,
|
||||
job_id: r.job_id,
|
||||
job_type: r.job_type,
|
||||
state: r.state,
|
||||
column: r.column,
|
||||
created_ms: r.created_ms,
|
||||
updated_ms: r.updated_ms,
|
||||
completed_ms: r.completed_ms,
|
||||
rows_processed: r.rows_processed,
|
||||
rows_skipped: r.rows_skipped,
|
||||
error: r.error,
|
||||
events: r.events,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None, table=None))]
|
||||
pub fn errors(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner
|
||||
.errors(job_id.as_deref(), table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|e| JobErrorEntry {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (cur_name, new_name, cur_namespace_path=None, new_namespace_path=None))]
|
||||
pub fn rename_table(
|
||||
self_: PyRef<'_, Self>,
|
||||
|
||||
+11
-2
@@ -15,8 +15,8 @@ use pyo3::{
|
||||
use query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
||||
use session::Session;
|
||||
use table::{
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, LsmWriteSpec,
|
||||
MergeResult, Table, UpdateFieldMetadataResult, UpdateResult,
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
||||
LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult,
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
@@ -42,8 +42,15 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
.write_style("LANCEDB_LOG_STYLE");
|
||||
env_logger::init_from_env(env);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<connection::FunctionInfo>()?;
|
||||
m.add_class::<connection::MaterializedViewInfo>()?;
|
||||
m.add_class::<connection::JobInfo>()?;
|
||||
m.add_class::<connection::PlatformJobDescription>()?;
|
||||
m.add_class::<connection::JobHistoryEntry>()?;
|
||||
m.add_class::<connection::JobErrorEntry>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
m.add_class::<FTSQuery>()?;
|
||||
@@ -59,6 +66,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<DeleteResult>()?;
|
||||
m.add_class::<DropColumnsResult>()?;
|
||||
m.add_class::<UpdateResult>()?;
|
||||
m.add_class::<FtsToken>()?;
|
||||
m.add_class::<PyAsyncPermutationBuilder>()?;
|
||||
m.add_class::<PyPermutationReader>()?;
|
||||
m.add_class::<PyExpr>()?;
|
||||
@@ -74,6 +82,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_function(wrap_pyfunction!(connect, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace_client, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(table::tokenize, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?;
|
||||
|
||||
+48
-8
@@ -19,6 +19,7 @@ use lancedb::index::scalar::{
|
||||
BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
|
||||
Operator, PhraseQuery,
|
||||
};
|
||||
use lancedb::query::AnalyzePlanDistributedMetrics;
|
||||
use lancedb::query::QueryBase;
|
||||
use lancedb::query::QueryExecutionOptions;
|
||||
use lancedb::query::QueryFilter;
|
||||
@@ -42,6 +43,25 @@ use pyo3::{Borrowed, FromPyObject, exceptions::PyRuntimeError};
|
||||
use pyo3::{PyErr, pyclass};
|
||||
use pyo3::{exceptions::PyValueError, intern};
|
||||
|
||||
fn analyze_plan_options(distributed_metrics: Option<&str>) -> PyResult<QueryExecutionOptions> {
|
||||
let analyze_plan_distributed_metrics = match distributed_metrics.unwrap_or("aggregate") {
|
||||
"aggregate" => AnalyzePlanDistributedMetrics::Aggregate,
|
||||
"per_worker" => AnalyzePlanDistributedMetrics::PerWorker,
|
||||
"full" => AnalyzePlanDistributedMetrics::Full,
|
||||
mode => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"Invalid distributed_metrics value '{}'. Expected one of: \
|
||||
'aggregate', 'per_worker', 'full'",
|
||||
mode
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let mut options = QueryExecutionOptions::default();
|
||||
options.analyze_plan_distributed_metrics = analyze_plan_distributed_metrics;
|
||||
Ok(options)
|
||||
}
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
|
||||
type Error = PyErr;
|
||||
|
||||
@@ -571,11 +591,16 @@ impl Query {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
#[pyo3(signature = (distributed_metrics=None))]
|
||||
pub fn analyze_plan(
|
||||
self_: PyRef<'_, Self>,
|
||||
distributed_metrics: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let options = analyze_plan_options(distributed_metrics.as_deref())?;
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.analyze_plan()
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
|
||||
})
|
||||
@@ -650,11 +675,16 @@ impl TakeQuery {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
#[pyo3(signature = (distributed_metrics=None))]
|
||||
pub fn analyze_plan(
|
||||
self_: PyRef<'_, Self>,
|
||||
distributed_metrics: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let options = analyze_plan_options(distributed_metrics.as_deref())?;
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.analyze_plan()
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
|
||||
})
|
||||
@@ -777,14 +807,19 @@ impl FTSQuery {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
#[pyo3(signature = (distributed_metrics=None))]
|
||||
pub fn analyze_plan(
|
||||
self_: PyRef<'_, Self>,
|
||||
distributed_metrics: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_
|
||||
.inner
|
||||
.clone()
|
||||
.full_text_search(self_.fts_query.clone());
|
||||
let options = analyze_plan_options(distributed_metrics.as_deref())?;
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.analyze_plan()
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
|
||||
})
|
||||
@@ -958,11 +993,16 @@ impl VectorQuery {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
#[pyo3(signature = (distributed_metrics=None))]
|
||||
pub fn analyze_plan(
|
||||
self_: PyRef<'_, Self>,
|
||||
distributed_metrics: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
let options = analyze_plan_options(distributed_metrics.as_deref())?;
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.analyze_plan()
|
||||
.analyze_plan_with_options(options)
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(e.to_string()))
|
||||
})
|
||||
|
||||
+305
-7
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use crate::runtime::future_into_py;
|
||||
use crate::runtime::{block_on, future_into_py};
|
||||
use crate::{
|
||||
connection::Connection,
|
||||
error::PythonErrorExt,
|
||||
@@ -12,19 +12,24 @@ use crate::{
|
||||
table::scannable::PyScannable,
|
||||
};
|
||||
use arrow::{
|
||||
array::{Array, LargeBinaryArray},
|
||||
datatypes::{DataType, Schema},
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||
};
|
||||
use lancedb::blob::BlobFile;
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
LoadColumnsRequest, NewColumnTransform, OptimizeAction, OptimizeOptions, Ref,
|
||||
Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
pyclass, pymethods,
|
||||
types::{IntoPyDict, PyAnyMethods, PyDict, PyDictMethods},
|
||||
pyclass, pyfunction, pymethods,
|
||||
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
mod scannable;
|
||||
@@ -412,6 +417,150 @@ impl From<lancedb::table::DropColumnsResult> for DropColumnsResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Lazy blob handle from ``Table.fetch_blob_files``.
|
||||
#[pyclass(name = "BlobFile")]
|
||||
pub struct PyBlobFile {
|
||||
inner: Arc<BlobFile>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let bytes = inner
|
||||
.read()
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind()))
|
||||
})
|
||||
}
|
||||
|
||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.close().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.is_closed().await })
|
||||
}
|
||||
|
||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.seek(position).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.tell().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
}
|
||||
|
||||
fn size(self_: PyRef<'_, Self>) -> u64 {
|
||||
self_.inner.size()
|
||||
}
|
||||
|
||||
/// Read a blob-local byte range without moving the cursor.
|
||||
fn read_range(self_: PyRef<'_, Self>, offset: u64, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let end = offset
|
||||
.checked_add(length as u64)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FtsToken {
|
||||
pub text: String,
|
||||
pub position: u32,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl FtsToken {
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!("FtsToken(text={:?}, position={})", self.text, self.position)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<LanceDbFtsToken> for FtsToken {
|
||||
fn from(token: LanceDbFtsToken) -> Self {
|
||||
Self {
|
||||
text: token.text,
|
||||
position: token.position,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction(signature = (
|
||||
query,
|
||||
*,
|
||||
base_tokenizer = "simple".to_string(),
|
||||
language = "English".to_string(),
|
||||
max_token_length = Some(40),
|
||||
lower_case = true,
|
||||
stem = true,
|
||||
remove_stop_words = true,
|
||||
ascii_folding = true,
|
||||
ngram_min_length = 3,
|
||||
ngram_max_length = 3,
|
||||
prefix_only = false
|
||||
))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn tokenize(
|
||||
query: String,
|
||||
base_tokenizer: String,
|
||||
language: String,
|
||||
max_token_length: Option<u32>,
|
||||
lower_case: bool,
|
||||
stem: bool,
|
||||
remove_stop_words: bool,
|
||||
ascii_folding: bool,
|
||||
ngram_min_length: u32,
|
||||
ngram_max_length: u32,
|
||||
prefix_only: bool,
|
||||
) -> PyResult<Vec<FtsToken>> {
|
||||
let params = FtsIndexBuilder::default()
|
||||
.base_tokenizer(base_tokenizer)
|
||||
.language(&language)
|
||||
.map_err(|_| {
|
||||
PyValueError::new_err(format!(
|
||||
"LanceDB does not support the requested language: '{}'",
|
||||
language
|
||||
))
|
||||
})?
|
||||
.max_token_length(max_token_length.map(|value| value as usize))
|
||||
.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)
|
||||
.ngram_prefix_only(prefix_only);
|
||||
let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect())
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Table {
|
||||
// We keep a copy of the name to use if the inner table is dropped
|
||||
@@ -645,8 +794,8 @@ impl Table {
|
||||
}
|
||||
|
||||
future_into_py(self_.py(), async move {
|
||||
op.execute().await.infer_error()?;
|
||||
Ok(())
|
||||
let job_id = op.execute().await.infer_error()?;
|
||||
Ok(job_id)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -710,6 +859,29 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (query, *, column=None, index_name=None))]
|
||||
pub fn tokenize(
|
||||
self_: PyRef<'_, Self>,
|
||||
query: String,
|
||||
column: Option<String>,
|
||||
index_name: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let tokens = match (column.as_deref(), index_name.as_deref()) {
|
||||
(Some(_), Some(_)) | (None, None) => {
|
||||
return Err(PyValueError::new_err(
|
||||
"Specify exactly one of 'column' or 'index_name'",
|
||||
));
|
||||
}
|
||||
(Some(column), None) => inner.tokenize_with_column(&query, column).await,
|
||||
(None, Some(index_name)) => inner.tokenize(&query, index_name).await,
|
||||
}
|
||||
.infer_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn index_stats(self_: PyRef<'_, Self>, index_name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -901,6 +1073,55 @@ impl Table {
|
||||
))
|
||||
}
|
||||
|
||||
/// Names of the blob v2 columns declared on this table, in declaration order.
|
||||
pub fn blob_columns(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.blob_columns().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Read blob bytes for `row_ids` from blob v2 column `column`.
|
||||
#[pyo3(signature = (column, row_ids))]
|
||||
pub fn fetch_blobs(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
row_ids: Vec<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let blobs: LargeBinaryArray = inner
|
||||
.fetch_blobs(column.as_str(), &row_ids)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Python::attach(|py| blobs.to_data().to_pyarrow(py).map(|obj| obj.unbind()))
|
||||
})
|
||||
}
|
||||
|
||||
/// Open lazy blob handles for `row_ids` from blob v2 column `column`.
|
||||
#[pyo3(signature = (column, row_ids))]
|
||||
pub fn fetch_blob_files(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
row_ids: Vec<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let handles = inner
|
||||
.fetch_blob_files(column.as_str(), &row_ids)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
.map(|handle| {
|
||||
handle.map(|file| PyBlobFile {
|
||||
inner: Arc::new(file),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
||||
pub fn optimize(
|
||||
@@ -1074,6 +1295,83 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_computed_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<(String, String)>,
|
||||
expression: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.add_computed_columns(&columns, &expression)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (columns, where_clause=None, num_workers=None, max_workers=None, batch_size=None, priority=None))]
|
||||
pub fn refresh_column(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<String>,
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_column(
|
||||
&columns,
|
||||
where_clause,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
priority,
|
||||
)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[pyo3(signature = (source_uris, source_format, target_key, columns, source_key=None, source_storage_options=None, on_missing=None, num_workers=None, max_workers=None, batch_size=None, commit_granularity=None, priority=None))]
|
||||
pub fn load_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
source_uris: Vec<String>,
|
||||
source_format: String,
|
||||
target_key: String,
|
||||
columns: Vec<(String, Option<String>)>,
|
||||
source_key: Option<String>,
|
||||
source_storage_options: Option<std::collections::HashMap<String, String>>,
|
||||
on_missing: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
commit_granularity: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let request = LoadColumnsRequest {
|
||||
source_uris,
|
||||
source_format,
|
||||
source_storage_options,
|
||||
target_key,
|
||||
source_key,
|
||||
columns,
|
||||
on_missing,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
commit_granularity,
|
||||
priority,
|
||||
};
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.load_columns(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
definitions: Vec<(String, String)>,
|
||||
|
||||
Generated
+40
-23
@@ -657,6 +657,15 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/0d/67e5b4109ea4a837e80daa87c2c696711955e40449a97e8926672534def2/click-8.4.1-py3-none-any.whl", hash = "sha256:482be17c6991b8c19c5429a1e995d9b0efdbb63172824c41f99965dc0ade8ec2", size = 116639, upload-time = "2026-05-22T04:08:35.26Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cloudpickle"
|
||||
version = "3.1.2"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cohere"
|
||||
version = "7.0.3"
|
||||
@@ -850,19 +859,25 @@ nvtx = [
|
||||
|
||||
[[package]]
|
||||
name = "datafusion"
|
||||
version = "52.3.0"
|
||||
version = "54.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cloudpickle" },
|
||||
{ name = "pyarrow" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/db/d4/a5ad7b665a80008901892fde61dc667318db0652a955d706ddca3a224b5a/datafusion-52.3.0.tar.gz", hash = "sha256:2e8b02ad142b1a0d673f035d96a0944a640ac78275003d7e453cee4afe4a20a4", size = 205026, upload-time = "2026-03-16T10:54:07.739Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/60/90/886f7e9cf827f07ebd60bd293e54e0a028a50dd49bbaef0ee42aae1981ea/datafusion-54.0.0.tar.gz", hash = "sha256:cfe7e8dfc026efc05824f49b53ad6a72caf5c2d6820759b6212a09e245a427ed", size = 276448, upload-time = "2026-06-29T11:19:34.816Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/55/63/1bb0737988cefa77274b459d64fa4b57ba4cf755639a39733e9581b5d599/datafusion-52.3.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a73f02406b2985b9145dd97f8221a929c9ef3289a8ba64c6b52043e240938528", size = 31503230, upload-time = "2026-03-16T10:53:50.312Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d6/e3/ea3b79239953c3044d19d8e9581015da025b6640796db03799e435b17910/datafusion-52.3.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:118a1f0add6a3f91fcbc90c71819fe08750e2981637d5e7b346e099e94a20b8b", size = 28159497, upload-time = "2026-03-16T10:53:54.032Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/24/c8/7d325feb4b7509ae03857fd7e164e95ec72e8c9f3dfd3178ec7f80d53977/datafusion-52.3.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:253ce7aee5fe84bd6ee290c20608114114bdb5115852617f97d3855d36ad9341", size = 30769154, upload-time = "2026-03-16T10:53:57.835Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/37/ee/478689c69b3cb1ccabb2d52feac0c181f6cdf20b51a81df35344b1dab9a6/datafusion-52.3.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2af3469d2f06959bec88579ab107a72f965de18b32e607069bbdd0b859ed8dbb", size = 33060335, upload-time = "2026-03-16T10:54:01.715Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1c/48/01906ab5c1a70373c6874ac5192d03646fa7b94d9ff06e3f676cb6b0f43f/datafusion-52.3.0-cp310-abi3-win_amd64.whl", hash = "sha256:9fb35738cf4dbff672dbcfffc7332813024cb0ad2ab8cda1fb90b9054277ab0c", size = 33765807, upload-time = "2026-03-16T10:54:05.728Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/46/58/4c5b981e3d9ade32a906c15a4941eef50c9b862781cdc14bf4dff48d026a/datafusion-54.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:946f55e48b8d523d7b4ac106bdf588b4493c2c66f81877d6952aafeaf7c3ec73", size = 39810553, upload-time = "2026-06-29T11:19:02.1Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/66/e5/5e4dbd42ce9a2affb3be90d9ab17cebde1a6f28b0d9fb4b83d612d5c8e42/datafusion-54.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:2a3bf43185c7e43e25242e5fb17b6a11b86bf976434c0bc493fdedbd9a080969", size = 37145255, upload-time = "2026-06-29T11:19:05.491Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c6/5e/dbb9e6e3e5006d34f295d7ac73f1302c8f2df140666402a06e6c55028edb/datafusion-54.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:9432bf162381e9282cbc74915b8b773895de18be836f7e3f6d0de4d981f24630", size = 38853856, upload-time = "2026-06-29T11:19:08.732Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a8/81/e69008e3479f4d0134875bc4ae39503bedcd55ca2597e71392c963c651b4/datafusion-54.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3bcd4d213fa74710e75e6e182cc468c2bdbc5ffc74a08c8155d414fbbfa1b3f6", size = 41050149, upload-time = "2026-06-29T11:19:12.108Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/61/d4/8ba6e3fe3291c9ccc94b5ca3ec3c1fbcbfbe5ece5ffb965e4550844e2c56/datafusion-54.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:b934e097e1bdca7d5768a81ac1bc4a1812cb459269f8b1a5d892a5d930f18376", size = 43444869, upload-time = "2026-06-29T11:19:15.963Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9d/41/5608323226f21a0fa180823c531dbc0ed270e9b694f299b7647505cb6a06/datafusion-54.0.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:c4e79048da82ad89b768bd0be7df39254cd2a0afe2b719d1f129e8a7229af683", size = 39796248, upload-time = "2026-06-29T11:19:19.208Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/18/81/392ee323104ab14ca689384723b69e137064a828233c165574f97a74c0e9/datafusion-54.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:fe57038003b18e28b90752c1e32b44af74ec4f552a1904aee725e1129a00c447", size = 37153577, upload-time = "2026-06-29T11:19:22.397Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/40/c4/ebd5ef5349ecbea7f5f9da76c213581c13e7bbe1b5735c9925b279eeb4eb/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:574f642832a106456cfc4f32aa82484c504fc32f4be2b510202bcb579de8e6d1", size = 38849839, upload-time = "2026-06-29T11:19:25.783Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5a/b9/2383d30d317bb913cab97dbf2e6e1d5f37f594860d5c5bc176e025cf7d4a/datafusion-54.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:796fd5683927443c5bc61999d00b9007ef9b5ce107725ea8d241df718860985d", size = 41074623, upload-time = "2026-06-29T11:19:29.119Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/35/5c/553fd1107dede0a56727fda7216a7198d41394f2d19697f4fb104cc695ea/datafusion-54.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:64973c63874ec31670dd97b32b18af7b07fad679cb20d58ed154038e3a5c204e", size = 43438801, upload-time = "2026-06-29T11:19:32.799Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1828,19 +1843,19 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "0.7.7"
|
||||
version = "0.8.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace-urllib3-client" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/06/5c/9822af615fc1bd3ee1073994696c739aecde377be32435ec3303aed1bc5d/lance_namespace-0.7.7.tar.gz", hash = "sha256:d00b525f2e26993a6c61668e798bca6c808605ab8a79f29f86a1a1af92d91ae2", size = 10754, upload-time = "2026-05-20T17:32:59.45Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/12/f7ab93b29be3edbf5fc3610714bf2d06088e7f4524bfb38dfd6852458b08/lance_namespace-0.8.6.tar.gz", hash = "sha256:18232e721c8188145f4ec9389cc2dfbeeabf54a619d94885ea1b3375bee9f4af", size = 11529, upload-time = "2026-06-12T17:36:41.651Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/11/43/186acc1156da20c351db196e2b6241b2453b16dc1b4cc8e0a626667ca471/lance_namespace-0.7.7-py3-none-any.whl", hash = "sha256:477a7ca6b5e1f673a2c9ba52f42d6e8e3ff7c27a601392a21eb90fba98d0309b", size = 12581, upload-time = "2026-05-20T17:32:57.389Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/1b/5b1668ee2dc8910965f390640359112a31157092fcf8e000b89c79b58708/lance_namespace-0.8.6-py3-none-any.whl", hash = "sha256:571eae34f9aad70e5b05020416c2860889b9ec82993ccd0eb015e7b39c3ea309", size = 13383, upload-time = "2026-06-12T17:36:43.456Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-urllib3-client"
|
||||
version = "0.7.7"
|
||||
version = "0.8.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
@@ -1848,9 +1863,9 @@ dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/07/95/38ab81ccc1e09beeecd8ddfc61b8bc73831dc5053db1e3f9021f64a4896b/lance_namespace_urllib3_client-0.7.7.tar.gz", hash = "sha256:4d8c066628c17c6a10cf643b51a7f7ae1bfb8a614d9cc54a5af38a4ba2b4b102", size = 202930, upload-time = "2026-05-20T17:32:58.308Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/80/fb224b4a89c1c1638cde949cb6cce6c3aca7759effbfea46a3d9c3960b21/lance_namespace_urllib3_client-0.8.6.tar.gz", hash = "sha256:b6fb1d306e74a7576e5309919020be744527de484a63dbf5eed10f8b368548df", size = 228772, upload-time = "2026-06-12T17:36:42.609Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/35/96/5483e48e40433b1d078183c15a92c99e59a156041b0260e7f18ee34e7c08/lance_namespace_urllib3_client-0.7.7-py3-none-any.whl", hash = "sha256:9221c3e00fd89f0c811953d94b32d2ea527765280460a174f5872dc8a74c0ed6", size = 334767, upload-time = "2026-05-20T17:32:55.883Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/90/1e27de15cd1b16785a1c7312beb0a59e75c8344a815f600f58173a565bd1/lance_namespace_urllib3_client-0.8.6-py3-none-any.whl", hash = "sha256:9d78249c3fb15aa3d15d668f78f04a275af3d08d800a7027492f37996ac4968b", size = 369950, upload-time = "2026-06-12T17:36:40.438Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1931,6 +1946,7 @@ tests = [
|
||||
{ name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" },
|
||||
{ name = "polars" },
|
||||
{ name = "pyarrow" },
|
||||
{ name = "pyarrow-stubs" },
|
||||
{ name = "pylance" },
|
||||
{ name = "pytest" },
|
||||
@@ -1950,7 +1966,7 @@ requires-dist = [
|
||||
{ name = "botocore", marker = "extra == 'embeddings'", specifier = ">=1.31.57" },
|
||||
{ name = "cohere", marker = "extra == 'embeddings'", specifier = ">=4.0" },
|
||||
{ name = "colpali-engine", marker = "extra == 'embeddings'", specifier = ">=0.3.10" },
|
||||
{ name = "datafusion", marker = "extra == 'tests'", specifier = ">=52,<53" },
|
||||
{ name = "datafusion", marker = "extra == 'tests'", specifier = ">=54,<55" },
|
||||
{ name = "deprecation", specifier = ">=2.1.0" },
|
||||
{ name = "duckdb", marker = "extra == 'tests'", specifier = ">=0.9.0" },
|
||||
{ name = "google-genai", marker = "extra == 'embeddings'", specifier = ">=1.0.0" },
|
||||
@@ -1978,10 +1994,11 @@ requires-dist = [
|
||||
{ name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.3.0" },
|
||||
{ name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.5.0" },
|
||||
{ name = "pyarrow", specifier = ">=16" },
|
||||
{ name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" },
|
||||
{ name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" },
|
||||
{ name = "pydantic", specifier = ">=1.10" },
|
||||
{ name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = ">=5.0.0b5" },
|
||||
{ name = "pylance", marker = "extra == 'tests'", specifier = "==9.0.0rc1" },
|
||||
{ name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" },
|
||||
{ name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" },
|
||||
{ name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" },
|
||||
@@ -3837,8 +3854,8 @@ crypto = [
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "7.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
version = "9.0.0rc1"
|
||||
source = { registry = "https://pypi.fury.io/lance-format" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
@@ -3846,12 +3863,12 @@ dependencies = [
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ac/ad/2f64921bf346e7075aef24a72595db44821724a3d89a9a92dd24e79632aa/pylance-7.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:98422021975be76e72b1572f41b8c9abb3bee5bdc9bfa5e9ce731110a65ed4d1", size = 62134146, upload-time = "2026-05-27T21:59:37.459Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/73/1c/c5a01bee0160b55d9a98895cbd33091d038f0a0995b121ab72e629008d02/pylance-7.0.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4bec86ee5b6fbd8bfc493e653f0a1fba0303cfe5492b9b46fc25ab908edc7183", size = 65373684, upload-time = "2026-05-27T22:04:01.584Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/da/1fe8b8f7dbfe734d76af76acc994fc360a0d0c79a4874ef69f5a72a58fe3/pylance-7.0.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:881491432c53184e52f8d1db8d5f872f39a03f36fb104bec77b33d379519d8b5", size = 69458555, upload-time = "2026-05-27T22:16:50.567Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/76/f0/dd505cf3fd0226ab9d94759acd713125af1d3bfacfd80bbd52e3b9f89509/pylance-7.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:18453999e7fff4f76b16d6b7882c9df0628bd142ff95e2461bd7dd5ee3fe0af3", size = 65394430, upload-time = "2026-05-27T22:05:30.923Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/17/ba/2357b81034f28eb00790e258ed140289a6a887a7468ca9df6349fd186b27/pylance-7.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:04a58051d408c60fe76d41a220dcaf8fea8fb6d1aa0ca78a709b60bc3cc8d19a", size = 69473470, upload-time = "2026-05-27T22:17:18.935Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/1f/ec/5c00b6303a67d787f9475141832cbdc513d674ac3dcaeef8a7b169905e65/pylance-7.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:467d4864af047eaab4e1370e2f1e88e2c6f507c079874421116cb41d78bc3629", size = 74792863, upload-time = "2026-05-27T22:19:23.875Z" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_vEHBE/pylance-9.0.0rc1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0b6b02a1808bb3072ee7fe4e36614cae6f86302513e73ec7f55b2234a963b24" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1Jipm4/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:30f0ebf0d88034301819eb964f9236ce555aaa58e7ab89c5975a3e2250bbb405" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_IvKxo/pylance-9.0.0rc1-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:44609ea2615ea6e684b85478d1694af2026458f61cf7895ecc75e238bfd17aa8" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2hidj1/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:182167a8dba9eeabffbffd53bd5b8548613d4d459b7cd7b34a840dd00cbb806f" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_1dFx3r/pylance-9.0.0rc1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8a63b11e814b7eab758bcaf0d6f97eb05ea86203d9fb0af718c462c24c7d6c9c" },
|
||||
{ url = "https://pypi.fury.io/lance-format/-/ver_2a8dSh/pylance-9.0.0rc1-cp310-abi3-win_amd64.whl", hash = "sha256:2ff8b953ae2b0550490c1a7efd210aa91bc223d200ffac28849056cfd7436d97" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.32.0-beta.0"
|
||||
version = "0.32.0-beta.2"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
@@ -44,13 +44,14 @@ lance-io = { workspace = true }
|
||||
lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lindera"] }
|
||||
lance-table = { workspace = true }
|
||||
lance-linalg = { workspace = true }
|
||||
lance-testing = { workspace = true }
|
||||
lance-encoding = { workspace = true }
|
||||
lance-arrow = { workspace = true }
|
||||
lance-namespace = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
# Pin the transitive GooseFS SDK until the 0.1.6 compile break is fixed upstream.
|
||||
goosefs-sdk = { version = "=0.1.5", optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
@@ -95,6 +96,7 @@ semver = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
anyhow = "1"
|
||||
lance-testing = { workspace = true }
|
||||
tempfile = "3.5.0"
|
||||
random_word = { version = "0.4.3", features = ["en"] }
|
||||
tokio = { version = "1.23", features = ["io-util", "macros", "net", "rt-multi-thread", "sync"] }
|
||||
@@ -132,6 +134,7 @@ azure = [
|
||||
]
|
||||
cos = ["lance/tencent", "lance-io/tencent"]
|
||||
goosefs = [
|
||||
"dep:goosefs-sdk",
|
||||
"lance/goosefs",
|
||||
"lance-io/goosefs",
|
||||
"lance-namespace-impls/dir-goosefs",
|
||||
|
||||
@@ -118,8 +118,12 @@ async fn create_empty_table(db: &Connection) -> Result<LanceDbTable> {
|
||||
|
||||
async fn create_index(table: &LanceDbTable) -> Result<()> {
|
||||
// --8<-- [start:create_index]
|
||||
table.create_index(&["vector"], Index::Auto).execute().await
|
||||
table
|
||||
.create_index(&["vector"], Index::Auto)
|
||||
.execute()
|
||||
.await?;
|
||||
// --8<-- [end:create_index]
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn search(table: &LanceDbTable) -> Result<Vec<RecordBatch>> {
|
||||
|
||||
@@ -23,8 +23,10 @@ use crate::connection::create_table::CreateTableBuilder;
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
|
||||
TableNamesRequest,
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, Database,
|
||||
DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo, JobInfo, MaterializedViewInfo,
|
||||
MvRefreshPlan, OpenTableRequest, PlatformJobDescription, ReadConsistency,
|
||||
RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
@@ -488,6 +490,140 @@ impl Connection {
|
||||
)
|
||||
}
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
// Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
// databases return NotSupported for now.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
pub async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
self.internal.create_function(request).await
|
||||
}
|
||||
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
pub async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
self.internal.list_functions().await
|
||||
}
|
||||
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
pub async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_function(name).await
|
||||
}
|
||||
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
pub async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
self.internal.create_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
pub async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
self.internal.refresh_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
pub async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
self.internal.table_lineage(request).await
|
||||
}
|
||||
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH).
|
||||
pub async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
self.internal
|
||||
.explain_refresh_materialized_view(name, full, src_version)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
pub async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
self.internal
|
||||
.alter_materialized_view(name, auto_refresh)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
pub async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_materialized_view(name).await
|
||||
}
|
||||
|
||||
/// List registered materialized view definitions.
|
||||
pub async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
self.internal.list_materialized_views().await
|
||||
}
|
||||
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
self.internal.list_jobs().await
|
||||
}
|
||||
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was flagged for cancellation.
|
||||
pub async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
self.internal.cancel_job(job_id).await
|
||||
}
|
||||
|
||||
/// Describe a platform job (`POST /v1/jobs/describe`): registry-backed
|
||||
/// lifecycle state plus the owner-written status payload. `None` when the
|
||||
/// registry has no such job.
|
||||
pub async fn describe_platform_job(
|
||||
&self,
|
||||
platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
self.internal.describe_platform_job(platform_job_id).await
|
||||
}
|
||||
|
||||
/// Resolve a submission (manifest) job id to its platform job id. `None`
|
||||
/// until the job has registered (dispatch is async).
|
||||
pub async fn resolve_platform_job_id(
|
||||
&self,
|
||||
manifest_job_id: &str,
|
||||
table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
self.internal
|
||||
.resolve_platform_job_id(manifest_job_id, table_hint)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Cancel a platform job. Idempotent on already-terminal jobs.
|
||||
pub async fn cancel_platform_job(&self, platform_job_id: &str) -> Result<()> {
|
||||
self.internal.cancel_platform_job(platform_job_id).await
|
||||
}
|
||||
|
||||
/// Look up a single server-side job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table) enables an O(1) server-side lookup; `None`
|
||||
/// scans the database's active jobs. A `None` result means unknown / not
|
||||
/// active.
|
||||
pub async fn get_job(&self, job_id: &str, table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
self.internal.get_job(job_id, table_hint).await
|
||||
}
|
||||
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables.
|
||||
/// Pass `job_id` to narrow to a single job.
|
||||
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
self.internal.job_history(job_id).await
|
||||
}
|
||||
|
||||
/// Per-row UDF errors (SHOW ERRORS) across the database's tables, optionally
|
||||
/// filtered by `job_id` and/or `table`.
|
||||
pub async fn errors(
|
||||
&self,
|
||||
job_id: Option<&str>,
|
||||
table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
self.internal.errors(job_id, table).await
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
|
||||
@@ -27,7 +27,7 @@ use lance_namespace::models::{
|
||||
};
|
||||
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::error::Result;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::table::{BaseTable, WriteOptions};
|
||||
|
||||
pub mod listing;
|
||||
@@ -200,6 +200,222 @@ pub enum ReadConsistency {
|
||||
Strong,
|
||||
}
|
||||
|
||||
/// A request to register a UDF (CREATE FUNCTION).
|
||||
///
|
||||
/// Functions are first-class database objects, decoupled from any
|
||||
/// column; computed columns and materialized views reference them by
|
||||
/// name. Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateFunctionRequest {
|
||||
/// Function name.
|
||||
pub name: String,
|
||||
/// Implementation language (currently "python").
|
||||
pub language: String,
|
||||
/// SQL return type, e.g. `FLOAT`, `FLOAT[1536]`,
|
||||
/// `STRUCT(a FLOAT, b VARCHAR)`, `TABLE(chunk VARCHAR, idx INT)`.
|
||||
pub return_type: String,
|
||||
/// Function body: source text, or base64 cloudpickle bytes when
|
||||
/// `options["body_format"] = "cloudpickle"`.
|
||||
pub body: String,
|
||||
/// Options: input_columns, pip, num_gpus, batch_size, timeout,
|
||||
/// error_policy, docker_image, body_format, ...
|
||||
pub options: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A request to create a materialized view (CREATE MATERIALIZED VIEW).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// The view's SELECT statement, e.g.
|
||||
/// `SELECT id, embed(body) AS vec FROM articles WHERE id > 1`.
|
||||
/// Bare columns project through; function-call columns compute via
|
||||
/// registered UDFs (a RETURNS TABLE function makes a row-expanding
|
||||
/// chunker view).
|
||||
pub query: String,
|
||||
/// Refresh automatically when the source table changes.
|
||||
pub auto_refresh: bool,
|
||||
/// Register the definition only; skip the initial population.
|
||||
pub with_no_data: bool,
|
||||
/// Optional source column to partition the view's table function on. If the
|
||||
/// column has an IVF vector index the server partitions by its clusters
|
||||
/// (image-dedup style); otherwise it groups by distinct value.
|
||||
pub partition_by: Option<String>,
|
||||
}
|
||||
|
||||
impl CreateMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>, query: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
query: query.into(),
|
||||
auto_refresh: false,
|
||||
with_no_data: false,
|
||||
partition_by: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A request to refresh a materialized view.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RefreshMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// Force a full rebuild (recompute and replace every row) instead of the
|
||||
/// default incremental refresh.
|
||||
pub full: bool,
|
||||
/// Pin the refresh to a source-table version; latest when absent.
|
||||
pub src_version: Option<u64>,
|
||||
/// Initial worker count.
|
||||
pub num_workers: Option<u32>,
|
||||
/// Elastic worker ceiling.
|
||||
pub max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
/// A request for the derived-compute lineage of a table/view (or one of its
|
||||
/// columns). The response is server-defined lineage JSON, returned opaque so
|
||||
/// this client need not model the server's lineage schema.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TableLineageRequest {
|
||||
/// Table or view name.
|
||||
pub name: String,
|
||||
/// Column for column-level lineage; whole table/view when absent.
|
||||
pub column: Option<String>,
|
||||
/// "upstream" | "downstream" | "both" (server default when absent).
|
||||
pub direction: Option<String>,
|
||||
/// Column-hops to walk; transitive when absent.
|
||||
pub depth: Option<u32>,
|
||||
}
|
||||
|
||||
impl RefreshMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
full: false,
|
||||
src_version: None,
|
||||
num_workers: None,
|
||||
max_workers: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A registered materialized view definition, as returned by
|
||||
/// `list_materialized_views`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
/// Source columns projected through.
|
||||
pub projection: Vec<String>,
|
||||
/// `alias=expression` per UDF-computed column.
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// A described platform job (`POST /v1/jobs/describe`): the job registry's
|
||||
/// lifecycle state plus the owner-written status payload.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PlatformJobDescription {
|
||||
/// The platform (registry) job id -- what describe/cancel accept.
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub job_subtype: String,
|
||||
/// "IN_PROGRESS" | "CANCELLED" | "FAILED" | "DONE".
|
||||
pub job_state: String,
|
||||
pub creation_ms: i64,
|
||||
/// The owner-written status payload -- `units_done` / `units_total` /
|
||||
/// `rows_committed` / `error` when present. Records whose owner has not
|
||||
/// written a payload yet carry the raw status-store URI string instead.
|
||||
pub status: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A row from `list_jobs`: one inflight server-side job (index build,
|
||||
/// compaction, column refresh, view refresh, ...).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
/// Lifecycle state: "running", "cancelling", or "stale".
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
/// Whether the job's final commit has completed (output visible).
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `job_history`: one durable, completed/terminal server-side job
|
||||
/// record (SHOW JOB HISTORY), read from a table's `_job_history` store. Unlike
|
||||
/// `JobInfo` (live, inflight jobs) this carries created/updated/completed
|
||||
/// timestamps and the lifecycle event log.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobHistoryInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
/// Newline-joined lifecycle event log, oldest first.
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `errors`: one per-row UDF failure recorded by `error_policy=skip`
|
||||
/// (SHOW ERRORS).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobErrorInfo {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a `REFRESH MATERIALIZED VIEW` would execute, as returned by
|
||||
/// `explain_refresh_materialized_view` (EXPLAIN REFRESH). No work is run.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
/// Whether a refresh would do anything (rebuild or non-empty units).
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
/// Source changed non-append-only since the last refresh -> rebuild.
|
||||
pub rebuild: bool,
|
||||
/// Number of row-range work units the refresh would process.
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
fn not_supported<T>(what: &str) -> Result<T> {
|
||||
Err(Error::NotSupported {
|
||||
message: format!("{} is not supported by this database", what),
|
||||
})
|
||||
}
|
||||
|
||||
/// The `Database` trait defines the interface for database implementations.
|
||||
///
|
||||
/// A database is responsible for managing tables and their metadata.
|
||||
@@ -245,6 +461,126 @@ pub trait Database:
|
||||
///
|
||||
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
||||
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
//
|
||||
// Server-backed features (LanceDB Enterprise / Cloud). The defaults
|
||||
// return NotSupported; the remote database overrides them. Local
|
||||
// single-node implementations are planned.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
async fn create_function(&self, _request: CreateFunctionRequest) -> Result<()> {
|
||||
not_supported("create_function")
|
||||
}
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
not_supported("list_functions")
|
||||
}
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
async fn drop_function(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_function")
|
||||
}
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
_request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
not_supported("create_materialized_view")
|
||||
}
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
_request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
not_supported("refresh_materialized_view")
|
||||
}
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
async fn table_lineage(&self, _request: TableLineageRequest) -> Result<String> {
|
||||
not_supported("table_lineage")
|
||||
}
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH). `full` plans a full rebuild (incremental
|
||||
/// planning requires stable row IDs on the source).
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
_name: &str,
|
||||
_full: bool,
|
||||
_src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
not_supported("explain_refresh_materialized_view")
|
||||
}
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
async fn alter_materialized_view(&self, _name: &str, _auto_refresh: bool) -> Result<()> {
|
||||
not_supported("alter_materialized_view")
|
||||
}
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
async fn drop_materialized_view(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_materialized_view")
|
||||
}
|
||||
/// List registered materialized view definitions.
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
not_supported("list_materialized_views")
|
||||
}
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
not_supported("list_jobs")
|
||||
}
|
||||
|
||||
/// Describe a platform job (`POST /v1/jobs/describe`): registry-backed
|
||||
/// lifecycle state plus the owner-written status payload. `None` when the
|
||||
/// registry has no such job.
|
||||
async fn describe_platform_job(
|
||||
&self,
|
||||
_platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
not_supported("describe_platform_job")
|
||||
}
|
||||
|
||||
/// Resolve a submission (manifest) job id to its platform job id via the
|
||||
/// registry's manifest-id filter (`POST /v1/jobs/list`). `None` until the
|
||||
/// job has registered (dispatch is async).
|
||||
async fn resolve_platform_job_id(
|
||||
&self,
|
||||
_manifest_job_id: &str,
|
||||
_table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
not_supported("resolve_platform_job_id")
|
||||
}
|
||||
|
||||
/// Cancel a platform job (`POST /v1/jobs/cancel`). Idempotent: cancelling
|
||||
/// an already-terminal job is a no-op success.
|
||||
async fn cancel_platform_job(&self, _platform_job_id: &str) -> Result<()> {
|
||||
not_supported("cancel_platform_job")
|
||||
}
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was found and flagged for cancellation,
|
||||
/// false if none was inflight (best-effort, like SQL `CANCEL JOB`).
|
||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
||||
not_supported("cancel_job")
|
||||
}
|
||||
/// Point-access for a single job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table, which `wait()` callers know) enables an
|
||||
/// O(1) server-side lookup. `None` if the job is unknown or not active.
|
||||
async fn get_job(&self, _job_id: &str, _table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
not_supported("get_job")
|
||||
}
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables,
|
||||
/// optionally narrowed to a single `job_id`.
|
||||
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
not_supported("job_history")
|
||||
}
|
||||
/// Per-row UDF errors (SHOW ERRORS) recorded by `error_policy=skip` across
|
||||
/// the database's tables, optionally filtered by `job_id` and/or `table`.
|
||||
async fn errors(
|
||||
&self,
|
||||
_job_id: Option<&str>,
|
||||
_table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
not_supported("errors")
|
||||
}
|
||||
|
||||
/// Open a table in the database
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Rename a table in the database
|
||||
|
||||
@@ -8,13 +8,13 @@ use arrow_array::{RecordBatch, UInt64Array};
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use lance::io::ObjectStore;
|
||||
use lance_core::{cache::LanceCache, utils::futures::FinallyStreamExt};
|
||||
use lance_encoding::decoder::DecoderPlugins;
|
||||
use lance_encoding::decoder::{DecoderPlugins, FilterExpression};
|
||||
use lance_file::{
|
||||
reader::{FileReader, FileReaderOptions},
|
||||
writer::{FileWriter, FileWriterOptions},
|
||||
};
|
||||
use lance_index::scalar::IndexReader;
|
||||
use lance_io::{
|
||||
ReadBatchParams,
|
||||
scheduler::{ScanScheduler, SchedulerConfig},
|
||||
utils::CachedFileSize,
|
||||
};
|
||||
@@ -216,6 +216,7 @@ impl Shuffler {
|
||||
let scan_scheduler = ScanScheduler::new(Arc::new(object_store), scheduler_config);
|
||||
let job_id = self.id.clone();
|
||||
let rng = Arc::new(Mutex::new(rng));
|
||||
let read_schema = arrow_schema.clone();
|
||||
|
||||
// Second pass, read each file as a single batch and shuffle
|
||||
let stream = futures::stream::iter(0..num_files)
|
||||
@@ -224,6 +225,7 @@ impl Shuffler {
|
||||
let rng = rng.clone();
|
||||
let tmp_dir = tmp_dir.clone();
|
||||
let job_id = job_id.clone();
|
||||
let read_schema = read_schema.clone();
|
||||
async move {
|
||||
let path = tmp_dir.join(format!("shuffle_{}_{file_index}.lance", job_id));
|
||||
let path = object_store::path::Path::from_absolute_path(path).unwrap();
|
||||
@@ -239,7 +241,19 @@ impl Shuffler {
|
||||
)
|
||||
.await?;
|
||||
// Need to read the entire file in a single batch for in-memory shuffling
|
||||
let batch = reader.read_record_batch(0, reader.num_rows()).await?;
|
||||
let batches = reader
|
||||
.read_stream(
|
||||
ReadBatchParams::RangeFull,
|
||||
reader.num_rows() as u32,
|
||||
1,
|
||||
FilterExpression::no_filter(),
|
||||
)
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
// An empty file yields no batches; fall back to an empty batch
|
||||
// with the expected schema so shuffling handles it gracefully.
|
||||
let batch = concat_batches(&read_schema, &batches)?;
|
||||
let mut rng = rng.lock().unwrap_or_else(|e| e.into_inner());
|
||||
Self::shuffle_batch(&batch, &mut rng, clump_size)
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ use std::sync::{
|
||||
|
||||
use arrow_array::{Array, BooleanArray, RecordBatch, UInt64Array};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use datafusion_common::hash_utils::create_hashes;
|
||||
use datafusion_common::hash_utils::{RandomState, create_hashes};
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use lance_arrow::SchemaExt;
|
||||
|
||||
@@ -234,7 +234,7 @@ impl Splitter {
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let mut hashes = vec![0; batch.num_rows()];
|
||||
let random_state = ahash::RandomState::with_seeds(0, 0, 0, 0);
|
||||
let random_state = RandomState::with_seed(0);
|
||||
create_hashes(&arrays, &random_state, &mut hashes).unwrap();
|
||||
// As an example, let's assume the weights are 1, 2. Our total weight is 3.
|
||||
//
|
||||
@@ -761,8 +761,8 @@ mod tests {
|
||||
verify_splitter(splitter, test_data(), 50, &[11, 8, 9], false).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hash_split() {
|
||||
async fn collect_hash_split() -> RecordBatch {
|
||||
let total_rows = 50;
|
||||
let data = lance_datagen::gen_batch()
|
||||
.with_seed(Seed::from(42))
|
||||
.col(
|
||||
@@ -783,7 +783,7 @@ mod tests {
|
||||
);
|
||||
|
||||
let split_batches = splitter
|
||||
.apply(data, 10)
|
||||
.apply(data, total_rows)
|
||||
.await
|
||||
.unwrap()
|
||||
.try_collect::<Vec<_>>()
|
||||
@@ -791,20 +791,35 @@ mod tests {
|
||||
.unwrap();
|
||||
|
||||
let schema = split_batches[0].schema();
|
||||
let split_batch = concat_batches(&schema, &split_batches).unwrap();
|
||||
concat_batches(&schema, &split_batches).unwrap()
|
||||
}
|
||||
|
||||
// These assertions are all based on fixed seed in data generation but they match
|
||||
// up roughly to what we expect (25% discarded, 25% in split 0, 50% in split 1)
|
||||
#[tokio::test]
|
||||
async fn test_hash_split() {
|
||||
let total_rows = 50;
|
||||
let split_batch = collect_hash_split().await;
|
||||
let split_batch_again = collect_hash_split().await;
|
||||
|
||||
// 14 rows (28%) are discarded because discard_weight is 1
|
||||
assert_eq!(split_batch.num_rows(), 36);
|
||||
assert_eq!(split_batch.num_rows(), split_batch_again.num_rows());
|
||||
assert_eq!(split_batch.num_columns(), split_batch_again.num_columns());
|
||||
for (left, right) in split_batch
|
||||
.columns()
|
||||
.iter()
|
||||
.zip(split_batch_again.columns().iter())
|
||||
{
|
||||
assert_eq!(left, right);
|
||||
}
|
||||
|
||||
assert!(split_batch.num_rows() > 0);
|
||||
assert!(split_batch.num_rows() < total_rows);
|
||||
assert_eq!(split_batch.num_columns(), 2);
|
||||
|
||||
let split_ids = split_batch.column(1).as_primitive::<UInt64Type>().values();
|
||||
let num_in_split_0 = split_ids.iter().filter(|v| **v == 0).count();
|
||||
let num_in_split_1 = split_ids.iter().filter(|v| **v == 1).count();
|
||||
|
||||
assert_eq!(num_in_split_0, 11); // 22%
|
||||
assert_eq!(num_in_split_1, 25); // 50%
|
||||
assert_eq!(num_in_split_0 + num_in_split_1, split_batch.num_rows());
|
||||
assert!(num_in_split_0 > 0);
|
||||
assert!(num_in_split_1 > num_in_split_0);
|
||||
}
|
||||
}
|
||||
|
||||
+125
-16
@@ -198,28 +198,36 @@ fn compute_embedding_arrays(
|
||||
batch: &RecordBatch,
|
||||
embeddings: &[(EmbeddingDefinition, Arc<dyn EmbeddingFunction>)],
|
||||
) -> Result<Vec<Arc<dyn Array>>> {
|
||||
if embeddings.len() == 1 {
|
||||
let (fld, func) = &embeddings[0];
|
||||
let src_column =
|
||||
batch
|
||||
.column_by_name(&fld.source_column)
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!("Source column '{}' not found", fld.source_column),
|
||||
})?;
|
||||
let input_columns = embeddings
|
||||
.iter()
|
||||
.map(|(fld, func)| {
|
||||
let src_column =
|
||||
batch
|
||||
.column_by_name(&fld.source_column)
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!("Source column '{}' not found", fld.source_column),
|
||||
})?;
|
||||
Ok((src_column.clone(), func))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
|
||||
if batch.num_rows() == 0 {
|
||||
return input_columns
|
||||
.iter()
|
||||
.map(|(_, func)| Ok(arrow_array::new_empty_array(func.dest_type()?.as_ref())))
|
||||
.collect();
|
||||
}
|
||||
|
||||
if input_columns.len() == 1 {
|
||||
let (src_column, func) = &input_columns[0];
|
||||
return Ok(vec![func.compute_source_embeddings(src_column.clone())?]);
|
||||
}
|
||||
|
||||
// Parallel path: multiple embeddings
|
||||
std::thread::scope(|s| {
|
||||
let handles: Vec<_> = embeddings
|
||||
let handles: Vec<_> = input_columns
|
||||
.iter()
|
||||
.map(|(fld, func)| {
|
||||
let src_column = batch.column_by_name(&fld.source_column).ok_or_else(|| {
|
||||
Error::InvalidInput {
|
||||
message: format!("Source column '{}' not found", fld.source_column),
|
||||
}
|
||||
})?;
|
||||
|
||||
.map(|(src_column, func)| {
|
||||
let handle = s.spawn(move || func.compute_source_embeddings(src_column.clone()));
|
||||
|
||||
Ok(handle)
|
||||
@@ -392,3 +400,104 @@ impl<R: RecordBatchReader> RecordBatchReader for WithEmbeddings<R> {
|
||||
.into_rich_schema()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
};
|
||||
|
||||
use arrow_array::{Array, ArrayRef, FixedSizeListArray, RecordBatch, StringArray};
|
||||
use arrow_schema::DataType;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FailingEmbedding {
|
||||
calls: AtomicUsize,
|
||||
}
|
||||
|
||||
impl EmbeddingFunction for FailingEmbedding {
|
||||
fn name(&self) -> &str {
|
||||
"failing"
|
||||
}
|
||||
|
||||
fn source_type(&self) -> Result<Cow<'_, DataType>> {
|
||||
Ok(Cow::Owned(DataType::Utf8))
|
||||
}
|
||||
|
||||
fn dest_type(&self) -> Result<Cow<'_, DataType>> {
|
||||
Ok(Cow::Owned(DataType::new_fixed_size_list(
|
||||
DataType::Float32,
|
||||
3,
|
||||
false,
|
||||
)))
|
||||
}
|
||||
|
||||
fn compute_source_embeddings(&self, _source: Arc<dyn Array>) -> Result<Arc<dyn Array>> {
|
||||
self.calls.fetch_add(1, Ordering::SeqCst);
|
||||
Err(Error::Runtime {
|
||||
message: "embedding function must not receive an empty batch".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn compute_query_embeddings(&self, _input: Arc<dyn Array>) -> Result<Arc<dyn Array>> {
|
||||
unreachable!("query embeddings are not exercised by this test")
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_batch_skips_embedding_functions() {
|
||||
let embedding_function = Arc::new(FailingEmbedding {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let source: ArrayRef = Arc::new(StringArray::from(Vec::<&str>::new()));
|
||||
let batch = RecordBatch::try_from_iter([("text", source)]).unwrap();
|
||||
let embeddings = vec![(
|
||||
EmbeddingDefinition::new("text", "failing", Some("text_embedding")),
|
||||
embedding_function.clone() as Arc<dyn EmbeddingFunction>,
|
||||
)];
|
||||
|
||||
let result = compute_embeddings_for_batch(batch, &embeddings).unwrap();
|
||||
|
||||
assert_eq!(embedding_function.calls.load(Ordering::SeqCst), 0);
|
||||
assert_eq!(result.num_rows(), 0);
|
||||
|
||||
let embedding = result.column_by_name("text_embedding").unwrap();
|
||||
assert_eq!(
|
||||
embedding.data_type(),
|
||||
&DataType::new_fixed_size_list(DataType::Float32, 3, false)
|
||||
);
|
||||
assert_eq!(embedding.null_count(), 0);
|
||||
|
||||
let embedding = embedding
|
||||
.as_any()
|
||||
.downcast_ref::<FixedSizeListArray>()
|
||||
.unwrap();
|
||||
assert_eq!(embedding.len(), 0);
|
||||
assert_eq!(embedding.value_length(), 3);
|
||||
assert_eq!(embedding.values().len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_batch_still_validates_source_column() {
|
||||
let embedding_function = Arc::new(FailingEmbedding {
|
||||
calls: AtomicUsize::new(0),
|
||||
});
|
||||
let source: ArrayRef = Arc::new(StringArray::from(Vec::<&str>::new()));
|
||||
let batch = RecordBatch::try_from_iter([("text", source)]).unwrap();
|
||||
let embeddings = vec![(
|
||||
EmbeddingDefinition::new("missing_column", "failing", Some("text_embedding")),
|
||||
embedding_function.clone() as Arc<dyn EmbeddingFunction>,
|
||||
)];
|
||||
|
||||
let result = compute_embeddings_for_batch(batch, &embeddings);
|
||||
assert!(result.is_err());
|
||||
assert!(
|
||||
matches!(result.unwrap_err(), Error::InvalidInput { .. }),
|
||||
"expected InvalidInput error when source column is missing"
|
||||
);
|
||||
assert_eq!(embedding_function.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -283,7 +283,10 @@ impl IndexBuilder {
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn execute(self) -> Result<()> {
|
||||
/// Returns the server-minted job id when the index build was deferred to
|
||||
/// a background job (remote tables only); `None` when the build completed
|
||||
/// synchronously within this call.
|
||||
pub async fn execute(self) -> Result<Option<String>> {
|
||||
self.parent.clone().create_index(self).await
|
||||
}
|
||||
}
|
||||
@@ -317,6 +320,8 @@ pub enum IndexType {
|
||||
// FTS
|
||||
#[serde(alias = "INVERTED", alias = "Inverted")]
|
||||
FTS,
|
||||
/// Catch-all for index types not recognized by this version of LanceDB.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for IndexType {
|
||||
@@ -334,6 +339,7 @@ impl std::fmt::Display for IndexType {
|
||||
Self::LabelList => write!(f, "LABEL_LIST"),
|
||||
Self::Fm => write!(f, "FM"),
|
||||
Self::FTS => write!(f, "FTS"),
|
||||
Self::Unknown => write!(f, "UNKNOWN"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -355,9 +361,7 @@ impl std::str::FromStr for IndexType {
|
||||
"IVF_HNSW_PQ" => Ok(Self::IvfHnswPq),
|
||||
"IVF_HNSW_SQ" => Ok(Self::IvfHnswSq),
|
||||
"IVF_HNSW_FLAT" => Ok(Self::IvfHnswFlat),
|
||||
_ => Err(Error::InvalidInput {
|
||||
message: format!("the input value {} is not a valid IndexType", value),
|
||||
}),
|
||||
_ => Ok(Self::Unknown),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -425,20 +429,15 @@ pub struct IndexConfig {
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct IndexMetadata {
|
||||
pub metric_type: Option<DistanceType>,
|
||||
// Sometimes the index type is provided at this level.
|
||||
pub index_type: Option<IndexType>,
|
||||
}
|
||||
|
||||
// This struct is used to deserialize the JSON data returned from the Lance API
|
||||
// Dataset::index_statistics().
|
||||
// Deserializes the JSON returned by Dataset::index_statistics().
|
||||
#[skip_serializing_none]
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct IndexStatisticsImpl {
|
||||
pub num_indexed_rows: usize,
|
||||
pub num_unindexed_rows: usize,
|
||||
pub indices: Vec<IndexMetadata>,
|
||||
// Sometimes, the index type is provided at this level.
|
||||
pub index_type: Option<IndexType>,
|
||||
pub num_indices: Option<u32>,
|
||||
}
|
||||
|
||||
|
||||
@@ -207,7 +207,15 @@ use lance_linalg::distance::DistanceType as LanceDistanceType;
|
||||
/// a built-in pull-based adapter.
|
||||
#[cfg(feature = "metrics")]
|
||||
pub use metrics;
|
||||
pub use table::Table;
|
||||
pub use table::{FtsToken, Table};
|
||||
|
||||
/// Tokenize a full-text search query using an explicit FTS tokenizer configuration.
|
||||
///
|
||||
/// This does not require a table or FTS index. The tokenizer options are the
|
||||
/// same [`index::scalar::FtsIndexBuilder`] values used when creating an FTS index.
|
||||
pub fn tokenize(query: &str, params: &index::scalar::FtsIndexBuilder) -> Result<Vec<FtsToken>> {
|
||||
table::tokenize(query, params)
|
||||
}
|
||||
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
#[non_exhaustive]
|
||||
|
||||
@@ -614,6 +614,12 @@ pub struct QueryExecutionOptions {
|
||||
pub max_batch_length: u32,
|
||||
/// Max duration to wait for the query to execute before timing out.
|
||||
pub timeout: Option<Duration>,
|
||||
/// How distributed worker metrics should be displayed by
|
||||
/// [`ExecutableQuery::analyze_plan`].
|
||||
///
|
||||
/// This only affects remote distributed query plans. Local query execution
|
||||
/// ignores this option.
|
||||
pub analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics,
|
||||
}
|
||||
|
||||
impl Default for QueryExecutionOptions {
|
||||
@@ -621,6 +627,7 @@ impl Default for QueryExecutionOptions {
|
||||
Self {
|
||||
max_batch_length: 1024,
|
||||
timeout: None,
|
||||
analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::Aggregate,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -633,6 +640,29 @@ impl QueryExecutionOptions {
|
||||
}
|
||||
}
|
||||
|
||||
/// How distributed worker metrics are displayed in analyzed query plans.
|
||||
#[non_exhaustive]
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub enum AnalyzePlanDistributedMetrics {
|
||||
/// Preserve the legacy output: aggregate worker metrics into one synthetic tree.
|
||||
#[default]
|
||||
Aggregate,
|
||||
/// Render one raw worker-side tree per distributed worker.
|
||||
PerWorker,
|
||||
/// Render the aggregate tree followed by the raw per-worker trees.
|
||||
Full,
|
||||
}
|
||||
|
||||
impl AnalyzePlanDistributedMetrics {
|
||||
pub(crate) fn as_query_param(self) -> &'static str {
|
||||
match self {
|
||||
Self::Aggregate => "aggregate",
|
||||
Self::PerWorker => "per_worker",
|
||||
Self::Full => "full",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A trait for a query object that can be executed to get results
|
||||
///
|
||||
/// There are various kinds of queries but they all return results
|
||||
|
||||
@@ -19,8 +19,10 @@ use lance_namespace::models::{
|
||||
|
||||
use crate::Error;
|
||||
use crate::database::{
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||
OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode,
|
||||
CreateTableRequest, Database, DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo,
|
||||
JobInfo, MaterializedViewInfo, MvRefreshPlan, OpenTableRequest, PlatformJobDescription,
|
||||
ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest,
|
||||
};
|
||||
use crate::error::Result;
|
||||
use crate::remote::util::stream_as_body;
|
||||
@@ -33,6 +35,210 @@ use super::client::{
|
||||
use super::table::RemoteTable;
|
||||
use super::util::parse_server_version;
|
||||
|
||||
// Wire types for the derived-compute routes (functions, materialized
|
||||
// views, jobs). Field shapes mirror the server's REST contract.
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateFunctionRequest {
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteFunctionEntry {
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListFunctionsResponse {
|
||||
functions: Vec<RemoteFunctionEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateMaterializedViewRequest {
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
partition_by: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteCreateMaterializedViewResponse {
|
||||
#[serde(default)]
|
||||
job_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteRefreshMaterializedViewRequest {
|
||||
#[serde(skip_serializing_if = "std::ops::Not::not")]
|
||||
full: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
num_workers: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteRefreshMaterializedViewResponse {
|
||||
job_id: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteExplainRefreshRequest {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
full: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteExplainRefreshResponse {
|
||||
table_name: String,
|
||||
has_work: bool,
|
||||
source_version: u64,
|
||||
last_refreshed_version: Option<u64>,
|
||||
full_refresh: bool,
|
||||
rebuild: bool,
|
||||
units_total: u64,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteAlterMaterializedViewRequest {
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteMaterializedViewEntry {
|
||||
name: String,
|
||||
source_table: String,
|
||||
#[serde(default)]
|
||||
projection: Vec<String>,
|
||||
#[serde(default)]
|
||||
udf_columns: Vec<String>,
|
||||
#[serde(default)]
|
||||
filter: Option<String>,
|
||||
#[serde(default)]
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListMaterializedViewsResponse {
|
||||
views: Vec<RemoteMaterializedViewEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteDescribePlatformJobResponse {
|
||||
job_id: String,
|
||||
job_type: String,
|
||||
#[serde(default)]
|
||||
job_subtype: String,
|
||||
job_state: String,
|
||||
#[serde(default)]
|
||||
creation_ms: i64,
|
||||
#[serde(default)]
|
||||
status: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListPlatformJobsResponse {
|
||||
#[serde(default)]
|
||||
jobs: Vec<RemotePlatformJobRow>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemotePlatformJobRow {
|
||||
job_id: String,
|
||||
#[serde(default)]
|
||||
table: String,
|
||||
#[serde(default)]
|
||||
job_subtype: String,
|
||||
#[serde(default)]
|
||||
state: String,
|
||||
#[serde(default)]
|
||||
created_at_millis: i64,
|
||||
#[serde(default)]
|
||||
status: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Platform list-row state -> the client's job vocabulary.
|
||||
fn platform_state_to_client(state: &str) -> String {
|
||||
match state {
|
||||
"in_progress" => "running",
|
||||
"done" => "finished",
|
||||
other => other,
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Describe job_state -> the client's job vocabulary.
|
||||
fn describe_state_to_client(state: &str) -> String {
|
||||
match state {
|
||||
"IN_PROGRESS" => "running",
|
||||
"DONE" => "finished",
|
||||
"FAILED" => "failed",
|
||||
"CANCELLED" => "cancelled",
|
||||
other => other,
|
||||
}
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn payload_i64(status: &serde_json::Value, key: &str) -> Option<i64> {
|
||||
status.get(key).and_then(serde_json::Value::as_i64)
|
||||
}
|
||||
|
||||
fn payload_error(status: &serde_json::Value) -> Option<String> {
|
||||
status
|
||||
.get("error")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_string)
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorEntry {
|
||||
job_id: String,
|
||||
table: String,
|
||||
column: String,
|
||||
error_type: String,
|
||||
error_message: String,
|
||||
#[serde(default)]
|
||||
fragment_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
source_row_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
table_version: Option<i64>,
|
||||
#[serde(default)]
|
||||
age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorsResponse {
|
||||
errors: Vec<RemoteErrorEntry>,
|
||||
}
|
||||
|
||||
impl From<RemoteErrorEntry> for JobErrorInfo {
|
||||
fn from(e: RemoteErrorEntry) -> Self {
|
||||
JobErrorInfo {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Request structure for the remote clone table API
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCloneTableRequest {
|
||||
@@ -641,6 +847,426 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
Ok(table)
|
||||
}
|
||||
|
||||
async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
let body = RemoteCreateFunctionRequest {
|
||||
language: request.language,
|
||||
return_type: request.return_type,
|
||||
body: request.body,
|
||||
options: request.options,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/function/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
let req = self.client.get("/v1/function/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListFunctionsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
let req = self.client.post(&format!("/v1/function/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
let body = RemoteCreateMaterializedViewRequest {
|
||||
query: request.query,
|
||||
auto_refresh: request.auto_refresh,
|
||||
with_no_data: request.with_no_data,
|
||||
partition_by: request.partition_by,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteCreateMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
let body = RemoteRefreshMaterializedViewRequest {
|
||||
full: request.full,
|
||||
src_version: request.src_version,
|
||||
num_workers: request.num_workers,
|
||||
max_workers: request.max_workers,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/refresh", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteRefreshMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
let mut req = self
|
||||
.client
|
||||
.get(&format!("/v1/table/{}/lineage", request.name));
|
||||
if let Some(column) = &request.column {
|
||||
req = req.query(&[("column", column)]);
|
||||
}
|
||||
if let Some(direction) = &request.direction {
|
||||
req = req.query(&[("direction", direction)]);
|
||||
}
|
||||
if let Some(depth) = request.depth {
|
||||
req = req.query(&[("depth", depth.to_string())]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
// Server-defined lineage JSON, returned opaque (the client does not
|
||||
// model the lineage schema; the Python layer deserializes it).
|
||||
rsp.text().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
let body = RemoteExplainRefreshRequest {
|
||||
full: Some(full),
|
||||
src_version,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/explain_refresh", name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteExplainRefreshResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: body.table_name,
|
||||
has_work: body.has_work,
|
||||
source_version: body.source_version,
|
||||
last_refreshed_version: body.last_refreshed_version,
|
||||
full_refresh: body.full_refresh,
|
||||
rebuild: body.rebuild,
|
||||
units_total: body.units_total,
|
||||
})
|
||||
}
|
||||
|
||||
async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/alter", name))
|
||||
.json(&RemoteAlterMaterializedViewRequest { auto_refresh });
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
let req = self.client.get("/v1/materialized_view/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListMaterializedViewsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/list")
|
||||
.json(&serde_json::json!({ "include_status": true }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListPlatformJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0);
|
||||
Ok(body
|
||||
.jobs
|
||||
.into_iter()
|
||||
.map(|row| JobInfo {
|
||||
table: row.table,
|
||||
job_id: row.job_id,
|
||||
// The platform job_type is always "indexer"; the subtype
|
||||
// (udf / mv_refresh / compaction / ...) is the useful label.
|
||||
job_type: row.job_subtype,
|
||||
state: platform_state_to_client(&row.state),
|
||||
column: None,
|
||||
age_seconds: (row.created_at_millis > 0)
|
||||
.then(|| (now_ms - row.created_at_millis) / 1000),
|
||||
command: None,
|
||||
units_done: payload_i64(&row.status, "units_done"),
|
||||
units_total: payload_i64(&row.status, "units_total"),
|
||||
committed: row.state == "done",
|
||||
rows_skipped: payload_i64(&row.status, "rows_skipped").unwrap_or(0) as u64,
|
||||
error: payload_error(&row.status),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn get_job(&self, job_id: &str, table: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
// A point snapshot from the platform API: resolve the submission id,
|
||||
// then describe. The snapshot keeps the caller's id.
|
||||
let Some(platform_id) = self.resolve_platform_job_id(job_id, table).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(described) = self.describe_platform_job(&platform_id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(JobInfo {
|
||||
table: table.unwrap_or_default().to_string(),
|
||||
job_id: job_id.to_string(),
|
||||
job_type: described.job_subtype,
|
||||
state: describe_state_to_client(&described.job_state),
|
||||
column: None,
|
||||
age_seconds: None,
|
||||
command: None,
|
||||
units_done: payload_i64(&described.status, "units_done"),
|
||||
units_total: payload_i64(&described.status, "units_total"),
|
||||
committed: described.job_state == "DONE",
|
||||
rows_skipped: payload_i64(&described.status, "rows_skipped").unwrap_or(0) as u64,
|
||||
error: payload_error(&described.status),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn describe_platform_job(
|
||||
&self,
|
||||
platform_job_id: &str,
|
||||
) -> Result<Option<PlatformJobDescription>> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/describe")
|
||||
.json(&serde_json::json!({ "job_id": platform_job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
if rsp.status().as_u16() == 404 {
|
||||
return Ok(None);
|
||||
}
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteDescribePlatformJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(Some(PlatformJobDescription {
|
||||
job_id: body.job_id,
|
||||
job_type: body.job_type,
|
||||
job_subtype: body.job_subtype,
|
||||
job_state: body.job_state,
|
||||
creation_ms: body.creation_ms,
|
||||
status: body.status,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn resolve_platform_job_id(
|
||||
&self,
|
||||
manifest_job_id: &str,
|
||||
table_hint: Option<&str>,
|
||||
) -> Result<Option<String>> {
|
||||
let req = self.client.post("/v1/jobs/list").json(&serde_json::json!({
|
||||
"manifest_job_id": manifest_job_id,
|
||||
"table_name": table_hint,
|
||||
"job_type": "indexer",
|
||||
}));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListPlatformJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.jobs.into_iter().next().map(|row| row.job_id))
|
||||
}
|
||||
|
||||
async fn cancel_platform_job(&self, platform_job_id: &str) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/cancel")
|
||||
.json(&serde_json::json!({ "job_id": platform_job_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
// Resolve the submission id and cancel through the platform API.
|
||||
// False when no matching job has registered (the legacy best-effort
|
||||
// contract).
|
||||
let Some(platform_id) = self.resolve_platform_job_id(job_id, None).await? else {
|
||||
return Ok(false);
|
||||
};
|
||||
self.cancel_platform_job(&platform_id).await?;
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
// One job: describe (identity) plus query_events (timeline). No id:
|
||||
// a registry listing, timeline-free -- pass an id for the event log.
|
||||
let Some(caller_id) = job_id else {
|
||||
let rows = self.list_jobs().await?;
|
||||
let now_ms = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as i64)
|
||||
.unwrap_or(0);
|
||||
return Ok(rows
|
||||
.into_iter()
|
||||
.map(|j| JobHistoryInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
created_ms: j.age_seconds.map(|a| now_ms - a * 1000).unwrap_or_default(),
|
||||
updated_ms: 0,
|
||||
completed_ms: None,
|
||||
rows_processed: None,
|
||||
rows_skipped: j.rows_skipped.try_into().ok(),
|
||||
error: j.error,
|
||||
events: None,
|
||||
})
|
||||
.collect());
|
||||
};
|
||||
// Accept either a platform id or a submission id.
|
||||
let platform_id = match self.describe_platform_job(caller_id).await? {
|
||||
Some(_) => caller_id.to_string(),
|
||||
None => match self.resolve_platform_job_id(caller_id, None).await? {
|
||||
Some(id) => id,
|
||||
None => return Ok(Vec::new()),
|
||||
},
|
||||
};
|
||||
let Some(described) = self.describe_platform_job(&platform_id).await? else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post("/v1/jobs/query_events")
|
||||
.json(&serde_json::json!({ "job_id": platform_id }));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body = rsp.bytes().await.err_to_http(request_id.clone())?;
|
||||
let reader = arrow_ipc::reader::StreamReader::try_new(std::io::Cursor::new(body), None)
|
||||
.map_err(|e| Error::Http {
|
||||
source: format!("failed to read job-events IPC stream: {e}").into(),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
})?;
|
||||
|
||||
let mut created_ms = i64::MAX;
|
||||
let mut updated_ms = 0i64;
|
||||
let mut completed_ms = None;
|
||||
let mut last_error = None;
|
||||
let mut events = Vec::new();
|
||||
for batch in reader {
|
||||
let batch = batch.map_err(|e| Error::Http {
|
||||
source: format!("failed to decode job-events batch: {e}").into(),
|
||||
request_id: request_id.clone(),
|
||||
status_code: None,
|
||||
})?;
|
||||
let states = batch
|
||||
.column_by_name("state")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
|
||||
let times = batch
|
||||
.column_by_name("updated_at_millis")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::Int64Array>());
|
||||
let payloads = batch
|
||||
.column_by_name("payload")
|
||||
.and_then(|c| c.as_any().downcast_ref::<arrow_array::StringArray>());
|
||||
let (Some(states), Some(times)) = (states, times) else {
|
||||
continue;
|
||||
};
|
||||
for i in 0..batch.num_rows() {
|
||||
let state = states.value(i);
|
||||
let ts = times.value(i);
|
||||
created_ms = created_ms.min(ts);
|
||||
updated_ms = updated_ms.max(ts);
|
||||
if matches!(state, "succeeded" | "failed" | "timed_out" | "canceled") {
|
||||
completed_ms = Some(ts);
|
||||
}
|
||||
if let Some(payloads) = payloads {
|
||||
if !arrow_array::Array::is_null(payloads, i) {
|
||||
if let Ok(payload) =
|
||||
serde_json::from_str::<serde_json::Value>(payloads.value(i))
|
||||
{
|
||||
if let Some(e) = payload_error(&payload) {
|
||||
last_error = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
events.push(format!("{state} {ts}"));
|
||||
}
|
||||
}
|
||||
Ok(vec![JobHistoryInfo {
|
||||
table: String::new(),
|
||||
job_id: caller_id.to_string(),
|
||||
job_type: described.job_subtype,
|
||||
state: describe_state_to_client(&described.job_state),
|
||||
column: None,
|
||||
created_ms: if created_ms == i64::MAX {
|
||||
described.creation_ms
|
||||
} else {
|
||||
created_ms
|
||||
},
|
||||
updated_ms,
|
||||
completed_ms,
|
||||
rows_processed: payload_i64(&described.status, "rows_committed"),
|
||||
rows_skipped: payload_i64(&described.status, "rows_skipped"),
|
||||
error: last_error.or_else(|| payload_error(&described.status)),
|
||||
events: (!events.is_empty()).then(|| events.join("\n")),
|
||||
}])
|
||||
}
|
||||
|
||||
async fn errors(&self, job_id: Option<&str>, table: Option<&str>) -> Result<Vec<JobErrorInfo>> {
|
||||
let mut req = self.client.get("/v1/errors");
|
||||
if let Some(j) = job_id {
|
||||
req = req.query(&[("job", j)]);
|
||||
}
|
||||
if let Some(t) = table {
|
||||
req = req.query(&[("table", t)]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteErrorsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.errors.into_iter().map(JobErrorInfo::from).collect())
|
||||
}
|
||||
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>> {
|
||||
let identifier = build_table_identifier(
|
||||
&request.name,
|
||||
@@ -1580,6 +2206,227 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_derived_compute_routes() {
|
||||
// create_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["language"], "python");
|
||||
assert_eq!(body["return_type"], "FLOAT[4]");
|
||||
assert_eq!(body["body"], "def embed(x): ...");
|
||||
assert_eq!(body["options"]["pip"], "torch");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.create_function(crate::database::CreateFunctionRequest {
|
||||
name: "embed".into(),
|
||||
language: "python".into(),
|
||||
return_type: "FLOAT[4]".into(),
|
||||
body: "def embed(x): ...".into(),
|
||||
options: [("pip".to_string(), "torch".to_string())].into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// list_functions
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/function/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"functions":[{"name":"embed","language":"python","return_type":"Float32","description":""}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let functions = conn.list_functions().await.unwrap();
|
||||
assert_eq!(functions.len(), 1);
|
||||
assert_eq!(functions[0].name, "embed");
|
||||
|
||||
// drop_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_function("embed").await.unwrap();
|
||||
|
||||
// create_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["query"], "SELECT id, embed(body) AS vec FROM docs");
|
||||
assert_eq!(body["auto_refresh"], true);
|
||||
assert_eq!(body["with_no_data"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","job_id":"j-1"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::CreateMaterializedViewRequest::new(
|
||||
"mv1",
|
||||
"SELECT id, embed(body) AS vec FROM docs",
|
||||
);
|
||||
request.auto_refresh = true;
|
||||
let job_id = conn.create_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id.as_deref(), Some("j-1"));
|
||||
|
||||
// refresh_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/refresh");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("src_version").is_none());
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-2"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::RefreshMaterializedViewRequest::new("mv1");
|
||||
request.num_workers = Some(2);
|
||||
let job_id = conn.refresh_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id, "j-2");
|
||||
|
||||
// alter_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/alter");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["auto_refresh"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.alter_materialized_view("mv1", false).await.unwrap();
|
||||
|
||||
// drop_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_materialized_view("mv1").await.unwrap();
|
||||
|
||||
// list_materialized_views
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"views":[{"name":"mv1","source_table":"docs","projection":["id"],"udf_columns":["vec=embed(body)"],"filter":null,"auto_refresh":true}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let views = conn.list_materialized_views().await.unwrap();
|
||||
assert_eq!(views.len(), 1);
|
||||
assert_eq!(views[0].source_table, "docs");
|
||||
assert!(views[0].auto_refresh);
|
||||
|
||||
// list_jobs: platform listing with status payloads
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"job_id":"plat-3","table":"docs","job_type":"indexer","job_subtype":"udf","state":"in_progress","created_at_millis":1000,"status":{"units_done":1,"units_total":2}}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let jobs = conn.list_jobs().await.unwrap();
|
||||
assert_eq!(jobs.len(), 1);
|
||||
assert_eq!(jobs[0].state, "running");
|
||||
assert_eq!(jobs[0].job_type, "udf");
|
||||
assert_eq!(jobs[0].units_total, Some(2));
|
||||
|
||||
// cancel_job: resolve via the manifest-id list filter, then cancel
|
||||
let conn = Connection::new_with_handler(|request| match request.url().path() {
|
||||
"/v1/jobs/list" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[{"job_id":"plat-3","state":"in_progress"}]}"#)
|
||||
.unwrap(),
|
||||
"/v1/jobs/cancel" => {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id":"plat-3"}"#)
|
||||
.unwrap()
|
||||
}
|
||||
other => panic!("unexpected path {other}"),
|
||||
});
|
||||
assert!(conn.cancel_job("j-3").await.unwrap());
|
||||
|
||||
// cancel_job: never registered -> false, and no cancel request
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[]}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(!conn.cancel_job("gone").await.unwrap());
|
||||
|
||||
// job_history(None): a registry listing, timeline-free
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/jobs/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"job_id":"plat-1","table":"docs","job_type":"indexer","job_subtype":"udf","state":"done","created_at_millis":1000,"status":{"rows_skipped":3}}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let hist = conn.job_history(None).await.unwrap();
|
||||
assert_eq!(hist.len(), 1);
|
||||
assert_eq!(hist[0].state, "finished");
|
||||
assert_eq!(hist[0].rows_skipped, Some(3));
|
||||
|
||||
// job_history(id): unknown everywhere -> empty
|
||||
let conn = Connection::new_with_handler(|request| match request.url().path() {
|
||||
"/v1/jobs/describe" => http::Response::builder().status(404).body("").unwrap(),
|
||||
"/v1/jobs/list" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[]}"#)
|
||||
.unwrap(),
|
||||
other => panic!("unexpected path {other}"),
|
||||
});
|
||||
assert!(conn.job_history(Some("j-1")).await.unwrap().is_empty());
|
||||
|
||||
// errors: GET /v1/errors with job + table filters
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/errors");
|
||||
assert_eq!(request.url().query(), Some("job=j-1&table=docs"));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"errors":[{"job_id":"j-1","table":"docs","column":"vec","error_type":"ValueError","error_message":"boom","fragment_id":0,"source_row_id":42,"table_version":7,"age_seconds":5}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let errs = conn.errors(Some("j-1"), Some("docs")).await.unwrap();
|
||||
assert_eq!(errs.len(), 1);
|
||||
assert_eq!(errs[0].error_type, "ValueError");
|
||||
assert_eq!(errs[0].source_row_id, Some(42));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clone_table() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
|
||||
@@ -36,7 +36,7 @@ use crate::{DistanceType, Error};
|
||||
use crate::{
|
||||
error::Result,
|
||||
index::{IndexBuilder, IndexConfig},
|
||||
query::QueryExecutionOptions,
|
||||
query::{AnalyzePlanDistributedMetrics, QueryExecutionOptions},
|
||||
table::{
|
||||
AddDataBuilder, BaseTable, OptimizeAction, OptimizeStats, TableDefinition, UpdateBuilder,
|
||||
merge::MergeInsertBuilder,
|
||||
@@ -1250,8 +1250,7 @@ impl<S: HttpSend + 'static> RemoteTable<S> {
|
||||
|
||||
match result {
|
||||
Ok(_) => {
|
||||
let add_result = insert
|
||||
.as_any()
|
||||
let add_result = (insert.as_ref() as &dyn std::any::Any)
|
||||
.downcast_ref::<RemoteInsertExec<S>>()
|
||||
.and_then(|i| i.add_result())
|
||||
.unwrap_or(AddResult { version: 0 });
|
||||
@@ -1993,9 +1992,16 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
async fn analyze_plan(
|
||||
&self,
|
||||
query: &AnyQuery,
|
||||
_options: QueryExecutionOptions,
|
||||
options: QueryExecutionOptions,
|
||||
) -> Result<String> {
|
||||
let request = self.post_read(&format!("/v1/table/{}/analyze_plan/", self.identifier));
|
||||
let mut request = self.post_read(&format!("/v1/table/{}/analyze_plan/", self.identifier));
|
||||
|
||||
if options.analyze_plan_distributed_metrics != AnalyzePlanDistributedMetrics::Aggregate {
|
||||
request = request.query(&[(
|
||||
"distributed_metrics",
|
||||
options.analyze_plan_distributed_metrics.as_query_param(),
|
||||
)]);
|
||||
}
|
||||
|
||||
let query_bodies = self.prepare_query_bodies(query).await?;
|
||||
let requests: Vec<reqwest::RequestBuilder> = query_bodies
|
||||
@@ -2103,7 +2109,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(delete_response)
|
||||
}
|
||||
|
||||
async fn create_index(&self, mut index: IndexBuilder) -> Result<()> {
|
||||
async fn create_index(&self, mut index: IndexBuilder) -> Result<Option<String>> {
|
||||
self.check_mutable().await?;
|
||||
let request = self
|
||||
.client
|
||||
@@ -2196,14 +2202,28 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
|
||||
// The server returns a job id only when the build was deferred to a
|
||||
// background job (pending vector index). Older servers return an
|
||||
// empty body; treat anything unparseable as "no job".
|
||||
#[derive(serde::Deserialize)]
|
||||
struct CreateIndexResponse {
|
||||
job_id: Option<String>,
|
||||
}
|
||||
let job_id = response
|
||||
.text()
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|body| serde_json::from_str::<CreateIndexResponse>(&body).ok())
|
||||
.and_then(|r| r.job_id);
|
||||
|
||||
if let Some(wait_timeout) = index.wait_timeout {
|
||||
let index_name = index.name.unwrap_or_else(|| format!("{}_idx", column));
|
||||
self.wait_for_index(&[&index_name], wait_timeout).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(job_id)
|
||||
}
|
||||
|
||||
/// Poll until the columns are fully indexed. Will return Error::Timeout if the columns
|
||||
@@ -2416,6 +2436,126 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
message: "optimize is not supported on LanceDB cloud.".into(),
|
||||
})
|
||||
}
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[(String, String)],
|
||||
expression: &str,
|
||||
) -> Result<()> {
|
||||
let new_columns: Vec<serde_json::Value> = columns
|
||||
.iter()
|
||||
.map(|(name, data_type)| {
|
||||
serde_json::json!({
|
||||
"name": name,
|
||||
"computed": { "data_type": data_type, "expression": expression },
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/add_columns/", self.identifier))
|
||||
.json(&serde_json::json!({ "new_columns": new_columns }));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn refresh_column(
|
||||
&self,
|
||||
columns: &[String],
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
let mut body = serde_json::json!({ "columns": columns });
|
||||
if let Some(w) = where_clause {
|
||||
body["where_clause"] = serde_json::Value::String(w);
|
||||
}
|
||||
if let Some(n) = num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(p) = priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/refresh_column", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RefreshColumnResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: RefreshColumnResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn load_columns(&self, request: crate::table::LoadColumnsRequest) -> Result<String> {
|
||||
let columns: Vec<serde_json::Value> = request
|
||||
.columns
|
||||
.iter()
|
||||
.map(|(target, source)| {
|
||||
serde_json::json!({
|
||||
"target": target,
|
||||
"source": source.clone().unwrap_or_else(|| target.clone()),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let mut source = serde_json::json!({
|
||||
"uris": request.source_uris,
|
||||
"format": request.source_format,
|
||||
});
|
||||
if let Some(opts) = request.source_storage_options {
|
||||
source["storage_options"] = serde_json::to_value(opts).unwrap_or_default();
|
||||
}
|
||||
let mut body = serde_json::json!({
|
||||
"columns": columns,
|
||||
"source": source,
|
||||
"target_key": request.target_key,
|
||||
});
|
||||
if let Some(k) = request.source_key {
|
||||
body["source_key"] = serde_json::Value::String(k);
|
||||
}
|
||||
if let Some(m) = request.on_missing {
|
||||
body["on_missing"] = serde_json::Value::String(m);
|
||||
}
|
||||
if let Some(n) = request.num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.commit_granularity {
|
||||
body["commit_granularity"] = n.into();
|
||||
}
|
||||
if let Some(p) = request.priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let http_request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/load_columns", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(http_request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct LoadColumnsResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: LoadColumnsResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn add_columns(
|
||||
&self,
|
||||
transforms: NewColumnTransform,
|
||||
@@ -2817,8 +2957,7 @@ mod tests {
|
||||
use super::*;
|
||||
|
||||
use crate::remote::client::{ClientConfig, RetryConfig};
|
||||
use crate::table::AddDataMode;
|
||||
use crate::table::FieldMetadataUpdate;
|
||||
use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken};
|
||||
|
||||
use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type};
|
||||
use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, record_batch};
|
||||
@@ -2841,7 +2980,10 @@ mod tests {
|
||||
use crate::{
|
||||
DistanceType, Error, Table,
|
||||
index::{Index, IndexStatistics, IndexType, vector::IvfPqIndexBuilder},
|
||||
query::{ColumnOrdering, ExecutableQuery, QueryBase},
|
||||
query::{
|
||||
AnalyzePlanDistributedMetrics, ColumnOrdering, ExecutableQuery, QueryBase,
|
||||
QueryExecutionOptions,
|
||||
},
|
||||
remote::ARROW_FILE_CONTENT_TYPE,
|
||||
};
|
||||
|
||||
@@ -2908,6 +3050,75 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_column() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/refresh_column");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["columns"], serde_json::json!(["vec"]));
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("where_clause").is_none());
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-9"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let job_id = table
|
||||
.refresh_column(&["vec".to_string()], None, Some(2), None, None, None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job_id, "j-9");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_columns() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/load_columns");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
body["columns"],
|
||||
serde_json::json!([{"target": "embedding", "source": "emb"}])
|
||||
);
|
||||
assert_eq!(body["source"]["format"], "parquet");
|
||||
assert_eq!(
|
||||
body["source"]["uris"],
|
||||
serde_json::json!(["s3://b/x.parquet"])
|
||||
);
|
||||
assert_eq!(body["target_key"], "document_id");
|
||||
assert_eq!(body["source_key"], "doc_id");
|
||||
assert_eq!(body["on_missing"], "null");
|
||||
assert_eq!(body["num_workers"], 4);
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"lc-7"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let request = crate::table::LoadColumnsRequest {
|
||||
source_uris: vec!["s3://b/x.parquet".to_string()],
|
||||
source_format: "parquet".to_string(),
|
||||
source_storage_options: None,
|
||||
target_key: "document_id".to_string(),
|
||||
source_key: Some("doc_id".to_string()),
|
||||
columns: vec![("embedding".to_string(), Some("emb".to_string()))],
|
||||
on_missing: Some("null".to_string()),
|
||||
num_workers: Some(4),
|
||||
max_workers: None,
|
||||
batch_size: None,
|
||||
commit_granularity: None,
|
||||
priority: None,
|
||||
};
|
||||
let job_id = table.load_columns(request).await.unwrap();
|
||||
assert_eq!(job_id, "lc-7");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_version() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
@@ -4049,6 +4260,42 @@ mod tests {
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_analyze_plan_distributed_metrics_query_param() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/");
|
||||
assert_eq!(
|
||||
request
|
||||
.url()
|
||||
.query_pairs()
|
||||
.find(|(k, _)| k == "distributed_metrics"),
|
||||
Some(("distributed_metrics".into(), "per_worker".into()))
|
||||
);
|
||||
|
||||
let body = request.body().unwrap().as_bytes().unwrap();
|
||||
let body: serde_json::Value = serde_json::from_slice(body).unwrap();
|
||||
assert_eq!(body["k"], serde_json::json!(1));
|
||||
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#""analyzed plan""#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let result = table
|
||||
.query()
|
||||
.limit(1)
|
||||
.analyze_plan_with_options(QueryExecutionOptions {
|
||||
analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, "analyzed plan");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_structured_fts() {
|
||||
let table =
|
||||
@@ -4406,6 +4653,42 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_returns_deferred_job_id() {
|
||||
let table =
|
||||
Table::new_with_handler("my_table", move |request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => {
|
||||
let schema = Schema::new(vec![Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
128,
|
||||
),
|
||||
false,
|
||||
)]);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap()
|
||||
}
|
||||
"/v1/table/my_table/create_index/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"job_id": "0a1b2c3d-4e5f-6789-abcd-ef0123456789"}"#.to_string())
|
||||
.unwrap(),
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
});
|
||||
|
||||
let job_id = table
|
||||
.create_index(&["vector"], Index::IvfPq(Default::default()))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
job_id.as_deref(),
|
||||
Some("0a1b2c3d-4e5f-6789-abcd-ef0123456789")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_nested_field_paths() {
|
||||
let schema = nested_index_schema();
|
||||
@@ -4888,6 +5171,141 @@ mod tests {
|
||||
assert_eq!(text_idx.created_at, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tokenize_uses_remote_index_details() {
|
||||
let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]);
|
||||
let index_details = serde_json::json!({
|
||||
"base_tokenizer": "icu",
|
||||
"language": "English",
|
||||
"with_position": false,
|
||||
"max_token_length": 40,
|
||||
"lower_case": true,
|
||||
"stem": false,
|
||||
"remove_stop_words": false,
|
||||
"ascii_folding": true,
|
||||
})
|
||||
.to_string();
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/index/list/" => {
|
||||
let body = serde_json::json!({
|
||||
"indexes": [
|
||||
{
|
||||
"index_name": "text_idx",
|
||||
"columns": ["text"],
|
||||
"index_type": "FTS",
|
||||
"index_details": index_details,
|
||||
},
|
||||
]
|
||||
});
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(serde_json::to_string(&body).unwrap())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
}
|
||||
});
|
||||
|
||||
let tokens = table
|
||||
.tokenize("Hello, こんにちは世界!", "text_idx")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
tokens,
|
||||
vec![
|
||||
FtsToken {
|
||||
text: "hello".to_string(),
|
||||
position: 0,
|
||||
},
|
||||
FtsToken {
|
||||
text: "こんにちは".to_string(),
|
||||
position: 1,
|
||||
},
|
||||
FtsToken {
|
||||
text: "世界".to_string(),
|
||||
position: 2,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tokenize_requires_existing_index_name() {
|
||||
let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]);
|
||||
let table = Table::new_with_handler("my_table", move |request| -> http::Response<String> {
|
||||
assert_eq!(request.method(), "POST");
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/index/list/" => {
|
||||
let body = serde_json::json!({ "indexes": [] });
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(serde_json::to_string(&body).unwrap())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
}
|
||||
});
|
||||
|
||||
let err = table.tokenize("hello", "text_idx").await.unwrap_err();
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::InvalidInput { message }
|
||||
if message.contains("No index named 'text_idx'")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tokenize_with_column_remote_requires_index_details() {
|
||||
let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]);
|
||||
let table = Table::new_with_handler("my_table", move |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap(),
|
||||
"/v1/table/my_table/index/list/" => {
|
||||
let body = serde_json::json!({
|
||||
"indexes": [
|
||||
{
|
||||
"index_name": "text_idx",
|
||||
"columns": ["text"],
|
||||
"index_type": "FTS",
|
||||
},
|
||||
]
|
||||
});
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(serde_json::to_string(&body).unwrap())
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected path: {}", path),
|
||||
}
|
||||
});
|
||||
|
||||
let err = table
|
||||
.tokenize_with_column("hello", "text")
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
Error::InvalidInput { message }
|
||||
if message.contains("does not include tokenizer details")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deserialize_created_at() {
|
||||
#[derive(Deserialize)]
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
|
||||
//! DataFusion ExecutionPlan for inserting data into remote LanceDB tables.
|
||||
|
||||
use std::any::Any;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use arrow_array::{ArrayRef, RecordBatch, UInt64Array};
|
||||
@@ -237,10 +236,6 @@ impl<S: HttpSend + 'static> ExecutionPlan for RemoteInsertExec<S> {
|
||||
Self::static_name()
|
||||
}
|
||||
|
||||
fn as_any(&self) -> &dyn Any {
|
||||
self
|
||||
}
|
||||
|
||||
fn properties(&self) -> &Arc<PlanProperties> {
|
||||
&self.properties
|
||||
}
|
||||
|
||||
+364
-46
@@ -23,10 +23,13 @@ use lance::dataset::{InsertBuilder, WriteParams};
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance::io::{ObjectStoreParams, WrappingObjectStore};
|
||||
use lance_datafusion::utils::StreamingWriteSource;
|
||||
use lance_index::IndexCriteria;
|
||||
use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsAccessor};
|
||||
pub use query::AnyQuery;
|
||||
|
||||
use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
|
||||
use lance_index::scalar::InvertedIndexParams;
|
||||
use lance_index::scalar::inverted::query::collect_query_tokens;
|
||||
use lance_namespace::LanceNamespace;
|
||||
use lance_namespace::error::NamespaceError;
|
||||
use lance_namespace::models::DescribeTableRequest;
|
||||
@@ -42,6 +45,7 @@ use std::sync::Arc;
|
||||
|
||||
use crate::connection::NamespaceClientPushdownOperation;
|
||||
|
||||
use crate::DistanceType;
|
||||
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
|
||||
use crate::database::Database;
|
||||
use crate::database::read_freshness::TableFreshness;
|
||||
@@ -49,10 +53,10 @@ use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
use crate::index::IndexStatistics;
|
||||
use crate::index::{Index, IndexBuilder};
|
||||
use crate::index::{IndexConfig, IndexStatisticsImpl};
|
||||
use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType};
|
||||
use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery};
|
||||
use crate::table::datafusion::insert::InsertExec;
|
||||
use crate::utils::{PatchReadParam, PatchWriteParam};
|
||||
use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path};
|
||||
|
||||
use self::dataset::DatasetConsistencyWrapper;
|
||||
use self::merge::MergeInsertBuilder;
|
||||
@@ -471,6 +475,60 @@ impl LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
/// A token produced by the tokenizer configured on a full-text search index.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
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,
|
||||
}
|
||||
|
||||
/// Tokenize a full-text search query using an explicit FTS tokenizer configuration.
|
||||
///
|
||||
/// This does not require a table or FTS index. Use
|
||||
/// [`crate::index::scalar::FtsIndexBuilder`] to supply the same tokenizer
|
||||
/// options used when creating an FTS index.
|
||||
pub fn tokenize(query: &str, params: &InvertedIndexParams) -> Result<Vec<FtsToken>> {
|
||||
let mut tokenizer = params.build().map_err(|err| Error::InvalidInput {
|
||||
message: format!("Failed to build tokenizer: {}", err),
|
||||
})?;
|
||||
let tokens = collect_query_tokens(query, &mut tokenizer);
|
||||
Ok((0..tokens.len())
|
||||
.map(|idx| FtsToken {
|
||||
text: tokens.get_token(idx).to_string(),
|
||||
position: tokens.position(idx),
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Request to fill existing table columns from an external source by
|
||||
/// primary-key join (Geneva `Table.load_columns()` parity). Server-backed
|
||||
/// feature (LanceDB Enterprise / Cloud).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoadColumnsRequest {
|
||||
/// External source URIs.
|
||||
pub source_uris: Vec<String>,
|
||||
/// Source format: "parquet" | "lance" | "ipc".
|
||||
pub source_format: String,
|
||||
/// Source-only storage options (e.g. cloud credentials).
|
||||
pub source_storage_options: Option<HashMap<String, String>>,
|
||||
/// Destination primary-key column.
|
||||
pub target_key: String,
|
||||
/// Source primary-key column. Defaults to `target_key` when None.
|
||||
pub source_key: Option<String>,
|
||||
/// Value column mappings as `(target, source)`; a None source defaults to
|
||||
/// the target name.
|
||||
pub columns: Vec<(String, Option<String>)>,
|
||||
/// Missing-row policy: "carry" (default) | "null" | "error".
|
||||
pub on_missing: Option<String>,
|
||||
pub num_workers: Option<u32>,
|
||||
pub max_workers: Option<u32>,
|
||||
pub batch_size: Option<u32>,
|
||||
pub commit_granularity: Option<u32>,
|
||||
pub priority: Option<String>,
|
||||
}
|
||||
|
||||
/// A trait for anything "table-like". This is used for both native tables (which target
|
||||
/// Lance datasets) and remote tables (which target LanceDB cloud)
|
||||
///
|
||||
@@ -524,7 +582,10 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
/// Update rows in the table.
|
||||
async fn update(&self, update: UpdateBuilder) -> Result<UpdateResult>;
|
||||
/// Create an index on the provided column(s).
|
||||
async fn create_index(&self, index: IndexBuilder) -> Result<()>;
|
||||
///
|
||||
/// Returns the server-minted job id when the build was deferred to a
|
||||
/// background job (remote tables only); `None` for synchronous builds.
|
||||
async fn create_index(&self, index: IndexBuilder) -> Result<Option<String>>;
|
||||
/// List the indices on the table.
|
||||
async fn list_indices(&self) -> Result<Vec<IndexConfig>>;
|
||||
/// Drop an index from the table.
|
||||
@@ -630,6 +691,47 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
transforms: NewColumnTransform,
|
||||
read_columns: Option<Vec<String>>,
|
||||
) -> Result<AddColumnsResult>;
|
||||
/// Declare computed columns bound to a registered function: each
|
||||
/// `(name, sql_type)` is added all-null with the expression stored
|
||||
/// as its binding; no compute happens here (the server's lazy
|
||||
/// detector or refresh_column fills them). Several columns map a
|
||||
/// struct-returning function's fields positionally. Server-backed
|
||||
/// feature; the default returns NotSupported.
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
_columns: &[(String, String)],
|
||||
_expression: &str,
|
||||
) -> Result<()> {
|
||||
Err(Error::NotSupported {
|
||||
message: "computed columns are not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Trigger recompute of computed columns. The expression is
|
||||
/// resolved server-side from each column's stored binding; columns
|
||||
/// bound to the same struct-returning function refresh together.
|
||||
/// Returns the refresh job id. Server-backed feature (LanceDB
|
||||
/// Enterprise / Cloud); the default returns NotSupported.
|
||||
async fn refresh_column(
|
||||
&self,
|
||||
_columns: &[String],
|
||||
_where_clause: Option<String>,
|
||||
_num_workers: Option<u32>,
|
||||
_max_workers: Option<u32>,
|
||||
_batch_size: Option<u32>,
|
||||
_priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
Err(Error::NotSupported {
|
||||
message: "refresh_column is not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Fill existing columns from an external source by primary-key join
|
||||
/// (Geneva `load_columns`). Returns the load job id. Server-backed feature;
|
||||
/// the default returns NotSupported.
|
||||
async fn load_columns(&self, _request: LoadColumnsRequest) -> Result<String> {
|
||||
Err(Error::NotSupported {
|
||||
message: "load_columns is not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Alter columns in the table.
|
||||
async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>;
|
||||
/// Drop columns from the table.
|
||||
@@ -1471,6 +1573,48 @@ impl Table {
|
||||
self.inner.add_columns(transforms, read_columns).await
|
||||
}
|
||||
|
||||
/// Declare computed columns bound to a registered function
|
||||
/// (`(name, sql_type)` pairs + a `f(args)` expression). No compute
|
||||
/// happens here. Server-backed feature.
|
||||
pub async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[(String, String)],
|
||||
expression: &str,
|
||||
) -> Result<()> {
|
||||
self.inner.add_computed_columns(columns, expression).await
|
||||
}
|
||||
|
||||
/// Trigger recompute of computed columns (REFRESH COLUMN). The
|
||||
/// expression comes from each column's stored binding; columns
|
||||
/// bound to the same struct-returning function refresh together.
|
||||
/// Returns the refresh job id. Server-backed feature.
|
||||
pub async fn refresh_column(
|
||||
&self,
|
||||
columns: &[String],
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
self.inner
|
||||
.refresh_column(
|
||||
columns,
|
||||
where_clause,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
priority,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Fill existing columns from an external Parquet/Lance/IPC source by
|
||||
/// primary-key join (Geneva `Table.load_columns()`). Returns the job id.
|
||||
pub async fn load_columns(&self, request: LoadColumnsRequest) -> Result<String> {
|
||||
self.inner.load_columns(request).await
|
||||
}
|
||||
|
||||
/// Change a column's name or nullability.
|
||||
pub async fn alter_columns(
|
||||
&self,
|
||||
@@ -1659,6 +1803,111 @@ impl Table {
|
||||
self.inner.list_indices().await
|
||||
}
|
||||
|
||||
/// Tokenize a full-text search 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. For remote tables, this means the
|
||||
/// same tokenizer model files must also exist locally.
|
||||
pub async fn tokenize(&self, query: &str, index_name: &str) -> Result<Vec<FtsToken>> {
|
||||
let indices = self.inner.list_indices().await?;
|
||||
let matches = indices
|
||||
.iter()
|
||||
.filter(|idx| idx.name == index_name)
|
||||
.collect::<Vec<_>>();
|
||||
let index = match matches.as_slice() {
|
||||
[index] => *index,
|
||||
[] => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("No index named '{}'", index_name),
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("Index name '{}' is ambiguous", index_name),
|
||||
});
|
||||
}
|
||||
};
|
||||
if index.index_type != IndexType::FTS {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("Index '{}' is not a full text search index", index_name),
|
||||
});
|
||||
}
|
||||
self.tokenize_with_index(query, index, index_name)
|
||||
}
|
||||
|
||||
/// Tokenize a full-text search query using the tokenizer configured on the
|
||||
/// FTS index for a column.
|
||||
///
|
||||
/// The column must have exactly one FTS index. 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.
|
||||
pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result<Vec<FtsToken>> {
|
||||
let schema = self.inner.schema().await?;
|
||||
let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?;
|
||||
let indices = self.inner.list_indices().await?;
|
||||
let matches = indices
|
||||
.iter()
|
||||
.filter(|idx| {
|
||||
idx.index_type == IndexType::FTS
|
||||
&& idx.columns.len() == 1
|
||||
&& idx.columns[0] == column
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let index = match matches.as_slice() {
|
||||
[index] => *index,
|
||||
[] => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!("Column '{}' does not have a full text search index", column),
|
||||
});
|
||||
}
|
||||
_ => {
|
||||
return Err(Error::InvalidInput {
|
||||
message: format!(
|
||||
"Column '{}' has multiple full text search indexes; tokenization by column is ambiguous",
|
||||
column
|
||||
),
|
||||
});
|
||||
}
|
||||
};
|
||||
self.tokenize(query, &index.name).await
|
||||
}
|
||||
|
||||
fn tokenize_with_index(
|
||||
&self,
|
||||
query: &str,
|
||||
index: &IndexConfig,
|
||||
index_name: &str,
|
||||
) -> Result<Vec<FtsToken>> {
|
||||
let selector_description = format!("index name '{}'", index_name);
|
||||
let details = index
|
||||
.index_details
|
||||
.as_deref()
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: format!(
|
||||
"Full text search index '{}' for {} does not include tokenizer details",
|
||||
index.name, selector_description
|
||||
),
|
||||
})?;
|
||||
let params = serde_json::from_str::<InvertedIndexParams>(details).map_err(|err| {
|
||||
Error::InvalidInput {
|
||||
message: format!(
|
||||
"Failed to parse tokenizer details for full text search index '{}' for {}: {}",
|
||||
index.name, selector_description, err
|
||||
),
|
||||
}
|
||||
})?;
|
||||
tokenize(query, ¶ms).map_err(|err| match err {
|
||||
Error::InvalidInput { message } => Error::InvalidInput {
|
||||
message: format!(
|
||||
"{} for full text search index '{}' for {}",
|
||||
message, index.name, selector_description
|
||||
),
|
||||
},
|
||||
err => err,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get the table URI (storage location)
|
||||
///
|
||||
/// Returns the full storage location of the table (e.g., S3/GCS path).
|
||||
@@ -2785,7 +3034,7 @@ impl BaseTable for NativeTable {
|
||||
Ok(AddResult { version })
|
||||
}
|
||||
|
||||
async fn create_index(&self, opts: IndexBuilder) -> Result<()> {
|
||||
async fn create_index(&self, opts: IndexBuilder) -> Result<Option<String>> {
|
||||
if opts.columns.len() != 1 {
|
||||
return Err(Error::Schema {
|
||||
message: "Multi-column (composite) indices are not yet supported".to_string(),
|
||||
@@ -2808,7 +3057,8 @@ impl BaseTable for NativeTable {
|
||||
}
|
||||
builder.await?;
|
||||
self.dataset.update(dataset);
|
||||
Ok(())
|
||||
// Native builds are synchronous -- there is never a background job.
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn drop_index(&self, index_name: &str) -> Result<()> {
|
||||
@@ -2967,17 +3217,21 @@ impl BaseTable for NativeTable {
|
||||
.await?
|
||||
.into_iter()
|
||||
.filter_map(|idx_desc| {
|
||||
let index_type: crate::index::IndexType = match idx_desc.index_type().parse() {
|
||||
Ok(index_type) => index_type,
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"Failed to parse index type for index {}: {}",
|
||||
idx_desc.name(),
|
||||
e
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let index_type: crate::index::IndexType = idx_desc
|
||||
.index_type()
|
||||
.parse()
|
||||
.unwrap_or(crate::index::IndexType::Unknown);
|
||||
if index_type == crate::index::IndexType::Unknown {
|
||||
// Internal or future index types that this version doesn't recognize
|
||||
// (e.g. Lance's internal FragReuseIndex) are silently excluded from
|
||||
// the user-visible index listing.
|
||||
log::debug!(
|
||||
"Skipping unrecognized index '{}' (type '{}') in list_indices",
|
||||
idx_desc.name(),
|
||||
idx_desc.index_type(),
|
||||
);
|
||||
return None;
|
||||
}
|
||||
|
||||
let field_ids = idx_desc.field_ids();
|
||||
let mut columns = Vec::with_capacity(field_ids.len());
|
||||
@@ -3044,40 +3298,83 @@ impl BaseTable for NativeTable {
|
||||
}
|
||||
|
||||
async fn index_stats(&self, index_name: &str) -> Result<Option<IndexStatistics>> {
|
||||
let stats = match self
|
||||
.dataset
|
||||
.get()
|
||||
.await?
|
||||
.index_statistics(index_name.as_ref())
|
||||
.await
|
||||
{
|
||||
Ok(stats) => stats,
|
||||
Err(lance_core::Error::IndexNotFound { .. }) => return Ok(None),
|
||||
Err(e) => return Err(Error::from(e)),
|
||||
// describe_indices() reads only manifest-level metadata (no index file I/O).
|
||||
// VectorIndexDetails in the manifest carries distance_type for indices written
|
||||
// by recent Lance versions. For older datasets that didn't write those details
|
||||
// we fall back to index_statistics() for vector index types.
|
||||
let dataset = self.dataset.get().await?;
|
||||
|
||||
let mut descriptions = dataset
|
||||
.describe_indices(Some(IndexCriteria::default().with_name(index_name)))
|
||||
.await?;
|
||||
let Some(description) = descriptions.pop() else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let mut stats: IndexStatisticsImpl =
|
||||
serde_json::from_str(&stats).map_err(|e| Error::InvalidInput {
|
||||
message: format!("error deserializing index statistics: {}", e),
|
||||
})?;
|
||||
let index_type: crate::index::IndexType = description
|
||||
.index_type()
|
||||
.parse()
|
||||
.unwrap_or(crate::index::IndexType::Unknown);
|
||||
|
||||
let first_index = stats.indices.pop().ok_or_else(|| Error::InvalidInput {
|
||||
message: "index statistics is empty".to_string(),
|
||||
})?;
|
||||
// Index type should be present at one of the levels.
|
||||
let index_type =
|
||||
stats
|
||||
.index_type
|
||||
.or(first_index.index_type)
|
||||
.ok_or_else(|| Error::InvalidInput {
|
||||
message: "index statistics was missing index type".to_string(),
|
||||
})?;
|
||||
Ok(Some(IndexStatistics {
|
||||
num_indexed_rows: stats.num_indexed_rows,
|
||||
num_unindexed_rows: stats.num_unindexed_rows,
|
||||
let is_vector = matches!(
|
||||
index_type,
|
||||
distance_type: first_index.metric_type,
|
||||
num_indices: stats.num_indices,
|
||||
crate::index::IndexType::IvfFlat
|
||||
| crate::index::IndexType::IvfSq
|
||||
| crate::index::IndexType::IvfPq
|
||||
| crate::index::IndexType::IvfRq
|
||||
| crate::index::IndexType::IvfHnswPq
|
||||
| crate::index::IndexType::IvfHnswSq
|
||||
| crate::index::IndexType::IvfHnswFlat
|
||||
);
|
||||
|
||||
// details() serializes VectorIndexDetails to JSON with an uppercase "metric_type"
|
||||
// field (e.g. "L2", "COSINE"). Parse it with a case-insensitive match.
|
||||
let distance_type = description.details().ok().and_then(|json| {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Details {
|
||||
metric_type: Option<String>,
|
||||
}
|
||||
serde_json::from_str::<Details>(&json)
|
||||
.ok()
|
||||
.and_then(|d| d.metric_type)
|
||||
.and_then(|m| match m.to_uppercase().as_str() {
|
||||
"L2" => Some(DistanceType::L2),
|
||||
"COSINE" => Some(DistanceType::Cosine),
|
||||
"DOT" => Some(DistanceType::Dot),
|
||||
"HAMMING" => Some(DistanceType::Hamming),
|
||||
_ => None,
|
||||
})
|
||||
});
|
||||
|
||||
// Older Lance datasets didn't write VectorIndexDetails, so distance_type won't
|
||||
// be in the manifest. Fall back to index_statistics() only in that case.
|
||||
if is_vector && distance_type.is_none() {
|
||||
let stats = dataset.index_statistics(index_name).await?;
|
||||
let mut stats: IndexStatisticsImpl =
|
||||
serde_json::from_str(&stats).map_err(|e| Error::InvalidInput {
|
||||
message: format!("error deserializing index statistics: {}", e),
|
||||
})?;
|
||||
let first_index = stats.indices.pop().ok_or_else(|| Error::InvalidInput {
|
||||
message: "index statistics is empty".to_string(),
|
||||
})?;
|
||||
return Ok(Some(IndexStatistics {
|
||||
num_indexed_rows: stats.num_indexed_rows,
|
||||
num_unindexed_rows: stats.num_unindexed_rows,
|
||||
index_type,
|
||||
distance_type: first_index.metric_type,
|
||||
num_indices: stats.num_indices,
|
||||
}));
|
||||
}
|
||||
|
||||
let num_indexed_rows = description.rows_indexed() as usize;
|
||||
let total_rows = dataset.count_rows(None).await?;
|
||||
let num_unindexed_rows = total_rows.saturating_sub(num_indexed_rows);
|
||||
Ok(Some(IndexStatistics {
|
||||
num_indexed_rows,
|
||||
num_unindexed_rows,
|
||||
index_type,
|
||||
distance_type,
|
||||
num_indices: Some(description.metadata().len() as u32),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -3234,6 +3531,27 @@ mod tests {
|
||||
use crate::query::{ExecutableQuery, QueryBase};
|
||||
use crate::test_utils::connection::new_test_connection;
|
||||
|
||||
#[test]
|
||||
fn test_tokenize_uses_explicit_simple_tokenizer() {
|
||||
let params =
|
||||
crate::index::scalar::FtsIndexBuilder::default().base_tokenizer("simple".to_string());
|
||||
let tokens = crate::tokenize("Running in cafés", ¶ms).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
tokens,
|
||||
vec![
|
||||
FtsToken {
|
||||
text: "run".to_string(),
|
||||
position: 0,
|
||||
},
|
||||
FtsToken {
|
||||
text: "cafe".to_string(),
|
||||
position: 2,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_open() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
@@ -589,6 +589,7 @@ mod tests {
|
||||
let stats = table.index_stats(index_name).await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 512);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.distance_type, Some(crate::DistanceType::L2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -646,6 +647,7 @@ mod tests {
|
||||
let stats = table.index_stats(index_name).await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 512);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.distance_type, Some(crate::DistanceType::L2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -690,6 +692,10 @@ mod tests {
|
||||
assert_eq!(index.index_type, crate::index::IndexType::IvfHnswFlat);
|
||||
assert_eq!(index.columns, vec!["embeddings".to_string()]);
|
||||
assert_eq!(table.count_rows(None).await.unwrap(), 512);
|
||||
let stats = table.index_stats(&index.name).await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 512);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.distance_type, Some(crate::DistanceType::L2));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -747,6 +753,15 @@ mod tests {
|
||||
let stats = table.index_stats(index_name).await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 1);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.index_type, crate::index::IndexType::BTree);
|
||||
assert_eq!(stats.distance_type, None);
|
||||
|
||||
// Rows added after the index was built appear as unindexed.
|
||||
let new_batch = record_batch!(("i", Int32, [2])).unwrap();
|
||||
table.add(new_batch).execute().await.unwrap();
|
||||
let stats = table.index_stats(index_name).await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 1);
|
||||
assert_eq!(stats.num_unindexed_rows, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -795,6 +810,12 @@ mod tests {
|
||||
.map(|b| b.num_rows())
|
||||
.sum::<usize>();
|
||||
assert_eq!(count, 1);
|
||||
|
||||
let stats = table.index_stats("text_idx").await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 1);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.index_type, crate::index::IndexType::Fm);
|
||||
assert_eq!(stats.distance_type, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1188,6 +1209,12 @@ mod tests {
|
||||
let index = configs_iter.next().unwrap();
|
||||
assert_eq!(index.index_type, crate::index::IndexType::Bitmap);
|
||||
assert_eq!(index.columns, vec!["large_data".to_string()]);
|
||||
|
||||
let stats = table.index_stats("category_idx").await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 100);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.index_type, crate::index::IndexType::Bitmap);
|
||||
assert_eq!(stats.distance_type, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1256,6 +1283,12 @@ mod tests {
|
||||
let index = index_configs.into_iter().next().unwrap();
|
||||
assert_eq!(index.index_type, crate::index::IndexType::LabelList);
|
||||
assert_eq!(index.columns, vec!["tags".to_string()]);
|
||||
|
||||
let stats = table.index_stats("tags_idx").await.unwrap().unwrap();
|
||||
assert_eq!(stats.num_indexed_rows, 40);
|
||||
assert_eq!(stats.num_unindexed_rows, 0);
|
||||
assert_eq!(stats.index_type, crate::index::IndexType::LabelList);
|
||||
assert_eq!(stats.distance_type, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user