mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 16:38:31 +00:00
Compare commits
48 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 92db5a1aa6 | |||
| 3fd322a93a | |||
| d8f0982ee8 | |||
| 7276c34c51 | |||
| 1918d1a3b6 | |||
| 3b626efa47 | |||
| 137eac9b50 | |||
| 06b53c97d6 | |||
| 711e05619b | |||
| afc0e5f497 | |||
| c12a6dce9f | |||
| 8ea78e3fbc | |||
| 40238d240a | |||
| 60428e1a32 | |||
| 5b982f2f05 | |||
| cde48fad95 | |||
| 1f2068b9fe | |||
| 7527890607 | |||
| a548e59d49 | |||
| 104fc5a08e | |||
| 715be580d0 | |||
| 0d9c87a079 | |||
| 8e364e6812 | |||
| 32a2776446 | |||
| 285add40dd | |||
| 22bf091de1 | |||
| ff81428a9c | |||
| 75c5c83f12 | |||
| 291e9e37be | |||
| 6c066530e5 | |||
| f428c6a76c | |||
| df89c133ca | |||
| ec763521d4 | |||
| f8dc2f78ee | |||
| 3bcff0165e | |||
| c6db80dd0b | |||
| f84190fe12 | |||
| 122dcd0f66 | |||
| e6661a7285 | |||
| 37466a0390 | |||
| bfce8a510d | |||
| a1261e6299 | |||
| 17c499177f | |||
| d889321b5e | |||
| 8a37f2ad77 | |||
| f94673ae5e | |||
| 3b70fc4c9d | |||
| 3a7b02119b |
@@ -1,6 +1,14 @@
|
||||
---
|
||||
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.
|
||||
description: >-
|
||||
Manage LanceDB table branches through the REST API: list, create, and delete
|
||||
branches; target schema reads, field-metadata updates, and index creation to a
|
||||
named branch; and verify that branch changes remain isolated from main. Use
|
||||
when a task involves branch lifecycle, an experimental or isolated table
|
||||
version, directing an operation to a non-main branch, or confirming that a
|
||||
mutation did not affect main. This skill also explains that LanceDB has no
|
||||
checkout operation; each request selects its target branch in the request
|
||||
body.
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
---
|
||||
name: lancedb
|
||||
description: Use when writing, reviewing, debugging, or documenting LanceDB pipelines in Python or TypeScript, especially code that should work across local LanceDB OSS tables and remote LanceDB Enterprise/Cloud tables. Helps avoid non-portable full-table materialization, choose idiomatic query/search patterns, and apply LanceDB performance defaults for ingestion, indexing, filtering, and diagnostics.
|
||||
---
|
||||
|
||||
# Building LanceDB Pipelines
|
||||
|
||||
Use this skill to produce LanceDB pipelines that are portable between local and remote tables (for LanceDB Enterprise/Cloud) and idiomatic for the selected SDK.
|
||||
|
||||
## LanceDB Table Modes
|
||||
|
||||
LanceDB has two common execution modes:
|
||||
|
||||
- **Local table**: embedded, open source, in-process LanceDB. The client opens data from a local path or object storage URI and executes queries in the application process.
|
||||
- **Remote table**: LanceDB Enterprise/Cloud table opened through a `db://...` URI. The data may be very large, commonly backed by object storage, and queried through a remote service.
|
||||
|
||||
Do NOT assume local-only table helpers exist on remote tables. If the user asks for LanceDB Enterprise, Cloud, `db://...`, production remote access, or a remote table, focus on the remote table path: use `search()` / `query()`, keep reads bounded with `select()` and `limit()`, and avoid table-level full materialization APIs.
|
||||
|
||||
## Workflow
|
||||
|
||||
1. Identify the SDK: Python, TypeScript, or both.
|
||||
2. Identify the table mode: local/embedded OSS, remote Enterprise/Cloud, or portable across both. If the user says "LanceDB Enterprise", choose the remote table path.
|
||||
3. Read the matching language branch before writing or changing code:
|
||||
- Python patterns: `references/python/patterns.md`
|
||||
- Python API quick reference: `references/python/api_reference.md`
|
||||
- Python performance guidance: `references/python/performance.md`
|
||||
- TypeScript patterns: `references/typescript/patterns.md`
|
||||
- TypeScript API quick reference: `references/typescript/api_reference.md`
|
||||
- TypeScript performance guidance: `references/typescript/performance.md`
|
||||
4. Start with `patterns.md` for the selected SDK. Read `api_reference.md` when choosing method names or return collectors. Read `performance.md` when the task involves ingestion, indexing, filtering, query tuning, diagnostics, or large datasets.
|
||||
5. For Python schemas, favor Pydantic models and validate records before writing. Use PyArrow schemas when Arrow-native, streaming, or highly dynamic data makes them materially better suited.
|
||||
6. Prefer `search()` or `query()` builders with explicit `select()` and `limit()` for reads.
|
||||
7. Avoid table-level full materialization in remote or portable code. This is the main local-vs-remote read pitfall.
|
||||
8. After a successful embedded OSS ingestion, call `table.optimize()`. Do not call it for Enterprise/Cloud; remote maintenance is automatic.
|
||||
9. For remote Enterprise/Cloud writes, never drop-then-reuse or `mode="overwrite"` the same table name — see "Enterprise: never drop-then-reuse the same table name" below. This is the main local-vs-remote write pitfall.
|
||||
10. If reviewing an existing file or repo, run `scripts/check_materialization.py` on the relevant paths and inspect each finding before editing.
|
||||
11. Cross-check unfamiliar or non-trivial API claims against the source tree instead of relying on memory.
|
||||
|
||||
## Core Portability Rule
|
||||
|
||||
Do not write code that assumes a local table API will exist on a remote table. Remote tables can be very large, so whole-table materialization helpers are intentionally unavailable or unsafe.
|
||||
|
||||
This does **not** mean result conversion is forbidden. Bounded query/search result collection is normal:
|
||||
|
||||
- Python: `table.search(...).select([...]).limit(10).to_pandas()`
|
||||
- TypeScript: `await table.search(...).select([...]).limit(10).toArray()`
|
||||
|
||||
The unsafe pattern is table-level or unbounded collection, plus local-only dataset escape hatches in remote code:
|
||||
|
||||
- Python: `table.to_pandas()`, `table.to_arrow()`, `table.to_polars()`; `table.to_lance()` is local/OSS-only dataset access, not materialization
|
||||
- TypeScript: `await table.toArrow()`, `await table.query().toArray()` without `limit()`
|
||||
|
||||
## Enterprise: never drop-then-reuse the same table name
|
||||
|
||||
LanceDB Enterprise/Cloud splits a **control plane** (DDL: create/drop/rename) from a **data plane** (query nodes that serve reads). Query nodes cache the resolved dataset for a table name for up to `table_cache_ttl` — **default 300 seconds (5 minutes)**. After you drop or overwrite a table, the control plane updates immediately but the data plane keeps serving the *old* dataset until that cache entry expires. During the window the two planes disagree.
|
||||
|
||||
The failure this causes: you `drop_table("t")` then immediately `create_table("t", ...)` (or `create_table("t", ..., mode="overwrite")`). The DDL returns success, but every query against `t` returns **`500 Internal Server Error`** (the query node resolves the stale/deleted dataset), and a fresh `describe` may still show the *old* schema/version. It looks like your write silently failed; it didn't — the name is cached.
|
||||
|
||||
**`mode="overwrite"` has the same problem** — it is a drop+create of the same name under the hood.
|
||||
|
||||
Rules for portable Enterprise ingestion:
|
||||
|
||||
1. **Never reuse a table name you just dropped/overwrote within the cache TTL.** Do not use `mode="overwrite"` to replace an existing Enterprise table in place.
|
||||
2. To (re)load data, **write to a fresh table name** (e.g. `<table>_v2`, or a run-stamped suffix). A brand-new name has no cached data-plane entry, so writes and reads work immediately.
|
||||
3. Before creating, `list_tables()` and **fail loudly if the name already exists** rather than overwriting — prompt for a new name.
|
||||
4. To land on a specific final name that is currently occupied by an old table: drop the old table, **wait out the TTL (~5 min), then `rename_table(fresh_name, final_name)`**. Renaming onto a name whose old dataset is still cached hits the same race, so the wait is mandatory. `rename_table` is a supported control-plane op.
|
||||
5. When you hand a table name back to a human, tell them which step still needs the propagation wait (usually: "the old `t` was dropped; run the rename in ~5 minutes").
|
||||
|
||||
This is Enterprise/Cloud-specific. Local/OSS tables have no separate data plane, so `mode="overwrite"` and immediate same-name reuse are fine there.
|
||||
|
||||
## Script
|
||||
|
||||
Run the scanner when reviewing or modifying an existing codebase:
|
||||
|
||||
```bash
|
||||
python skills/lancedb/scripts/check_materialization.py path/to/file_or_dir
|
||||
```
|
||||
|
||||
The script reports likely unsafe full-table materialization in Python and TypeScript. Treat results as review prompts, not automatic proof of a bug.
|
||||
@@ -0,0 +1,105 @@
|
||||
# Python API Reference
|
||||
|
||||
Quick method reference for Python LanceDB code. Cross-check source for non-trivial claims.
|
||||
|
||||
## Connect
|
||||
|
||||
```python
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect("./camelot-db") # local/OSS
|
||||
db = lancedb.connect("db://my-db", api_key=api_key, region=region) # remote
|
||||
```
|
||||
|
||||
**Place the local database directory next to the script/entrypoint that opens it** (i.e. resolve the path relative to the script, `Path(__file__).parent / "camelot-db"`), not buried under a shared `data/` folder. The Lance dataset is the database, not a data file — keeping it beside its code makes ownership obvious and paths stable regardless of the working directory the script is launched from.
|
||||
|
||||
**Do not name the directory `lancedb`** (e.g. `./lancedb`, `./data/lancedb`). It collides with the imported `lancedb` package name, which is confusing to read and easy to shadow in scripts. Give it a name derived from the repo or dataset with a clear prefix/suffix — for example `./<dataset>-db`, `./<repo>_lancedb`, or `./vectordb`.
|
||||
|
||||
Async:
|
||||
|
||||
```python
|
||||
db = await lancedb.connect_async("./camelot-db")
|
||||
```
|
||||
|
||||
## Table Reads
|
||||
|
||||
| Task | Preferred API |
|
||||
| --- | --- |
|
||||
| Vector search | `table.search(query_vector).limit(k)` |
|
||||
| Full scan with filters/projection (sync) | `table.search().where(...).select(...).limit(...)` |
|
||||
| Full scan with filters/projection (async) | `table.query().where(...).select(...).limit(...)` |
|
||||
| Filter | `.where("col > 10")` |
|
||||
| Projection | `.select(["id", "text"])` |
|
||||
| Bound result count | `.limit(20)` |
|
||||
| Collect bounded result as Python objects (default, no extra deps) | `.to_list()` on query/search result |
|
||||
| Collect bounded result as Arrow (default, `pyarrow` always available) | `.to_arrow()` on query/search result |
|
||||
| Collect bounded result as pandas (only if project uses pandas) | `.to_pandas()` on query/search result |
|
||||
| Collect bounded result as Polars (only if project uses polars) | `.to_polars()` on query/search result |
|
||||
|
||||
## Sync vs Async Scan API
|
||||
|
||||
The plain-scan entry point differs between the sync and async clients. **Verified against `lancedb` 0.34.0** — re-check if the pinned version changes:
|
||||
|
||||
- **Sync** (`lancedb.connect(...)`): the table has **no `.query()` method**. Use `.search()` with no argument for a plain scan; it returns a query builder that supports `.where()`, `.select()`, `.limit()`, and the `.to_list()` / `.to_arrow()` / `.to_pandas()` / `.to_polars()` collectors.
|
||||
```python
|
||||
rows = table.search().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
- **Async** (`lancedb.connect_async(...)`): the table has **both** `.query()` and `.search()`. Use `.query()` for a plain scan.
|
||||
```python
|
||||
rows = await async_table.query().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
|
||||
Do not call `table.query()` on a sync table — it raises `AttributeError`.
|
||||
|
||||
## Local vs Remote Table Methods
|
||||
|
||||
| API | Local table | Remote table | Agent guidance |
|
||||
| --- | --- | --- | --- |
|
||||
| `table.search(...)` | Yes | Yes | Preferred read path (sync + async) |
|
||||
| `table.query()` | Async only | Async only | Sync scan path is `table.search()`; `.query()` is the async scan builder |
|
||||
| `table.to_pandas()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_arrow()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_polars()` | Yes | No / unsafe for portability | Avoid in portable code |
|
||||
| `table.to_lance()` | Yes | No | Local/OSS escape hatch only |
|
||||
|
||||
## Indexes
|
||||
|
||||
Use `create_index(...)` for vector indexes and modern index configs. Use scalar indexes for filtered or merge keys.
|
||||
|
||||
Common calls:
|
||||
|
||||
```python
|
||||
table.create_index("vector")
|
||||
table.create_scalar_index("status")
|
||||
table.create_fts_index("text")
|
||||
```
|
||||
|
||||
Check source docs before specifying advanced index config names or parameters.
|
||||
|
||||
## Filtering And Recall Knobs
|
||||
|
||||
```python
|
||||
table.search(query_vector).where("status = 'ready'") # pre-filter by default
|
||||
table.search(query_vector).where("status = 'ready'", prefilter=False)
|
||||
table.search(query_vector).limit(10).refine_factor(20)
|
||||
table.search(query_vector).limit(10).nprobes(50)
|
||||
```
|
||||
|
||||
Use post-filtering only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
```python
|
||||
print(table.search(query_vector).where("year > 2000").limit(10).analyze_plan())
|
||||
print(table.index_stats("vector_idx"))
|
||||
```
|
||||
|
||||
Use these before changing indexes or search tuning.
|
||||
|
||||
## Maintenance
|
||||
|
||||
```python
|
||||
table.optimize()
|
||||
```
|
||||
|
||||
Call this after every successful local/OSS ingestion. It handles compaction, cleanup of old versions according to retention, and index optimization. Do not add this for LanceDB Enterprise/Cloud remote tables; Enterprise handles compaction and cleanup automatically from cluster configuration.
|
||||
@@ -0,0 +1,173 @@
|
||||
# Python Patterns
|
||||
|
||||
Use these patterns when writing Python code with `lancedb`.
|
||||
|
||||
## Before Writing Code
|
||||
|
||||
Choose the output type from what the project actually depends on. **Do not assume `pandas` or `polars` is installed** — they are heavy dependencies that many LanceDB projects do not use. `pyarrow`, by contrast, ships as a LanceDB dependency and is always available, so it is a safe default to lean on.
|
||||
|
||||
Default output (after applying `select()` and `limit()`):
|
||||
|
||||
- **Python objects**: `.to_list()` — a list of dicts, no extra dependencies. Prefer this for scripts, examples, and agent-generated code unless there is a reason to do otherwise.
|
||||
- **PyArrow**: `.to_arrow()` — a `pyarrow.Table`, when the surrounding code is Arrow-native or you need columnar/zero-copy handoff.
|
||||
|
||||
Only reach for a DataFrame when the project *already* declares that dependency:
|
||||
|
||||
- Pandas projects (pandas in `pyproject.toml`/requirements): `.to_pandas()`.
|
||||
- Polars projects (polars declared): `.to_polars()`.
|
||||
|
||||
If unsure, check the dependency manifest or the imports in surrounding files. When in doubt, use `.to_list()` or `.to_arrow()`.
|
||||
|
||||
## Schema Design and Validation
|
||||
|
||||
Favor `LanceModel` and Pydantic validation for Python schemas. They keep field
|
||||
types readable, validate source records before a write, and map directly to a
|
||||
LanceDB schema. Use `Vector(dimension)` for fixed-size vectors:
|
||||
|
||||
```python
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
class Document(LanceModel):
|
||||
id: int
|
||||
text: str
|
||||
vector: Vector(384, nullable=False)
|
||||
|
||||
rows = [Document.model_validate(row) for row in source_rows]
|
||||
table = db.create_table("documents", schema=Document)
|
||||
table.add(rows)
|
||||
```
|
||||
|
||||
Use PyArrow schemas instead when the pipeline is already Arrow-native, needs
|
||||
record-batch streaming, or has runtime schema requirements that would make a
|
||||
Pydantic model harder to understand. Declare Pydantic as a direct project
|
||||
dependency when application code imports it, even if LanceDB also depends on it.
|
||||
|
||||
## Recommended Patterns
|
||||
|
||||
### Bounded search or query
|
||||
|
||||
Use this for application reads, examples, notebooks, and agent-generated scripts:
|
||||
|
||||
```python
|
||||
results = (
|
||||
table.search(query_vector)
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.to_list() # or .to_arrow(); .to_pandas()/.to_polars() only if the project uses them
|
||||
)
|
||||
```
|
||||
|
||||
Why: `search()` works across local and remote tables and on both the sync and async clients. `select()` avoids fetching unused columns. `limit()` prevents accidental full-table reads. `.to_list()` and `.to_arrow()` avoid assuming pandas/polars is installed (see "Before Writing Code").
|
||||
|
||||
For a **plain scan** (no query vector), the entry point differs by client:
|
||||
|
||||
```python
|
||||
# Sync client: no .query() method — use .search() with no argument.
|
||||
rows = table.search().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
|
||||
# Async client: use .query().
|
||||
rows = await async_table.query().where("status = 'ready'").select(["id", "text"]).limit(20).to_list()
|
||||
```
|
||||
|
||||
`table.query()` on a sync table raises `AttributeError` (verified on `lancedb` 0.34.0). See the "Sync vs Async Scan API" section in `api_reference.md`.
|
||||
|
||||
### Bounded query result conversion
|
||||
|
||||
It is fine to collect bounded query/search results:
|
||||
|
||||
```python
|
||||
arrow_table = table.search().select(["id"]).limit(100).to_arrow() # sync plain scan
|
||||
rows = table.search(query_vector).limit(10).to_list()
|
||||
df = table.search(query_vector).limit(10).to_pandas() # only if pandas is a project dep
|
||||
```
|
||||
|
||||
### Local-only Lance dataset API
|
||||
|
||||
`table.to_lance()` does not itself materialize the full dataset. It returns the underlying `lance.LanceDataset`, making the table accessible through the PyLance dataset API. Use it when the task is explicitly local/OSS and needs Lance dataset methods not exposed by LanceDB:
|
||||
|
||||
```python
|
||||
# Local/OSS only: RemoteTable does not expose table.to_lance().
|
||||
ds = table.to_lance()
|
||||
for batch in ds.to_batches(columns=["id", "text"], batch_size=10_000):
|
||||
process(batch)
|
||||
```
|
||||
|
||||
### Async Python
|
||||
|
||||
Keep the same shape and bound the result before collecting:
|
||||
|
||||
```python
|
||||
results = await (
|
||||
async_table.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.to_list() # or .to_arrow()
|
||||
)
|
||||
```
|
||||
|
||||
## Anti-Patterns
|
||||
|
||||
**Avoid the following anti-patterns in your code.**
|
||||
|
||||
### Table-level full materialization
|
||||
|
||||
Avoid whole-table collectors in portable or large-table code:
|
||||
|
||||
```python
|
||||
df = table.to_pandas()
|
||||
arrow_table = table.to_arrow()
|
||||
polars_df = table.to_polars()
|
||||
```
|
||||
|
||||
Why: local tables expose these whole-table collectors, but remote tables intentionally do not — a remote production table can be far larger than a local development table, so it is easy to accidentally pull the entire table into memory.
|
||||
|
||||
`table.to_lance()` is different: it is not a full materialization call, but it is still local/OSS-only and should not appear in code meant to run against remote Enterprise tables.
|
||||
|
||||
### Unbounded result collection
|
||||
|
||||
Avoid query/search collection without a meaningful limit:
|
||||
|
||||
```python
|
||||
rows = table.search().to_list() # unbounded plain scan
|
||||
rows = table.search(query_vector).to_list() # unbounded vector search
|
||||
```
|
||||
|
||||
Prefer `select(...).limit(...)` before collecting; for large reads, stream in batches instead.
|
||||
|
||||
### Per-row writes
|
||||
|
||||
Avoid loops that write one row per call:
|
||||
|
||||
```python
|
||||
for row in rows:
|
||||
table.add([row]) # one commit + fragment per row
|
||||
```
|
||||
|
||||
Each `add()` creates a new version and fragment. Pass the whole batch in a single call, or chunk very large inputs:
|
||||
|
||||
```python
|
||||
table.add(rows) # single commit
|
||||
# for very large inputs, add batches of several thousand rows
|
||||
```
|
||||
|
||||
After the final successful write to an embedded OSS table, call
|
||||
`table.optimize()`. Skip this for Enterprise/Cloud tables because their
|
||||
maintenance is automatic.
|
||||
|
||||
### Drop-then-reuse the same table name (Enterprise/Cloud)
|
||||
|
||||
Avoid dropping or overwriting a remote table and then reusing that name right away:
|
||||
|
||||
```python
|
||||
db.drop_table("my_table")
|
||||
table = db.create_table("my_table", data=rows) # reads 500 for ~5 min
|
||||
table = db.create_table("my_table", data=rows, mode="overwrite") # same problem
|
||||
```
|
||||
|
||||
Why: Enterprise/Cloud splits DDL (control plane) from query serving (data plane). The data plane caches the dataset behind a table name for up to `table_cache_ttl` (default 300s / 5 min), so after a drop/overwrite the DDL succeeds but queries against the reused name return `500 Internal Server Error` until the cache expires — and a fresh `describe` may still show the old schema. Instead, write to a **fresh name**, use `list_tables()` and fail if it already exists, then `rename_table(fresh, final)` onto the final name only after the old table's drop has propagated (~5 min). See the "Enterprise: never drop-then-reuse the same table name" section in `SKILL.md`. Local/OSS tables have no separate data plane — overwrite freely there.
|
||||
|
||||
### Guessing performance fixes
|
||||
|
||||
Avoid changing `nprobes`, `refine_factor`, or index types before checking the query plan and index stats. Diagnose first, then tune one knob at a time.
|
||||
@@ -0,0 +1,131 @@
|
||||
# Python Performance Guidance
|
||||
|
||||
Use this when writing Python code that ingests data, queries large tables, builds indexes, or investigates latency.
|
||||
|
||||
## Ingestion
|
||||
|
||||
### Recommended: validate schemas and records with Pydantic
|
||||
|
||||
Favor `LanceModel` for readable Python schema definitions and validate source
|
||||
records before writing. Use PyArrow directly for Arrow-native or streaming
|
||||
pipelines where it is the clearer representation.
|
||||
|
||||
```python
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
|
||||
class Document(LanceModel):
|
||||
id: int
|
||||
text: str
|
||||
vector: Vector(384, nullable=False)
|
||||
|
||||
rows = [Document.model_validate(row) for row in source_rows]
|
||||
table = db.create_table("documents", schema=Document)
|
||||
table.add(rows)
|
||||
```
|
||||
|
||||
### Recommended: bulk ingestion for materialized data
|
||||
|
||||
```python
|
||||
table.add(arrow_table)
|
||||
table.add(df)
|
||||
table.add(pa.dataset("data/", format="parquet"))
|
||||
```
|
||||
|
||||
For very large initial loads, create the table empty first, then call `add(...)`. Passing data directly to `create_table(name, data)` can skip the auto-parallel write path.
|
||||
|
||||
### Recommended: iterator ingestion for generated or streamed data
|
||||
|
||||
```python
|
||||
def batches():
|
||||
for raw in source:
|
||||
vectors = model.encode(raw["text"])
|
||||
yield pa.RecordBatch.from_pydict({**raw, "vector": vectors})
|
||||
|
||||
table.add(batches())
|
||||
```
|
||||
|
||||
Use chunks of several thousand rows or more when practical. Tiny batches and per-row writes create many small fragments.
|
||||
|
||||
### Anti-pattern: per-row `add()`
|
||||
|
||||
```python
|
||||
for row in rows:
|
||||
table.add([row])
|
||||
```
|
||||
|
||||
Each call creates a version and fragment. This slows ingestion and later queries.
|
||||
|
||||
## Indexing
|
||||
|
||||
- Build a vector index once brute-force vector search becomes too slow. As a rule of thumb, local brute force is fine below roughly 100K vectors; beyond that, build an index.
|
||||
- Use `IVF_PQ` as the general-purpose default. Enterprise builds this automatically.
|
||||
- Use scalar indexes for filtered columns and merge/upsert keys.
|
||||
- Use `BTREE` for mostly distinct numeric/string/temporal columns, `BITMAP` for booleans and low-cardinality columns, and `LABEL_LIST` for list membership queries.
|
||||
- Keep full-text defaults unless phrase queries require position data.
|
||||
|
||||
## Querying
|
||||
|
||||
Always be explicit:
|
||||
|
||||
```python
|
||||
table.search(query_vector).select(["id", "title"]).limit(20)
|
||||
```
|
||||
|
||||
- `select()` reduces bytes read and transferred.
|
||||
- `limit()` prevents accidental full-table materialization.
|
||||
- Pre-filtering is the default and guarantees returned rows satisfy the predicate.
|
||||
- Use post-filtering only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Recall Tuning
|
||||
|
||||
Tune one knob at a time:
|
||||
|
||||
- Quantized indexes: raise `refine_factor` to rescore more candidates on full vectors.
|
||||
- HNSW-backed indexes: raise `ef`; start around `1.5 * k`, increase toward `10 * k` if recall is short.
|
||||
- IVF candidate breadth: `nprobes` is auto-tuned; override only when a selective pre-filter leaves too few neighbors.
|
||||
|
||||
## Maintenance
|
||||
|
||||
After every successful embedded OSS/local ingestion, call `table.optimize()`.
|
||||
Do not add this to LanceDB Enterprise/Cloud remote table code; remote compaction
|
||||
and cleanup are handled automatically based on the Enterprise cluster
|
||||
configuration.
|
||||
|
||||
Why local maintenance is needed:
|
||||
|
||||
- Frequent writes can create many small fragments. Queries then need to scan across more files, which can increase latency.
|
||||
- Updates, deletes, and appends create new table versions. Old versions are retained for time travel and rollback, which can grow disk usage.
|
||||
- Indexes may have newly added rows that are not yet fully optimized into the index structure.
|
||||
|
||||
For local/OSS tables, run `optimize()` after the final successful ingestion
|
||||
write. Also run it after later batches of update/delete operations or on a
|
||||
regular maintenance schedule:
|
||||
|
||||
```python
|
||||
table.optimize()
|
||||
```
|
||||
|
||||
If the user wants more aggressive local disk cleanup, pass a shorter cleanup retention window:
|
||||
|
||||
```python
|
||||
from datetime import timedelta
|
||||
|
||||
table.optimize(cleanup_older_than=timedelta(days=1))
|
||||
```
|
||||
|
||||
Do not use very short cleanup windows when the application depends on time travel, rollback, or old versions.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
Before changing code or indexes, inspect:
|
||||
|
||||
```python
|
||||
print(table.search(query_vector).where("year > 2000").limit(10).analyze_plan())
|
||||
print(table.index_stats("vector_idx"))
|
||||
```
|
||||
|
||||
Look for high scan bytes, missing indexes, fragmented data, and unindexed rows.
|
||||
|
||||
## Python Multiprocessing
|
||||
|
||||
When using multiprocessing, use `spawn` rather than `fork`. LanceDB is multi-threaded internally, and `fork` plus a multi-threaded process is unsafe.
|
||||
@@ -0,0 +1,78 @@
|
||||
# TypeScript API Reference
|
||||
|
||||
Quick method reference for TypeScript LanceDB code. Cross-check source for non-trivial claims.
|
||||
|
||||
## Connect
|
||||
|
||||
```typescript
|
||||
import * as lancedb from "@lancedb/lancedb";
|
||||
|
||||
const db = await lancedb.connect("./camelot-db");
|
||||
```
|
||||
|
||||
**Place the local database directory next to the script/entrypoint that opens it** (resolve the path relative to the module, e.g. via `import.meta.dirname` / `__dirname`), not buried under a shared `data/` folder. The Lance dataset is the database, not a data file — keeping it beside its code makes ownership obvious and paths stable regardless of the working directory the script is launched from.
|
||||
|
||||
**Do not name the directory `lancedb`** (e.g. `./lancedb`, `./data/lancedb`). It collides with the imported `lancedb` package/namespace, which is confusing to read. Give it a name derived from the repo or dataset with a clear prefix/suffix — for example `./<dataset>-db`, `./<repo>_lancedb`, or `./vectordb`.
|
||||
|
||||
Remote connections use `db://...` plus Enterprise/Cloud credentials and deployment settings. Check current source/docs for exact connection options.
|
||||
|
||||
## Table Reads
|
||||
|
||||
| Task | Preferred API |
|
||||
| --- | --- |
|
||||
| Vector search | `table.search(queryVector).limit(k)` |
|
||||
| Full scan with filters/projection | `table.query().where(...).select(...).limit(...)` |
|
||||
| Filter | `.where("col > 10")` |
|
||||
| Projection | `.select(["id", "text"])` |
|
||||
| Bound result count | `.limit(20)` |
|
||||
| Collect bounded result as objects | `.toArray()` on query/search result |
|
||||
| Collect bounded result as Arrow | `.toArrow()` on query/search result |
|
||||
| Stream result batches | `for await (const batch of table.query()...)` |
|
||||
|
||||
## Local vs Remote Safety
|
||||
|
||||
| API | Agent guidance |
|
||||
| --- | --- |
|
||||
| `table.search(...)` | Preferred read path |
|
||||
| `table.query()` | Preferred scan/filter path |
|
||||
| `await table.toArrow()` | Avoid in portable or large-table code |
|
||||
| `await table.query().toArray()` with no `limit()` | Avoid; unbounded collection |
|
||||
| `await table.query().toArrow()` with no `limit()` | Avoid; unbounded collection |
|
||||
|
||||
## Indexes
|
||||
|
||||
```typescript
|
||||
await table.createIndex("vector");
|
||||
await table.createIndex("status");
|
||||
```
|
||||
|
||||
Use vector indexes for large vector search workloads and scalar indexes for filtered columns or merge/upsert keys. Check source/docs before specifying advanced index options.
|
||||
|
||||
## Filtering And Recall Knobs
|
||||
|
||||
```typescript
|
||||
await table.search(queryVector).where("status = 'ready'").limit(10).toArray();
|
||||
await table.search(queryVector).limit(10).refineFactor(20).toArray();
|
||||
await table.search(queryVector).limit(10).nprobes(50).toArray();
|
||||
await table.search(queryVector).limit(10).ef(100).toArray();
|
||||
await table.search(queryVector).where("status = 'ready'").postfilter().limit(10).toArray();
|
||||
```
|
||||
|
||||
Use `postfilter()` only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
```typescript
|
||||
console.log(await table.search(queryVector).where("year > 2000").limit(10).analyzePlan());
|
||||
console.log(await table.indexStats("vector_idx"));
|
||||
```
|
||||
|
||||
Use these before changing indexes or search tuning.
|
||||
|
||||
## Maintenance
|
||||
|
||||
```typescript
|
||||
await table.optimize();
|
||||
```
|
||||
|
||||
Call this after every successful local/OSS ingestion. It handles compaction, cleanup of old versions according to retention, and index optimization. Do not add this for LanceDB Enterprise/Cloud remote tables; Enterprise handles compaction and cleanup automatically from cluster configuration.
|
||||
@@ -0,0 +1,100 @@
|
||||
# TypeScript Patterns
|
||||
|
||||
Use these patterns when writing TypeScript code with `@lancedb/lancedb`.
|
||||
|
||||
## Recommended Patterns
|
||||
|
||||
### Bounded query
|
||||
|
||||
Use this for application reads, scripts, and examples:
|
||||
|
||||
```typescript
|
||||
const rows = await table
|
||||
.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.toArray();
|
||||
```
|
||||
|
||||
### Bounded vector search
|
||||
|
||||
```typescript
|
||||
const rows = await table
|
||||
.search(queryVector)
|
||||
.select(["id", "text"])
|
||||
.limit(20)
|
||||
.toArray();
|
||||
```
|
||||
|
||||
### Batch streaming for larger reads
|
||||
|
||||
When the task needs many rows, avoid collecting everything at once:
|
||||
|
||||
```typescript
|
||||
for await (const batch of table
|
||||
.query()
|
||||
.where("status = 'ready'")
|
||||
.select(["id", "text"])
|
||||
.limit(10_000)) {
|
||||
process(batch);
|
||||
}
|
||||
```
|
||||
|
||||
## Anti-Patterns
|
||||
|
||||
**Avoid the following anti-patterns in your code.**
|
||||
|
||||
### Table-level full materialization
|
||||
|
||||
Avoid whole-table collectors in portable or large-table code:
|
||||
|
||||
```typescript
|
||||
const tableArrow = await table.toArrow();
|
||||
```
|
||||
|
||||
Why: local tables expose these whole-table collectors, but remote tables intentionally do not — a remote production table can be far larger than a local development table, so it is easy to accidentally pull the entire table into memory.
|
||||
|
||||
### Unbounded result collection
|
||||
|
||||
Avoid query/search collection without a meaningful limit:
|
||||
|
||||
```typescript
|
||||
const rows = await table.query().toArray(); // unbounded plain scan
|
||||
const rows = await table.search(queryVector).toArray(); // unbounded vector search
|
||||
```
|
||||
|
||||
Prefer `select(...).limit(...)` before collecting; for large reads, stream in batches instead.
|
||||
|
||||
### Per-row writes
|
||||
|
||||
Avoid loops that write one row per call:
|
||||
|
||||
```typescript
|
||||
for (const row of rows) {
|
||||
await table.add([row]); // one commit + fragment per row
|
||||
}
|
||||
```
|
||||
|
||||
Each `add()` creates a new version and fragment. Pass the whole batch in a single call, or chunk very large inputs:
|
||||
|
||||
```typescript
|
||||
await table.add(rows); // single commit
|
||||
// for very large inputs, add in chunks of several thousand rows
|
||||
```
|
||||
|
||||
### Drop-then-reuse the same table name (Enterprise/Cloud)
|
||||
|
||||
Avoid dropping or overwriting a remote table and then reusing that name right away:
|
||||
|
||||
```typescript
|
||||
await db.dropTable("my_table");
|
||||
const table = await db.createTable("my_table", rows); // reads 500 for ~5 min
|
||||
const table = await db.createTable("my_table", rows, { mode: "overwrite" }); // same problem
|
||||
```
|
||||
|
||||
Why: Enterprise/Cloud splits DDL (control plane) from query serving (data plane). The data plane caches the dataset behind a table name for up to `table_cache_ttl` (default 300s / 5 min), so after a drop/overwrite the DDL succeeds but queries against the reused name return `500 Internal Server Error` until the cache expires — and a fresh `describe` may still show the old schema. Instead, write to a **fresh name**, use `tableNames()` and fail if it already exists, then `renameTable(fresh, final)` onto the final name only after the old table's drop has propagated (~5 min). See the "Enterprise: never drop-then-reuse the same table name" section in `SKILL.md`. Local/OSS tables have no separate data plane — overwrite freely there.
|
||||
|
||||
### Guessing performance fixes
|
||||
|
||||
Avoid changing `nprobes`, `refineFactor`, `ef`, or index settings before checking `analyzePlan()` and `indexStats(...)`. Diagnose first, then tune one knob at a time.
|
||||
@@ -0,0 +1,78 @@
|
||||
# TypeScript Performance Guidance
|
||||
|
||||
Use this when writing TypeScript code that ingests data, queries large tables, builds indexes, or investigates latency.
|
||||
|
||||
## Ingestion
|
||||
|
||||
- Prefer bulk or batched writes.
|
||||
- Avoid per-row write loops; they create many small commits/fragments.
|
||||
- For generated data, accumulate reasonable batches before adding.
|
||||
- For file-backed data, prefer APIs that stream from Arrow/Parquet-style inputs when available.
|
||||
|
||||
## Indexing
|
||||
|
||||
- Build a vector index once brute-force vector search becomes too slow. As a rule of thumb, local brute force is fine below roughly 100K vectors; beyond that, build an index.
|
||||
- Use the general-purpose vector index defaults unless the task has explicit recall/latency requirements.
|
||||
- Build scalar indexes for filtered columns and merge/upsert keys.
|
||||
- Use full-text index phrase options only when phrase queries require them.
|
||||
|
||||
## Querying
|
||||
|
||||
Always be explicit:
|
||||
|
||||
```typescript
|
||||
await table.search(queryVector).select(["id", "title"]).limit(20).toArray();
|
||||
```
|
||||
|
||||
- `select()` reduces bytes read and transferred.
|
||||
- `limit()` prevents accidental full-table collection.
|
||||
- Pre-filtering is the default behavior. Use `postfilter()` only when fewer than `limit` results are acceptable.
|
||||
|
||||
## Recall Tuning
|
||||
|
||||
Tune one knob at a time:
|
||||
|
||||
- Quantized indexes: raise `refineFactor(...)` to rescore more candidates on full vectors.
|
||||
- HNSW-backed indexes: raise `ef(...)`; start around `1.5 * k`, increase toward `10 * k` if recall is short.
|
||||
- IVF candidate breadth: `nprobes(...)` is usually auto-tuned; override only when a selective pre-filter leaves too few neighbors.
|
||||
|
||||
## Maintenance
|
||||
|
||||
After every successful embedded OSS/local ingestion, call `table.optimize()`.
|
||||
Do not add this to LanceDB Enterprise/Cloud remote table code; remote compaction
|
||||
and cleanup are handled automatically based on the Enterprise cluster
|
||||
configuration.
|
||||
|
||||
Why local maintenance is needed:
|
||||
|
||||
- Frequent writes can create many small fragments. Queries then need to scan across more files, which can increase latency.
|
||||
- Updates, deletes, and appends create new table versions. Old versions are retained for time travel and rollback, which can grow disk usage.
|
||||
- Indexes may have newly added rows that are not yet fully optimized into the index structure.
|
||||
|
||||
For local/OSS tables, run `optimize()` after the final successful ingestion
|
||||
write. Also run it after later batches of update/delete operations or on a
|
||||
regular maintenance schedule:
|
||||
|
||||
```typescript
|
||||
await table.optimize();
|
||||
```
|
||||
|
||||
If the user wants more aggressive local disk cleanup, pass a shorter cleanup retention window:
|
||||
|
||||
```typescript
|
||||
const olderThan = new Date(Date.now() - 24 * 60 * 60 * 1000);
|
||||
await table.optimize({ cleanupOlderThan: olderThan });
|
||||
```
|
||||
|
||||
Do not use very short cleanup windows when the application depends on time travel, rollback, or old versions.
|
||||
|
||||
## Diagnostics
|
||||
|
||||
Before changing code or indexes, inspect:
|
||||
|
||||
```typescript
|
||||
console.log(await table.search(queryVector).where("year > 2000").limit(10).analyzePlan());
|
||||
console.log(await table.indexStats("vector_idx"));
|
||||
```
|
||||
|
||||
Look for high scan cost, missing indexes, fragmented data, and unindexed rows.
|
||||
@@ -0,0 +1,135 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Scan Python and TypeScript for likely unsafe LanceDB materialization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PY_FULL_TABLE = re.compile(r"\b\w+\.(to_pandas|to_arrow|to_polars)\s*\(")
|
||||
TS_TABLE_TO_ARROW = re.compile(r"\b\w+\.toArrow\s*\(")
|
||||
TS_QUERY_COLLECTOR = re.compile(r"\.query\s*\(\s*\)[\s\S]*?\.to(Array|Arrow)\s*\(")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Finding:
|
||||
path: Path
|
||||
line: int
|
||||
message: str
|
||||
text: str
|
||||
|
||||
|
||||
def iter_files(paths: list[Path]) -> list[Path]:
|
||||
files: list[Path] = []
|
||||
for path in paths:
|
||||
if path.is_dir():
|
||||
files.extend(
|
||||
p
|
||||
for p in path.rglob("*")
|
||||
if p.suffix in {".py", ".ts", ".tsx"} and "node_modules" not in p.parts
|
||||
)
|
||||
elif path.suffix in {".py", ".ts", ".tsx"}:
|
||||
files.append(path)
|
||||
return sorted(set(files))
|
||||
|
||||
|
||||
def line_number(text: str, offset: int) -> int:
|
||||
return text.count("\n", 0, offset) + 1
|
||||
|
||||
|
||||
def scan_python(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in PY_FULL_TABLE.finditer(text):
|
||||
line_start = text.rfind("\n", 0, match.start()) + 1
|
||||
line_end = text.find("\n", match.start())
|
||||
if line_end == -1:
|
||||
line_end = len(text)
|
||||
line = text[line_start:line_end].strip()
|
||||
if ".search(" in line or ".query(" in line:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
f"Review Python `{match.group(1)}()` call; table-level materialization is not portable to remote tables.",
|
||||
line,
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def statement_around(text: str, start: int, end: int) -> str:
|
||||
before = max(text.rfind(";", 0, start), text.rfind("\n\n", 0, start))
|
||||
after_candidates = [pos for pos in (text.find(";", end), text.find("\n\n", end)) if pos != -1]
|
||||
after = min(after_candidates) if after_candidates else len(text)
|
||||
return text[before + 1 : after].strip()
|
||||
|
||||
|
||||
def scan_typescript(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in TS_TABLE_TO_ARROW.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".query(" in stmt or ".search(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review TypeScript `table.toArrow()`-style call; table-level materialization is not portable for large/remote tables.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in TS_QUERY_COLLECTOR.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".limit(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review unbounded TypeScript query collection; add `limit()` or stream batches.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def scan_file(path: Path) -> list[Finding]:
|
||||
text = path.read_text(encoding="utf-8", errors="replace")
|
||||
if path.suffix == ".py":
|
||||
return scan_python(path, text)
|
||||
if path.suffix in {".ts", ".tsx"}:
|
||||
return scan_typescript(path, text)
|
||||
return []
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("paths", nargs="+", type=Path)
|
||||
parser.add_argument(
|
||||
"--no-fail", action="store_true", help="Always exit 0 after reporting findings."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
findings: list[Finding] = []
|
||||
for path in iter_files(args.paths):
|
||||
findings.extend(scan_file(path))
|
||||
|
||||
for finding in findings:
|
||||
print(f"{finding.path}:{finding.line}: {finding.message}")
|
||||
print(f" {finding.text}")
|
||||
|
||||
if findings:
|
||||
print(
|
||||
f"\n{len(findings)} finding(s). Review manually; bounded query result conversion may be OK."
|
||||
)
|
||||
return 0 if args.no_fail or not findings else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.31.0-beta.4"
|
||||
current_version = "0.32.0-beta.2"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -34,15 +34,16 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
target: x86_64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
curl -L https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-$(uname -m).zip > /tmp/protoc.zip \
|
||||
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||
&& rm /tmp/protoc.zip
|
||||
curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-x86_64.zip -o /tmp/protoc.zip
|
||||
unzip /tmp/protoc.zip -d /usr/local
|
||||
rm /tmp/protoc.zip
|
||||
/usr/local/bin/protoc --version
|
||||
- name: Build Arm Manylinux Wheel
|
||||
if: ${{ inputs.arm-build == 'true' }}
|
||||
uses: PyO3/maturin-action@v1
|
||||
@@ -50,13 +51,14 @@ runs:
|
||||
maturin-version: "1.12.4"
|
||||
command: build
|
||||
working-directory: python
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/' -e PROTOC=/usr/local/bin/protoc"
|
||||
target: aarch64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
yum install -y clang \
|
||||
&& curl -L https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-aarch_64.zip > /tmp/protoc.zip \
|
||||
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||
&& rm /tmp/protoc.zip
|
||||
yum install -y clang
|
||||
curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-aarch_64.zip -o /tmp/protoc.zip
|
||||
unzip /tmp/protoc.zip -d /usr/local
|
||||
rm /tmp/protoc.zip
|
||||
/usr/local/bin/protoc --version
|
||||
|
||||
@@ -25,7 +25,7 @@ jobs:
|
||||
# Only runs on tags that matches the make-release action
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
workspaces: rust
|
||||
@@ -47,7 +47,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,14 +36,14 @@ jobs:
|
||||
echo "guidelines = ${{ inputs.guidelines }}"
|
||||
|
||||
- name: Checkout Repo
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.branch }}
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 (used by the nodejs install step below) requires
|
||||
# Node >= 22.13; use 24 since 22 hits EOL in October.
|
||||
@@ -82,7 +82,7 @@ jobs:
|
||||
cache: maven
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Install Node.js dependencies for TypeScript bindings
|
||||
|
||||
@@ -30,13 +30,13 @@ jobs:
|
||||
echo "tag = ${{ inputs.tag || 'latest' }}"
|
||||
|
||||
- name: Checkout Repo LanceDB
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
name: Verify PR title / description conforms to semantic-release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: "18"
|
||||
# These rules are disabled because Github will always ensure there
|
||||
|
||||
@@ -35,7 +35,7 @@ jobs:
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install dependencies needed for ubuntu
|
||||
run: |
|
||||
sudo apt install -y protobuf-compiler libssl-dev
|
||||
@@ -53,7 +53,7 @@ jobs:
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -e .
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -r ../docs/requirements.txt
|
||||
- name: Set up node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
cache: 'npm'
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Set up Java 8
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
@@ -73,7 +73,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,7 +36,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Set up Java 17
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
|
||||
@@ -19,7 +19,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
- name: Install license-header-checker
|
||||
working-directory: /tmp
|
||||
run: |
|
||||
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
steps:
|
||||
- name: Output Inputs
|
||||
run: echo "${{ toJSON(github.event.inputs) }}"
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
@@ -38,14 +38,14 @@ jobs:
|
||||
CC: gcc-12
|
||||
CXX: g++-12
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October. The library itself still supports Node >= 18
|
||||
@@ -86,14 +86,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
name: Setup Node.js 24 for build
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
@@ -130,7 +130,7 @@ jobs:
|
||||
echo "Run 'pnpm run docs', fix any warnings, and commit the changes."
|
||||
exit 1
|
||||
fi
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
name: Setup Node.js ${{ matrix.node-version }} for test
|
||||
with:
|
||||
node-version: ${{ matrix.node-version }}
|
||||
@@ -166,14 +166,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v4
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -103,7 +103,7 @@ jobs:
|
||||
features: fp16kernels
|
||||
pre_build: brew install protobuf
|
||||
- target: x86_64-pc-windows-msvc
|
||||
host: windows-latest
|
||||
host: windows-2025-8x-x64
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
@@ -111,12 +111,21 @@ jobs:
|
||||
# There is an issue where choco doesn't add nasm to the path
|
||||
export PATH="$PATH:/c/Program Files/NASM"
|
||||
nasm -v
|
||||
# Fat LTO of the cdylib is single-threaded and the peak-memory
|
||||
# step of the build, and had started hitting rustc-LLVM OOM on the
|
||||
# Windows runners. ThinLTO parallelizes it across the runner's
|
||||
# cores and keeps peak memory well under the limit.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: aarch64-pc-windows-msvc
|
||||
host: windows-latest
|
||||
host: windows-2025-8x-x64
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc
|
||||
rustup target add aarch64-pc-windows-msvc
|
||||
# See ThinLTO note on the x86_64-pc-windows-msvc target above.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-unknown-linux-gnu
|
||||
host: ubuntu-latest
|
||||
features: fp16kernels
|
||||
@@ -170,13 +179,13 @@ jobs:
|
||||
run:
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -190,7 +199,7 @@ jobs:
|
||||
toolchain: stable
|
||||
targets: ${{ matrix.settings.target }}
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v4
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry/index/
|
||||
@@ -244,7 +253,7 @@ jobs:
|
||||
if: ${{ !matrix.settings.docker }}
|
||||
shell: bash
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/*.node
|
||||
@@ -256,7 +265,7 @@ jobs:
|
||||
run: pnpm tsc
|
||||
- name: Upload Generic Artifacts
|
||||
if: ${{ matrix.settings.target == 'aarch64-apple-darwin' }}
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: |
|
||||
@@ -287,13 +296,13 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup Node.js 24 for install
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -303,18 +312,18 @@ jobs:
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Setup Node.js ${{ matrix.node }} for test
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: ${{ matrix.node }}
|
||||
- name: Download artifacts
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
@@ -339,13 +348,13 @@ jobs:
|
||||
needs:
|
||||
- test-lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
@@ -353,14 +362,14 @@ jobs:
|
||||
registry-url: "https://registry.npmjs.org"
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v4
|
||||
- uses: actions/download-artifact@v8
|
||||
name: Download arch-specific binaries
|
||||
with:
|
||||
pattern: lancedb-*
|
||||
@@ -398,7 +407,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -41,7 +41,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -66,7 +66,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -95,7 +95,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -160,7 +160,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -189,7 +189,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -212,7 +212,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
+27
-11
@@ -40,7 +40,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -65,7 +65,7 @@ jobs:
|
||||
timeout-minutes: 10
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: EmbarkStudios/cargo-deny-action@v2
|
||||
with:
|
||||
command: check advisories bans licenses sources
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
CC: clang
|
||||
CXX: clang++
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
# Building without a lock file often requires the latest Rust version since downstream
|
||||
# dependencies may have updated their minimum Rust version.
|
||||
- uses: actions-rust-lang/setup-rust-toolchain@v1
|
||||
@@ -113,7 +113,7 @@ jobs:
|
||||
CXX: clang++-18
|
||||
GH_TOKEN: ${{ secrets.SOPHON_READ_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -125,10 +125,26 @@ jobs:
|
||||
- uses: rui314/setup-mold@v1
|
||||
- name: Make Swap
|
||||
run: |
|
||||
sudo fallocate -l 16G /swapfile
|
||||
sudo chmod 600 /swapfile
|
||||
sudo mkswap /swapfile
|
||||
sudo swapon /swapfile
|
||||
swapfile=/swapfile
|
||||
min_swap_bytes=$((15 * 1024 * 1024 * 1024))
|
||||
active_swap_bytes="$(sudo swapon --show=NAME,SIZE --bytes --noheadings | awk '$1 == "/swapfile" { print $2 }')"
|
||||
if [ -n "$active_swap_bytes" ]; then
|
||||
if [ "$active_swap_bytes" -ge "$min_swap_bytes" ]; then
|
||||
echo "/swapfile is already active with enough space; skipping swap creation"
|
||||
exit 0
|
||||
fi
|
||||
echo "/swapfile is already active but smaller than 16G; using /mnt/lancedb-swapfile"
|
||||
swapfile=/mnt/lancedb-swapfile
|
||||
fi
|
||||
if sudo swapon --show=NAME --noheadings | grep -Fxq "$swapfile"; then
|
||||
echo "$swapfile is already active; skipping swap creation"
|
||||
exit 0
|
||||
fi
|
||||
sudo rm -f "$swapfile"
|
||||
sudo fallocate -l 16G "$swapfile"
|
||||
sudo chmod 600 "$swapfile"
|
||||
sudo mkswap "$swapfile"
|
||||
sudo swapon "$swapfile"
|
||||
- name: Build
|
||||
run: cargo build --profile ci --all-features --tests --locked --examples
|
||||
- name: Run feature tests
|
||||
@@ -152,7 +168,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: rust
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -181,7 +197,7 @@ jobs:
|
||||
run:
|
||||
working-directory: rust/lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
- name: Set target
|
||||
run: rustup target add ${{ matrix.target }}
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -210,7 +226,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
submodules: true
|
||||
- name: Install dependencies
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v4
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
Generated
+382
-122
File diff suppressed because it is too large
Load Diff
+17
-14
@@ -13,24 +13,25 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.0.0-beta.10", default-features = false, "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.10", default-features = false, "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.10", default-features = false, "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.10", "tag" = "v9.0.0-beta.10", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
arrow-array = "58.0.0"
|
||||
arrow-buffer = "58.0.0"
|
||||
arrow-data = "58.0.0"
|
||||
arrow-ipc = "58.0.0"
|
||||
arrow-ord = "58.0.0"
|
||||
@@ -53,6 +54,8 @@ half = { "version" = "2.7.1", default-features = false, features = [
|
||||
] }
|
||||
futures = "0"
|
||||
log = "0.4"
|
||||
metrics = "0.24"
|
||||
metrics-util = "0.19"
|
||||
moka = { version = "0.12", features = ["future"] }
|
||||
object_store = "0.13.2"
|
||||
pin-project = "1.0.7"
|
||||
|
||||
@@ -51,18 +51,6 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0436
|
||||
{ id = "RUSTSEC-2024-0436", reason = "transitive via datafusion; awaiting ecosystem migration" },
|
||||
|
||||
# encoding: unmaintained. Reached through lindera-dictionary, which is
|
||||
# required by the native Lindera tokenizer path. Lindera has not migrated
|
||||
# off this crate yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2021-0153
|
||||
{ id = "RUSTSEC-2021-0153", reason = "transitive via lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# fast-float: unsound and unmaintained. Reached only through polars-arrow
|
||||
# from the optional Polars integration; replacement requires a Polars
|
||||
# dependency upgrade.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0379
|
||||
{ id = "RUSTSEC-2024-0379", reason = "transitive via polars-arrow; waiting on Polars migration" },
|
||||
|
||||
# tantivy: segfault on malformed input due to missing bounds check.
|
||||
# Pulled in via lance for full-text search. We only feed tantivy
|
||||
# documents we construct ourselves, not attacker-controlled bytes.
|
||||
@@ -80,18 +68,6 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0119
|
||||
{ id = "RUSTSEC-2025-0119", reason = "transitive via hf-hub/indicatif; cosmetic formatting crate" },
|
||||
|
||||
# bincode: unmaintained. Reached through lindera and lindera-dictionary,
|
||||
# which are required by the native Lindera tokenizer path. Lindera has not
|
||||
# migrated to another serialization format yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0141
|
||||
{ id = "RUSTSEC-2025-0141", reason = "transitive via lindera/lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# lru: soundness issue in IterMut. Reached only through aws-sdk-s3 in
|
||||
# LanceDB's dev-dependency graph; LanceDB does not use that iterator
|
||||
# directly. Clearing this requires the AWS SDK chain to update lru.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0002
|
||||
{ id = "RUSTSEC-2026-0002", reason = "transitive via aws-sdk-s3 dev-dependency; waiting on AWS SDK lru upgrade" },
|
||||
|
||||
# rustls-webpki 0.101.7 (old major line): name-constraint checks for
|
||||
# URI / wildcard names. Pulled in only via the legacy rustls 0.21 chain
|
||||
# from aws-smithy-http-client. The 0.103 line we actively use is patched.
|
||||
@@ -108,17 +84,23 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0104
|
||||
{ id = "RUSTSEC-2026-0104", reason = "only affects rustls-webpki 0.101 from legacy aws-smithy/rustls 0.21 chain" },
|
||||
|
||||
# rand 0.8.5: soundness issue only when ThreadRng reseeds inside a custom
|
||||
# logger. Reached through several transitive chains. LanceDB does not use
|
||||
# rand from a custom logger; upgrade once all pinned chains accept 0.8.6+.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0097
|
||||
{ id = "RUSTSEC-2026-0097", reason = "transitive rand 0.8.5; LanceDB does not call ThreadRng from custom logging" },
|
||||
|
||||
# pyo3 advisories in the Python bindings; tracked pending a patched pyo3 release.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0176
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0177
|
||||
{ id = "RUSTSEC-2026-0176", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
{ id = "RUSTSEC-2026-0177", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
|
||||
# quick-xml < 0.41.0: quadratic runtime on duplicate attribute names (DoS).
|
||||
# quick-xml < 0.41.0: unbounded namespace-declaration allocation in NsReader (DoS).
|
||||
# Pulled in transitively by inferno (dev-only flame-graph dep), lance-namespace-impls
|
||||
# (git dep from lance), and opendal/reqsign (cloud storage XML parsing). The XML
|
||||
# parsed by opendal/reqsign comes from trusted cloud-storage endpoints (S3, GCS,
|
||||
# Azure), not attacker-controlled input. Clearing requires upstream crates to migrate
|
||||
# to quick-xml >= 0.41.0.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0194
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0195
|
||||
{ id = "RUSTSEC-2026-0194", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
{ id = "RUSTSEC-2026-0195", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.31.0-beta.4</version>
|
||||
<version>0.32.0-beta.2</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -518,6 +518,9 @@ x > 5 OR y = 'test'
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND rather
|
||||
than replacing the previous filter.
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
@@ -398,6 +398,26 @@ Drop an index from the table.
|
||||
|
||||
***
|
||||
|
||||
### getLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
abstract getLsmWriteSpec(): Promise<undefined | LsmWriteSpec>
|
||||
```
|
||||
|
||||
Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) currently installed on this table.
|
||||
|
||||
Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
spec has been set, or it was removed with [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)).
|
||||
The returned spec — including its `maintainedIndexes` and
|
||||
`writerConfigDefaults` — mirrors what was passed to
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec).
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`undefined` \| [`LsmWriteSpec`](../interfaces/LsmWriteSpec.md)>
|
||||
|
||||
***
|
||||
|
||||
### indexStats()
|
||||
|
||||
```ts
|
||||
@@ -914,6 +934,32 @@ Return the table as an arrow table
|
||||
|
||||
***
|
||||
|
||||
### tokenize()
|
||||
|
||||
```ts
|
||||
abstract tokenize(query, options): Promise<FtsToken[]>
|
||||
```
|
||||
|
||||
Tokenize a full-text search query using the tokenizer configured on an FTS index.
|
||||
|
||||
Specify exactly one of `column` or `indexName`.
|
||||
|
||||
Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
the client process from index metadata. For remote tables, this means the
|
||||
same tokenizer model files must also exist locally.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **query**: `string`
|
||||
|
||||
* **options**: [`TokenizeTableOptions`](../type-aliases/TokenizeTableOptions.md)
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]>
|
||||
|
||||
***
|
||||
|
||||
### unsetLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -767,6 +767,9 @@ x > 5 OR y = 'test'
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND rather
|
||||
than replacing the previous filter.
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / instrumentLanceDbMetrics
|
||||
|
||||
# Function: instrumentLanceDbMetrics()
|
||||
|
||||
```ts
|
||||
function instrumentLanceDbMetrics(meterProvider?): boolean
|
||||
```
|
||||
|
||||
Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
|
||||
Installs a process-global metrics recorder and creates one observable
|
||||
instrument per LanceDB metric (currently object store request counts, bytes,
|
||||
latency, errors, and throttles) on the given (or global) `MeterProvider`. The
|
||||
configured `MetricReader` then collects them on its own schedule.
|
||||
|
||||
Counters and gauges map directly to observable counters/gauges. Because
|
||||
OpenTelemetry has no asynchronous histogram instrument, each histogram is
|
||||
exported Prometheus-style as cumulative `le` bucket counts (`<name>_bucket`,
|
||||
with an `le` attribute) plus `<name>_count` and `<name>_sum`.
|
||||
|
||||
Requires `@opentelemetry/api` (a dependency) and, to actually export, an
|
||||
OpenTelemetry SDK such as `@opentelemetry/sdk-metrics`.
|
||||
|
||||
## Parameters
|
||||
|
||||
* **meterProvider?**: `MeterProvider`
|
||||
The provider to register instruments on. Defaults to the
|
||||
global provider from `@opentelemetry/api`.
|
||||
|
||||
## Returns
|
||||
|
||||
`boolean`
|
||||
|
||||
`true` if the recorder is installed and instruments are registered.
|
||||
`false` if a different `metrics` recorder is already installed in this
|
||||
process (only one global recorder is permitted), in which case a warning is
|
||||
emitted and no instruments are created. Calling this more than once is safe;
|
||||
instruments are created only on the first successful call.
|
||||
@@ -0,0 +1,26 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / tokenize
|
||||
|
||||
# Function: tokenize()
|
||||
|
||||
```ts
|
||||
function tokenize(query, options?): Promise<FtsToken[]>
|
||||
```
|
||||
|
||||
Tokenize a full-text search query using an explicit tokenizer.
|
||||
|
||||
This does not require a table or FTS index. The tokenizer options match
|
||||
[Index.fts](../classes/Index.md#fts).
|
||||
|
||||
## Parameters
|
||||
|
||||
* **query**: `string`
|
||||
|
||||
* **options?**: `Partial`<[`TokenizeOptions`](../interfaces/TokenizeOptions.md)>
|
||||
|
||||
## Returns
|
||||
|
||||
`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]>
|
||||
@@ -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,7 @@
|
||||
|
||||
## Type Aliases
|
||||
|
||||
- [BaseTokenizer](type-aliases/BaseTokenizer.md)
|
||||
- [Data](type-aliases/Data.md)
|
||||
- [DataLike](type-aliases/DataLike.md)
|
||||
- [FieldLike](type-aliases/FieldLike.md)
|
||||
@@ -125,12 +128,15 @@
|
||||
- [RecordBatchLike](type-aliases/RecordBatchLike.md)
|
||||
- [SchemaLike](type-aliases/SchemaLike.md)
|
||||
- [TableLike](type-aliases/TableLike.md)
|
||||
- [TokenizeTableOptions](type-aliases/TokenizeTableOptions.md)
|
||||
|
||||
## Functions
|
||||
|
||||
- [RecordBatchIterator](functions/RecordBatchIterator.md)
|
||||
- [connect](functions/connect.md)
|
||||
- [connectNamespace](functions/connectNamespace.md)
|
||||
- [instrumentLanceDbMetrics](functions/instrumentLanceDbMetrics.md)
|
||||
- [makeArrowTable](functions/makeArrowTable.md)
|
||||
- [packBits](functions/packBits.md)
|
||||
- [permutationBuilder](functions/permutationBuilder.md)
|
||||
- [tokenize](functions/tokenize.md)
|
||||
|
||||
@@ -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.
|
||||
@@ -8,6 +8,14 @@
|
||||
|
||||
## Properties
|
||||
|
||||
### clumpSize?
|
||||
|
||||
```ts
|
||||
optional clumpSize: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### counts?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenizeOptions
|
||||
|
||||
# Interface: TokenizeOptions
|
||||
|
||||
Options for tokenizing a full-text search query without a table index.
|
||||
|
||||
## Properties
|
||||
|
||||
### asciiFolding?
|
||||
|
||||
```ts
|
||||
optional asciiFolding: boolean;
|
||||
```
|
||||
|
||||
Whether to fold ASCII characters.
|
||||
|
||||
***
|
||||
|
||||
### baseTokenizer?
|
||||
|
||||
```ts
|
||||
optional baseTokenizer: BaseTokenizer;
|
||||
```
|
||||
|
||||
The tokenizer to use. The default is "simple".
|
||||
|
||||
***
|
||||
|
||||
### language?
|
||||
|
||||
```ts
|
||||
optional language: string;
|
||||
```
|
||||
|
||||
Language for stemming and stop words.
|
||||
|
||||
***
|
||||
|
||||
### lowercase?
|
||||
|
||||
```ts
|
||||
optional lowercase: boolean;
|
||||
```
|
||||
|
||||
Whether to lowercase tokens.
|
||||
|
||||
***
|
||||
|
||||
### maxTokenLength?
|
||||
|
||||
```ts
|
||||
optional maxTokenLength: number;
|
||||
```
|
||||
|
||||
Maximum token length; tokens longer than this are ignored.
|
||||
|
||||
***
|
||||
|
||||
### ngramMaxLength?
|
||||
|
||||
```ts
|
||||
optional ngramMaxLength: number;
|
||||
```
|
||||
|
||||
N-gram maximum length.
|
||||
|
||||
***
|
||||
|
||||
### ngramMinLength?
|
||||
|
||||
```ts
|
||||
optional ngramMinLength: number;
|
||||
```
|
||||
|
||||
N-gram minimum length.
|
||||
|
||||
***
|
||||
|
||||
### prefixOnly?
|
||||
|
||||
```ts
|
||||
optional prefixOnly: boolean;
|
||||
```
|
||||
|
||||
Whether to only emit token prefixes for the n-gram tokenizer.
|
||||
|
||||
***
|
||||
|
||||
### removeStopWords?
|
||||
|
||||
```ts
|
||||
optional removeStopWords: boolean;
|
||||
```
|
||||
|
||||
Whether to remove stop words.
|
||||
|
||||
***
|
||||
|
||||
### stem?
|
||||
|
||||
```ts
|
||||
optional stem: boolean;
|
||||
```
|
||||
|
||||
Whether to stem tokens.
|
||||
@@ -0,0 +1,19 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / BaseTokenizer
|
||||
|
||||
# Type Alias: BaseTokenizer
|
||||
|
||||
```ts
|
||||
type BaseTokenizer:
|
||||
| "simple"
|
||||
| "whitespace"
|
||||
| "raw"
|
||||
| "ngram"
|
||||
| "icu"
|
||||
| "icu/split"
|
||||
| `jieba/${string}`
|
||||
| `lindera/${string}`;
|
||||
```
|
||||
@@ -0,0 +1,11 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenizeTableOptions
|
||||
|
||||
# Type Alias: TokenizeTableOptions
|
||||
|
||||
```ts
|
||||
type TokenizeTableOptions: object | object;
|
||||
```
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.31.0-beta.4</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.31.0-beta.4</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.10</lance-core.version>
|
||||
<lance-core.version>9.0.0-beta.23</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.31.0-beta.4"
|
||||
version = "0.32.0-beta.2"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
@@ -44,6 +44,6 @@ aws-lc-rs = "=1.16.3"
|
||||
napi-build = "2.3.1"
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
MeterProvider,
|
||||
type MetricData,
|
||||
MetricReader,
|
||||
} from "@opentelemetry/sdk-metrics";
|
||||
import * as tmp from "tmp";
|
||||
import { connect, instrumentLanceDbMetrics } from "../lancedb";
|
||||
// snapshotLancedbMetrics is internal plumbing (not part of the public API), so
|
||||
// it is imported from the native module rather than the package entry point.
|
||||
import { snapshotLancedbMetrics } from "../lancedb/native";
|
||||
|
||||
// The metrics recorder is process-global and installed once, so the whole
|
||||
// bridge is exercised in a single test to avoid cross-test global-state coupling.
|
||||
|
||||
// A minimal pull-based reader whose `collect()` we drive directly, invoking the
|
||||
// observable-instrument callbacks. `@opentelemetry/sdk-metrics` ships no
|
||||
// in-memory reader, so we subclass the abstract base.
|
||||
class TestMetricReader extends MetricReader {
|
||||
protected async onForceFlush(): Promise<void> {
|
||||
// no-op: collection is driven directly via collect()
|
||||
}
|
||||
protected async onShutdown(): Promise<void> {
|
||||
// no-op: nothing to release
|
||||
}
|
||||
}
|
||||
|
||||
async function metricsByName(
|
||||
reader: TestMetricReader,
|
||||
): Promise<Map<string, MetricData>> {
|
||||
const collected = await reader.collect();
|
||||
const result = new Map<string, MetricData>();
|
||||
for (const scope of collected.resourceMetrics.scopeMetrics) {
|
||||
for (const metric of scope.metrics) {
|
||||
result.set(metric.descriptor.name, metric);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
describe("OpenTelemetry metrics bridge", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
afterEach(() => tmpDir.removeCallback());
|
||||
|
||||
it("snapshot is safe to call regardless of install state", () => {
|
||||
expect(Array.isArray(snapshotLancedbMetrics())).toBe(true);
|
||||
});
|
||||
|
||||
it("exports object store metrics via observable instruments", async () => {
|
||||
const reader = new TestMetricReader();
|
||||
const provider = new MeterProvider({ readers: [reader] });
|
||||
expect(instrumentLanceDbMetrics(provider)).toBe(true);
|
||||
|
||||
// Generate object store activity on the local filesystem (scheme "file").
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = Array.from({ length: 256 }, (_, i) => ({ id: i }));
|
||||
const table = await db.createTable("t", data);
|
||||
expect(await table.countRows()).toBe(256);
|
||||
|
||||
const metrics = await metricsByName(reader);
|
||||
|
||||
const requests = metrics.get("lance_object_store_requests_total");
|
||||
expect(requests).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const requestPoints = (requests!.dataPoints as any[]) ?? [];
|
||||
expect(requestPoints.length).toBeGreaterThan(0);
|
||||
for (const p of requestPoints) {
|
||||
// Labelled by `operation` and `base` (the store scheme by default).
|
||||
expect(p.attributes).toHaveProperty("base");
|
||||
expect(p.attributes).toHaveProperty("operation");
|
||||
}
|
||||
const totalRequests = requestPoints.reduce((acc, p) => acc + p.value, 0);
|
||||
expect(totalRequests).toBeGreaterThan(0);
|
||||
|
||||
// Histograms are decomposed into bucket / count / sum observable counters.
|
||||
const bucket = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_bucket",
|
||||
);
|
||||
expect(bucket).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const bucketPoints = (bucket!.dataPoints as any[]) ?? [];
|
||||
expect(bucketPoints.length).toBeGreaterThan(0);
|
||||
expect(bucketPoints.every((p) => "le" in p.attributes)).toBe(true);
|
||||
// The implicit +Inf bucket must be present.
|
||||
expect(bucketPoints.some((p) => p.attributes.le === "+Inf")).toBe(true);
|
||||
|
||||
const count = metrics.get(
|
||||
"lance_object_store_request_duration_seconds_count",
|
||||
);
|
||||
expect(count).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const countPoints = (count!.dataPoints as any[]) ?? [];
|
||||
expect(countPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
const sum = metrics.get("lance_object_store_request_duration_seconds_sum");
|
||||
expect(sum).toBeDefined();
|
||||
// biome-ignore lint/suspicious/noExplicitAny: SDK point shape
|
||||
const sumPoints = (sum!.dataPoints as any[]) ?? [];
|
||||
expect(sumPoints.reduce((acc, p) => acc + p.value, 0)).toBeGreaterThan(0);
|
||||
|
||||
// Unit handling: only `_sum` keeps the histogram's unit (seconds); `_bucket`
|
||||
// and `_count` observe cumulative counts and are unitless.
|
||||
expect(sum!.descriptor.unit).toBe("s");
|
||||
expect(bucket!.descriptor.unit).toBe("");
|
||||
expect(count!.descriptor.unit).toBe("");
|
||||
|
||||
await provider.shutdown();
|
||||
});
|
||||
});
|
||||
@@ -215,6 +215,20 @@ describe("Query orderBy", () => {
|
||||
expect(results[2].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should combine repeated where clauses with AND", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 1.0")
|
||||
.where("score < 3.0")
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
// Only rows matching both predicates should be returned, rather than the
|
||||
// second where() silently replacing the first.
|
||||
expect(results.length).toBe(2);
|
||||
expect(results[0].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(2.8, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with limit", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
|
||||
@@ -16,6 +16,7 @@ import {
|
||||
PhraseQuery,
|
||||
Table,
|
||||
connect,
|
||||
tokenize,
|
||||
} from "../lancedb";
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
@@ -2307,6 +2308,75 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
expect(results2[0].text).toBe(data[1].text);
|
||||
});
|
||||
|
||||
test("tokenizes FTS queries by column or index name", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
{
|
||||
text: "Running in cafés",
|
||||
japanese: "Hello, こんにちは世界!",
|
||||
vector: [0.1, 0.2, 0.3],
|
||||
},
|
||||
];
|
||||
const table = await db.createTable("test", data);
|
||||
await table.createIndex("text", {
|
||||
config: Index.fts({ baseTokenizer: "simple" }),
|
||||
});
|
||||
await table.createIndex("japanese", {
|
||||
config: Index.fts({
|
||||
baseTokenizer: "icu",
|
||||
stem: false,
|
||||
removeStopWords: false,
|
||||
}),
|
||||
name: "japanese_icu_idx",
|
||||
});
|
||||
|
||||
await expect(table.tokenize("hello", {} as never)).rejects.toThrow(
|
||||
"Specify exactly one",
|
||||
);
|
||||
await expect(
|
||||
table.tokenize("hello", {
|
||||
column: "text",
|
||||
indexName: "text_idx",
|
||||
} as never),
|
||||
).rejects.toThrow("Specify exactly one");
|
||||
|
||||
const simpleTokens = await table.tokenize("Running in cafés", {
|
||||
column: "text",
|
||||
});
|
||||
expect(simpleTokens).toEqual([
|
||||
{ text: "run", position: 0 },
|
||||
{ text: "cafe", position: 2 },
|
||||
]);
|
||||
|
||||
const icuTokens = await table.tokenize("Hello, こんにちは世界!", {
|
||||
indexName: "japanese_icu_idx",
|
||||
});
|
||||
expect(icuTokens).toEqual([
|
||||
{ text: "hello", position: 0 },
|
||||
{ text: "こんにちは", position: 1 },
|
||||
{ text: "世界", position: 2 },
|
||||
]);
|
||||
|
||||
const directSimpleTokens = await tokenize("Running in cafés", {
|
||||
baseTokenizer: "simple",
|
||||
});
|
||||
expect(directSimpleTokens).toEqual([
|
||||
{ text: "run", position: 0 },
|
||||
{ text: "cafe", position: 2 },
|
||||
]);
|
||||
|
||||
const directIcuTokens = await tokenize("Hello, こんにちは世界!", {
|
||||
baseTokenizer: "icu",
|
||||
stem: false,
|
||||
removeStopWords: false,
|
||||
});
|
||||
expect(directIcuTokens).toEqual([
|
||||
{ text: "hello", position: 0 },
|
||||
{ text: "こんにちは", position: 1 },
|
||||
{ text: "世界", position: 2 },
|
||||
]);
|
||||
});
|
||||
|
||||
test("full text search fast search", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [{ text: "hello world", vector: [0.1, 0.2, 0.3], id: 1 }];
|
||||
@@ -2992,6 +3062,56 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => {
|
||||
}),
|
||||
).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("reads back the installed spec via getLsmWriteSpec", async () => {
|
||||
const conn = await connect(tmpDir.name);
|
||||
const table = await makeTable(conn);
|
||||
await table.setUnenforcedPrimaryKey("id");
|
||||
|
||||
// Nothing installed yet.
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// A real scalar index is needed to name it as a maintained index.
|
||||
await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]);
|
||||
await table.createIndex("id");
|
||||
const indexName = (await table.listIndices())[0].name;
|
||||
|
||||
// Bucket spec round-trips, including maintained indexes and writer config
|
||||
// defaults. Lance writer-config keys are canonically snake_case.
|
||||
// biome-ignore lint/style/useNamingConvention: Lance writer-config keys are snake_case
|
||||
const writerConfigDefaults = { durable_write: "false" };
|
||||
await table.setLsmWriteSpec({
|
||||
specType: "bucket",
|
||||
column: "id",
|
||||
numBuckets: 4,
|
||||
maintainedIndexes: [indexName],
|
||||
writerConfigDefaults,
|
||||
});
|
||||
const spec = await table.getLsmWriteSpec();
|
||||
expect(spec).toBeDefined();
|
||||
expect(spec?.specType).toBe("bucket");
|
||||
expect(spec?.column).toBe("id");
|
||||
expect(spec?.numBuckets).toBe(4);
|
||||
expect(spec?.maintainedIndexes).toEqual([indexName]);
|
||||
expect(spec?.writerConfigDefaults).toEqual(writerConfigDefaults);
|
||||
|
||||
// After unset, undefined again.
|
||||
await table.unsetLsmWriteSpec();
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// Identity round-trips (column recovered from the schema).
|
||||
await table.setLsmWriteSpec({ specType: "identity", column: "id" });
|
||||
const identity = await table.getLsmWriteSpec();
|
||||
expect(identity?.specType).toBe("identity");
|
||||
expect(identity?.column).toBe("id");
|
||||
await table.unsetLsmWriteSpec();
|
||||
|
||||
// Unsharded round-trips (no routing column).
|
||||
await table.setLsmWriteSpec({ specType: "unsharded" });
|
||||
const unsharded = await table.getLsmWriteSpec();
|
||||
expect(unsharded?.specType).toBe("unsharded");
|
||||
expect(unsharded?.column).toBeFalsy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("LSM merge insert", () => {
|
||||
|
||||
@@ -13,13 +13,21 @@ import {
|
||||
Connection as LanceDbConnection,
|
||||
JsHeaderProvider as NativeJsHeaderProvider,
|
||||
Session,
|
||||
tokenize as nativeTokenize,
|
||||
} from "./native.js";
|
||||
|
||||
import { HeaderProvider } from "./header";
|
||||
import type { BaseTokenizer } from "./indices";
|
||||
import type { FtsToken } from "./table";
|
||||
|
||||
// Re-export native header provider for use with connectWithHeaderProvider
|
||||
export { JsHeaderProvider as NativeJsHeaderProvider } from "./native.js";
|
||||
|
||||
// OpenTelemetry metrics bridge. Only the high-level entry point is public; the
|
||||
// underlying recorder/catalog/snapshot functions remain internal plumbing that
|
||||
// `otel.ts` consumes from the native module.
|
||||
export { instrumentLanceDbMetrics } from "./otel";
|
||||
|
||||
export {
|
||||
AddColumnsSql,
|
||||
ConnectionOptions,
|
||||
@@ -109,6 +117,7 @@ export {
|
||||
HnswPqOptions,
|
||||
HnswSqOptions,
|
||||
FtsOptions,
|
||||
BaseTokenizer,
|
||||
} from "./indices";
|
||||
|
||||
export {
|
||||
@@ -119,6 +128,8 @@ export {
|
||||
OptimizeOptions,
|
||||
Version,
|
||||
WriteProgress,
|
||||
FtsToken,
|
||||
TokenizeTableOptions,
|
||||
LsmWriteSpec,
|
||||
ColumnAlteration,
|
||||
FieldMetadataUpdate,
|
||||
@@ -150,6 +161,68 @@ export {
|
||||
} from "./arrow";
|
||||
export { IntoSql, packBits } from "./util";
|
||||
|
||||
/**
|
||||
* Options for tokenizing a full-text search query without a table index.
|
||||
*/
|
||||
export interface TokenizeOptions {
|
||||
/**
|
||||
* The tokenizer to use. The default is "simple".
|
||||
*/
|
||||
baseTokenizer?: BaseTokenizer;
|
||||
|
||||
/** Language for stemming and stop words. */
|
||||
language?: string;
|
||||
|
||||
/** Maximum token length; tokens longer than this are ignored. */
|
||||
maxTokenLength?: number;
|
||||
|
||||
/** Whether to lowercase tokens. */
|
||||
lowercase?: boolean;
|
||||
|
||||
/** Whether to stem tokens. */
|
||||
stem?: boolean;
|
||||
|
||||
/** Whether to remove stop words. */
|
||||
removeStopWords?: boolean;
|
||||
|
||||
/** Whether to fold ASCII characters. */
|
||||
asciiFolding?: boolean;
|
||||
|
||||
/** N-gram minimum length. */
|
||||
ngramMinLength?: number;
|
||||
|
||||
/** N-gram maximum length. */
|
||||
ngramMaxLength?: number;
|
||||
|
||||
/** Whether to only emit token prefixes for the n-gram tokenizer. */
|
||||
prefixOnly?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize a full-text search query using an explicit tokenizer.
|
||||
*
|
||||
* This does not require a table or FTS index. The tokenizer options match
|
||||
* {@link Index.fts}.
|
||||
*/
|
||||
export async function tokenize(
|
||||
query: string,
|
||||
options?: Partial<TokenizeOptions>,
|
||||
): Promise<FtsToken[]> {
|
||||
return await nativeTokenize(
|
||||
query,
|
||||
options?.baseTokenizer,
|
||||
options?.language,
|
||||
options?.maxTokenLength,
|
||||
options?.lowercase,
|
||||
options?.stem,
|
||||
options?.removeStopWords,
|
||||
options?.asciiFolding,
|
||||
options?.ngramMinLength,
|
||||
options?.ngramMaxLength,
|
||||
options?.prefixOnly,
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Connect to a LanceDB instance at the given URI.
|
||||
*
|
||||
|
||||
@@ -486,6 +486,16 @@ export interface IvfFlatOptions {
|
||||
sampleRate?: number;
|
||||
}
|
||||
|
||||
export type BaseTokenizer =
|
||||
| "simple"
|
||||
| "whitespace"
|
||||
| "raw"
|
||||
| "ngram"
|
||||
| "icu"
|
||||
| "icu/split"
|
||||
| `jieba/${string}`
|
||||
| `lindera/${string}`;
|
||||
|
||||
/**
|
||||
* Options to create a full text search index
|
||||
*/
|
||||
@@ -509,8 +519,12 @@ export interface FtsOptions {
|
||||
* "whitespace" - Whitespace tokenizer. This tokenizer splits the text into tokens using whitespace as a delimiter.
|
||||
*
|
||||
* "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token.
|
||||
*
|
||||
* "icu" - ICU dictionary-based word segmentation.
|
||||
*
|
||||
* "icu/split" - ICU segmentation with simple-style delimiter splitting.
|
||||
*/
|
||||
baseTokenizer?: "simple" | "whitespace" | "raw" | "ngram";
|
||||
baseTokenizer?: BaseTokenizer;
|
||||
|
||||
/**
|
||||
* language for stemming and stop words
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
type Attributes,
|
||||
type MeterProvider,
|
||||
type ObservableResult,
|
||||
metrics,
|
||||
} from "@opentelemetry/api";
|
||||
|
||||
import {
|
||||
lancedbMetricsCatalog,
|
||||
registerLancedbMetricsRecorder,
|
||||
snapshotLancedbMetrics,
|
||||
} from "./native";
|
||||
|
||||
let instrumented = false;
|
||||
|
||||
/**
|
||||
* Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
*
|
||||
* Installs a process-global metrics recorder and creates one observable
|
||||
* instrument per LanceDB metric (currently object store request counts, bytes,
|
||||
* latency, errors, and throttles) on the given (or global) `MeterProvider`. The
|
||||
* configured `MetricReader` then collects them on its own schedule.
|
||||
*
|
||||
* Counters and gauges map directly to observable counters/gauges. Because
|
||||
* OpenTelemetry has no asynchronous histogram instrument, each histogram is
|
||||
* exported Prometheus-style as cumulative `le` bucket counts (`<name>_bucket`,
|
||||
* with an `le` attribute) plus `<name>_count` and `<name>_sum`.
|
||||
*
|
||||
* Requires `@opentelemetry/api` (a dependency) and, to actually export, an
|
||||
* OpenTelemetry SDK such as `@opentelemetry/sdk-metrics`.
|
||||
*
|
||||
* @param meterProvider The provider to register instruments on. Defaults to the
|
||||
* global provider from `@opentelemetry/api`.
|
||||
* @returns `true` if the recorder is installed and instruments are registered.
|
||||
* `false` if a different `metrics` recorder is already installed in this
|
||||
* process (only one global recorder is permitted), in which case a warning is
|
||||
* emitted and no instruments are created. Calling this more than once is safe;
|
||||
* instruments are created only on the first successful call.
|
||||
*/
|
||||
export function instrumentLanceDbMetrics(
|
||||
meterProvider?: MeterProvider,
|
||||
): boolean {
|
||||
if (!registerLancedbMetricsRecorder()) {
|
||||
console.warn(
|
||||
"Could not install the LanceDB metrics recorder: another `metrics` " +
|
||||
"recorder is already installed in this process. LanceDB metrics will " +
|
||||
"not be exported via OpenTelemetry.",
|
||||
);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (instrumented) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const provider = meterProvider ?? metrics.getMeterProvider();
|
||||
const meter = provider.getMeter("lancedb");
|
||||
|
||||
const scalarCallback = (metricName: string) => (result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name === metricName && point.value != null) {
|
||||
result.observe(point.value, point.attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const bucketCallback = (metricName: string) => (result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name !== metricName || point.buckets == null) {
|
||||
continue;
|
||||
}
|
||||
for (const bucket of point.buckets) {
|
||||
const attributes: Attributes = {
|
||||
...point.attributes,
|
||||
le: bucket.le,
|
||||
};
|
||||
result.observe(bucket.cumulativeCount, attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const fieldCallback =
|
||||
(metricName: string, field: "count" | "sum") =>
|
||||
(result: ObservableResult) => {
|
||||
for (const point of snapshotLancedbMetrics()) {
|
||||
if (point.name !== metricName) {
|
||||
continue;
|
||||
}
|
||||
const value = point[field];
|
||||
if (value != null) {
|
||||
result.observe(value, point.attributes);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
for (const desc of lancedbMetricsCatalog()) {
|
||||
const unit = desc.unit ?? "";
|
||||
if (desc.kind === "counter") {
|
||||
const counter = meter.createObservableCounter(desc.name, {
|
||||
unit,
|
||||
description: desc.description,
|
||||
});
|
||||
counter.addCallback(scalarCallback(desc.name));
|
||||
} else if (desc.kind === "gauge") {
|
||||
const gauge = meter.createObservableGauge(desc.name, {
|
||||
unit,
|
||||
description: desc.description,
|
||||
});
|
||||
gauge.addCallback(scalarCallback(desc.name));
|
||||
} else if (desc.kind === "histogram") {
|
||||
// `_bucket` and `_count` observe cumulative sample counts, not the
|
||||
// histogram's measured quantity, so they are unitless; only `_sum`
|
||||
// carries the histogram's unit.
|
||||
const bucket = meter.createObservableCounter(`${desc.name}_bucket`, {
|
||||
description: `${desc.description} (cumulative buckets)`,
|
||||
});
|
||||
bucket.addCallback(bucketCallback(desc.name));
|
||||
|
||||
const count = meter.createObservableCounter(`${desc.name}_count`, {
|
||||
description: `${desc.description} (count)`,
|
||||
});
|
||||
count.addCallback(fieldCallback(desc.name, "count"));
|
||||
|
||||
const sum = meter.createObservableCounter(`${desc.name}_sum`, {
|
||||
unit,
|
||||
description: `${desc.description} (sum)`,
|
||||
});
|
||||
sum.addCallback(fieldCallback(desc.name, "sum"));
|
||||
}
|
||||
}
|
||||
|
||||
instrumented = true;
|
||||
return true;
|
||||
}
|
||||
@@ -362,6 +362,9 @@ export class StandardQueryBase<
|
||||
*
|
||||
* Filtering performance can often be improved by creating a scalar index
|
||||
* on the filter column(s).
|
||||
*
|
||||
* Calling this multiple times combines the filters with a logical AND rather
|
||||
* than replacing the previous filter.
|
||||
*/
|
||||
where(predicate: string): this {
|
||||
this.doCall((inner: NativeQueryType) => inner.onlyIf(predicate));
|
||||
|
||||
@@ -158,6 +158,26 @@ export interface Version {
|
||||
metadata: Record<string, string>;
|
||||
}
|
||||
|
||||
/** Token produced by the tokenizer configured on a full-text search index. */
|
||||
export interface FtsToken {
|
||||
/** Token text after tokenizer filters have been applied. */
|
||||
text: string;
|
||||
/** Token position used by full-text query matching. */
|
||||
position: number;
|
||||
}
|
||||
|
||||
export type TokenizeTableOptions =
|
||||
| {
|
||||
/** FTS-indexed column whose tokenizer should be used. */
|
||||
column: string;
|
||||
indexName?: never;
|
||||
}
|
||||
| {
|
||||
/** Name of the FTS index whose tokenizer should be used. */
|
||||
indexName: string;
|
||||
column?: never;
|
||||
};
|
||||
|
||||
/**
|
||||
* Specification selecting Lance's MemWAL LSM-style write path for
|
||||
* `mergeInsert`.
|
||||
@@ -585,6 +605,17 @@ export abstract class Table {
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
abstract unsetLsmWriteSpec(): Promise<void>;
|
||||
/**
|
||||
* Read the {@link LsmWriteSpec} currently installed on this table.
|
||||
*
|
||||
* Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
* spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}).
|
||||
* The returned spec — including its `maintainedIndexes` and
|
||||
* `writerConfigDefaults` — mirrors what was passed to
|
||||
* {@link Table#setLsmWriteSpec}.
|
||||
* @returns {Promise<LsmWriteSpec | undefined>}
|
||||
*/
|
||||
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
|
||||
/**
|
||||
* Drain and close any cached MemWAL shard writers held for this table.
|
||||
*
|
||||
@@ -705,6 +736,19 @@ export abstract class Table {
|
||||
abstract optimize(options?: Partial<OptimizeOptions>): Promise<OptimizeStats>;
|
||||
/** List all indices that have been created with {@link Table.createIndex} */
|
||||
abstract listIndices(): Promise<IndexConfig[]>;
|
||||
/**
|
||||
* Tokenize a full-text search query using the tokenizer configured on an FTS index.
|
||||
*
|
||||
* Specify exactly one of `column` or `indexName`.
|
||||
*
|
||||
* Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in
|
||||
* the client process from index metadata. For remote tables, this means the
|
||||
* same tokenizer model files must also exist locally.
|
||||
*/
|
||||
abstract tokenize(
|
||||
query: string,
|
||||
options: TokenizeTableOptions,
|
||||
): Promise<FtsToken[]>;
|
||||
/** Return the table as an arrow table */
|
||||
abstract toArrow(): Promise<ArrowTable>;
|
||||
|
||||
@@ -1091,6 +1135,15 @@ export class LocalTable extends Table {
|
||||
return await this.inner.unsetLsmWriteSpec();
|
||||
}
|
||||
|
||||
async getLsmWriteSpec(): Promise<LsmWriteSpec | undefined> {
|
||||
// The native binding types `specType` as a plain `string`; narrow it back
|
||||
// to the public union. The Rust `From` impl only ever emits one of the
|
||||
// three valid values, so the cast is safe.
|
||||
return ((await this.inner.getLsmWriteSpec()) ?? undefined) as
|
||||
| LsmWriteSpec
|
||||
| undefined;
|
||||
}
|
||||
|
||||
async closeLsmWriters(): Promise<void> {
|
||||
return await this.inner.closeLsmWriters();
|
||||
}
|
||||
@@ -1153,6 +1206,17 @@ export class LocalTable extends Table {
|
||||
return await this.inner.listIndices();
|
||||
}
|
||||
|
||||
async tokenize(
|
||||
query: string,
|
||||
options: TokenizeTableOptions,
|
||||
): Promise<FtsToken[]> {
|
||||
return await this.inner.tokenize(
|
||||
query,
|
||||
options?.column,
|
||||
options?.indexName,
|
||||
);
|
||||
}
|
||||
|
||||
async toArrow(): Promise<ArrowTable> {
|
||||
return await this.query().toArrow();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.31.0-beta.4",
|
||||
"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.31.0-beta.4",
|
||||
"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.31.0-beta.4",
|
||||
"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.31.0-beta.4",
|
||||
"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.31.0-beta.4",
|
||||
"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.31.0-beta.4",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.31.0-beta.4",
|
||||
"version": "0.32.0-beta.2",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+73
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.4",
|
||||
"version": "0.32.0-beta.1",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.4",
|
||||
"version": "0.32.0-beta.1",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
@@ -18,6 +18,7 @@
|
||||
"win32"
|
||||
],
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"devDependencies": {
|
||||
@@ -27,6 +28,7 @@
|
||||
"@biomejs/biome": "^1.7.3",
|
||||
"@jest/globals": "^29.7.0",
|
||||
"@napi-rs/cli": "3.7.0",
|
||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
||||
"@types/axios": "^0.14.0",
|
||||
"@types/jest": "^29.1.2",
|
||||
"@types/node": "22.7.4",
|
||||
@@ -4148,6 +4150,75 @@
|
||||
"@octokit/openapi-types": "^27.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/api": {
|
||||
"version": "1.9.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.1.tgz",
|
||||
"integrity": "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==",
|
||||
"license": "Apache-2.0",
|
||||
"engines": {
|
||||
"node": ">=8.0.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/core": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/core/-/core-1.30.1.tgz",
|
||||
"integrity": "sha512-OOCM2C/QIURhJMuKaekP3TRBxBKxG/TWWA0TL2J6nXUtDnuCtccy49LUJF8xPFXMX+0LMcxFpCo8M9cGY1W6rQ==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/semantic-conventions": "1.28.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.0.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/resources": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/resources/-/resources-1.30.1.tgz",
|
||||
"integrity": "sha512-5UxZqiAgLYGFjS4s9qm5mBVo433u+dSPUFWVWXmLAD4wB65oMCoXaJP1KJa9DIYYMeHu3z4BZcStG3LC593cWA==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/core": "1.30.1",
|
||||
"@opentelemetry/semantic-conventions": "1.28.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.0.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/sdk-metrics": {
|
||||
"version": "1.30.1",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/sdk-metrics/-/sdk-metrics-1.30.1.tgz",
|
||||
"integrity": "sha512-q9zcZ0Okl8jRgmy7eNW3Ku1XSgg3sDLa5evHZpCwjspw7E8Is4K/haRPDJrBcX3YSn/Y7gUvFnByNYEKQNbNog==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@opentelemetry/core": "1.30.1",
|
||||
"@opentelemetry/resources": "1.30.1"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@opentelemetry/api": ">=1.3.0 <1.10.0"
|
||||
}
|
||||
},
|
||||
"node_modules/@opentelemetry/semantic-conventions": {
|
||||
"version": "1.28.0",
|
||||
"resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.28.0.tgz",
|
||||
"integrity": "sha512-lp4qAiMTD4sNWW4DbKLBkfiMZ4jbAboJIGOQr5DvciMRI494OapieI9qiODpOt0XBr1LjIDy1xAGAnVs5supTA==",
|
||||
"dev": true,
|
||||
"license": "Apache-2.0",
|
||||
"engines": {
|
||||
"node": ">=14"
|
||||
}
|
||||
},
|
||||
"node_modules/@protobufjs/aspromise": {
|
||||
"version": "1.1.2",
|
||||
"resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz",
|
||||
|
||||
+3
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.31.0-beta.4",
|
||||
"version": "0.32.0-beta.2",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
@@ -44,6 +44,7 @@
|
||||
"@biomejs/biome": "^1.7.3",
|
||||
"@jest/globals": "^29.7.0",
|
||||
"@napi-rs/cli": "3.7.0",
|
||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
||||
"@types/axios": "^0.14.0",
|
||||
"@types/jest": "^29.1.2",
|
||||
"@types/node": "22.7.4",
|
||||
@@ -92,6 +93,7 @@
|
||||
"version": "napi version"
|
||||
},
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
|
||||
Generated
+53
@@ -8,6 +8,9 @@ importers:
|
||||
|
||||
.:
|
||||
dependencies:
|
||||
'@opentelemetry/api':
|
||||
specifier: ^1.9.0
|
||||
version: 1.9.1
|
||||
apache-arrow:
|
||||
specifier: '>=15.0.0 <=18.1.0'
|
||||
version: 18.1.0
|
||||
@@ -33,6 +36,9 @@ importers:
|
||||
'@napi-rs/cli':
|
||||
specifier: 3.7.0
|
||||
version: 3.7.0(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)(@types/node@22.7.4)
|
||||
'@opentelemetry/sdk-metrics':
|
||||
specifier: ^1.30.0
|
||||
version: 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@types/axios':
|
||||
specifier: ^0.14.0
|
||||
version: 0.14.4
|
||||
@@ -1307,6 +1313,32 @@ packages:
|
||||
'@octokit/types@16.0.0':
|
||||
resolution: {integrity: sha512-sKq+9r1Mm4efXW1FCk7hFSeJo4QKreL/tTbR0rz/qx/r1Oa2VV83LTA/H/MuCOX7uCIJmQVRKBcbmWoySjAnSg==}
|
||||
|
||||
'@opentelemetry/api@1.9.1':
|
||||
resolution: {integrity: sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==}
|
||||
engines: {node: '>=8.0.0'}
|
||||
|
||||
'@opentelemetry/core@1.30.1':
|
||||
resolution: {integrity: sha512-OOCM2C/QIURhJMuKaekP3TRBxBKxG/TWWA0TL2J6nXUtDnuCtccy49LUJF8xPFXMX+0LMcxFpCo8M9cGY1W6rQ==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.0.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/resources@1.30.1':
|
||||
resolution: {integrity: sha512-5UxZqiAgLYGFjS4s9qm5mBVo433u+dSPUFWVWXmLAD4wB65oMCoXaJP1KJa9DIYYMeHu3z4BZcStG3LC593cWA==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.0.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/sdk-metrics@1.30.1':
|
||||
resolution: {integrity: sha512-q9zcZ0Okl8jRgmy7eNW3Ku1XSgg3sDLa5evHZpCwjspw7E8Is4K/haRPDJrBcX3YSn/Y7gUvFnByNYEKQNbNog==}
|
||||
engines: {node: '>=14'}
|
||||
peerDependencies:
|
||||
'@opentelemetry/api': '>=1.3.0 <1.10.0'
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.28.0':
|
||||
resolution: {integrity: sha512-lp4qAiMTD4sNWW4DbKLBkfiMZ4jbAboJIGOQr5DvciMRI494OapieI9qiODpOt0XBr1LjIDy1xAGAnVs5supTA==}
|
||||
engines: {node: '>=14'}
|
||||
|
||||
'@protobufjs/aspromise@1.1.2':
|
||||
resolution: {integrity: sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==}
|
||||
|
||||
@@ -4925,6 +4957,27 @@ snapshots:
|
||||
dependencies:
|
||||
'@octokit/openapi-types': 27.0.0
|
||||
|
||||
'@opentelemetry/api@1.9.1': {}
|
||||
|
||||
'@opentelemetry/core@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/semantic-conventions': 1.28.0
|
||||
|
||||
'@opentelemetry/resources@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/core': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@opentelemetry/semantic-conventions': 1.28.0
|
||||
|
||||
'@opentelemetry/sdk-metrics@1.30.1(@opentelemetry/api@1.9.1)':
|
||||
dependencies:
|
||||
'@opentelemetry/api': 1.9.1
|
||||
'@opentelemetry/core': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
'@opentelemetry/resources': 1.30.1(@opentelemetry/api@1.9.1)
|
||||
|
||||
'@opentelemetry/semantic-conventions@1.28.0': {}
|
||||
|
||||
'@protobufjs/aspromise@1.1.2':
|
||||
optional: true
|
||||
|
||||
|
||||
@@ -9,8 +9,11 @@ use lancedb::index::vector::{
|
||||
IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
|
||||
IvfRqIndexBuilder,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use napi_derive::napi;
|
||||
|
||||
use crate::error::NapiErrorExt;
|
||||
use crate::table::FtsToken;
|
||||
use crate::util::parse_distance_type;
|
||||
|
||||
#[napi]
|
||||
@@ -30,6 +33,65 @@ impl Index {
|
||||
}
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
#[allow(dead_code, clippy::too_many_arguments)]
|
||||
pub fn tokenize(
|
||||
query: String,
|
||||
base_tokenizer: Option<String>,
|
||||
language: Option<String>,
|
||||
max_token_length: Option<u32>,
|
||||
lower_case: Option<bool>,
|
||||
stem: Option<bool>,
|
||||
remove_stop_words: Option<bool>,
|
||||
ascii_folding: Option<bool>,
|
||||
ngram_min_length: Option<u32>,
|
||||
ngram_max_length: Option<u32>,
|
||||
prefix_only: Option<bool>,
|
||||
) -> napi::Result<Vec<FtsToken>> {
|
||||
let mut opts = FtsIndexBuilder::default();
|
||||
if let Some(base_tokenizer) = base_tokenizer {
|
||||
opts = opts.base_tokenizer(base_tokenizer);
|
||||
}
|
||||
if let Some(language) = language {
|
||||
opts = opts.language(&language).map_err(|_| {
|
||||
napi::Error::from_reason(format!(
|
||||
"LanceDB does not support the requested language: '{}'",
|
||||
language
|
||||
))
|
||||
})?;
|
||||
}
|
||||
if let Some(max_token_length) = max_token_length {
|
||||
opts = opts.max_token_length(Some(max_token_length as usize));
|
||||
}
|
||||
if let Some(lower_case) = lower_case {
|
||||
opts = opts.lower_case(lower_case);
|
||||
}
|
||||
if let Some(stem) = stem {
|
||||
opts = opts.stem(stem);
|
||||
}
|
||||
if let Some(remove_stop_words) = remove_stop_words {
|
||||
opts = opts.remove_stop_words(remove_stop_words);
|
||||
}
|
||||
if let Some(ascii_folding) = ascii_folding {
|
||||
opts = opts.ascii_folding(ascii_folding);
|
||||
}
|
||||
if let Some(ngram_min_length) = ngram_min_length {
|
||||
opts = opts.ngram_min_length(ngram_min_length);
|
||||
}
|
||||
if let Some(ngram_max_length) = ngram_max_length {
|
||||
opts = opts.ngram_max_length(ngram_max_length);
|
||||
}
|
||||
if let Some(prefix_only) = prefix_only {
|
||||
opts = opts.ngram_prefix_only(prefix_only);
|
||||
}
|
||||
|
||||
Ok(lancedb_tokenize(&query, &opts)
|
||||
.default_error()?
|
||||
.into_iter()
|
||||
.map(FtsToken::from)
|
||||
.collect())
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl Index {
|
||||
#[napi(factory)]
|
||||
|
||||
@@ -12,6 +12,7 @@ mod header;
|
||||
mod index;
|
||||
mod iterator;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
mod query;
|
||||
pub mod remote;
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Node.js bindings over [`lancedb::metrics_otel`].
|
||||
//!
|
||||
//! The aggregation, catalog, and histogram bucketing all live in the LanceDB
|
||||
//! core crate; this module only converts the core snapshot types into napi
|
||||
//! objects and exposes the three entry points to JavaScript, where
|
||||
//! `lancedb/otel.ts` bridges them into the user's OpenTelemetry `MeterProvider`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use lancedb::metrics_otel::{MetricPoint as CoreMetricPoint, MetricValue};
|
||||
use napi_derive::napi;
|
||||
|
||||
/// One cumulative histogram bucket: all samples with value `<= le`.
|
||||
#[napi(object)]
|
||||
pub struct MetricBucket {
|
||||
/// The inclusive upper bound of the bucket, or `"+Inf"` for the final bucket.
|
||||
pub le: String,
|
||||
/// Cumulative number of samples less than or equal to `le`.
|
||||
pub cumulative_count: f64,
|
||||
}
|
||||
|
||||
/// One aggregated metric data point. For counters and gauges only `value` is
|
||||
/// set; for histograms `buckets` (cumulative `le` counts), `count`, and `sum`
|
||||
/// are set.
|
||||
#[napi(object)]
|
||||
pub struct MetricPoint {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub attributes: HashMap<String, String>,
|
||||
pub value: Option<f64>,
|
||||
pub buckets: Option<Vec<MetricBucket>>,
|
||||
pub count: Option<f64>,
|
||||
pub sum: Option<f64>,
|
||||
}
|
||||
|
||||
impl From<CoreMetricPoint> for MetricPoint {
|
||||
fn from(point: CoreMetricPoint) -> Self {
|
||||
let kind = point.kind.as_str().to_string();
|
||||
let (value, buckets, count, sum) = match point.value {
|
||||
MetricValue::Scalar(v) => (Some(v), None, None, None),
|
||||
MetricValue::Histogram {
|
||||
buckets,
|
||||
count,
|
||||
sum,
|
||||
} => (
|
||||
None,
|
||||
Some(
|
||||
buckets
|
||||
.into_iter()
|
||||
// Counts stay well within the f64-exact integer range
|
||||
// (2^53), so this cast is lossless in practice and keeps
|
||||
// the values plain JS numbers for OpenTelemetry.
|
||||
.map(|(le, cumulative_count)| MetricBucket {
|
||||
le,
|
||||
cumulative_count: cumulative_count as f64,
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
Some(count as f64),
|
||||
Some(sum),
|
||||
),
|
||||
};
|
||||
Self {
|
||||
name: point.name,
|
||||
kind,
|
||||
attributes: point.attributes,
|
||||
value,
|
||||
buckets,
|
||||
count,
|
||||
sum,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A described metric, used by the JavaScript layer to create instruments up front.
|
||||
#[napi(object)]
|
||||
pub struct MetricDescription {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub unit: Option<String>,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// Install the LanceDB metrics recorder as the process-global `metrics` recorder.
|
||||
///
|
||||
/// Returns `true` if the recorder is installed (now or previously). Returns
|
||||
/// `false` if a *different* recorder is already installed — `metrics` allows
|
||||
/// only one global recorder per process, so LanceDB cannot coexist with another.
|
||||
#[napi]
|
||||
pub fn register_lancedb_metrics_recorder() -> bool {
|
||||
lancedb::metrics_otel::register_metrics_recorder()
|
||||
}
|
||||
|
||||
/// The catalog of described LanceDB metrics. Empty until the recorder is installed.
|
||||
#[napi]
|
||||
pub fn lancedb_metrics_catalog() -> Vec<MetricDescription> {
|
||||
lancedb::metrics_otel::metrics_catalog()
|
||||
.into_iter()
|
||||
.map(|desc| MetricDescription {
|
||||
name: desc.name,
|
||||
kind: desc.kind.as_str().to_string(),
|
||||
unit: desc.unit,
|
||||
description: desc.description,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// A point-in-time snapshot of every recorded metric. Empty until the recorder
|
||||
/// is installed.
|
||||
#[napi]
|
||||
pub fn snapshot_lancedb_metrics() -> Vec<MetricPoint> {
|
||||
lancedb::metrics_otel::snapshot_metrics()
|
||||
.into_iter()
|
||||
.map(MetricPoint::from)
|
||||
.collect()
|
||||
}
|
||||
@@ -16,6 +16,7 @@ pub struct SplitRandomOptions {
|
||||
pub counts: Option<Vec<i64>>,
|
||||
pub fixed: Option<i64>,
|
||||
pub seed: Option<i64>,
|
||||
pub clump_size: Option<i64>,
|
||||
pub split_names: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -125,10 +126,15 @@ impl PermutationBuilder {
|
||||
};
|
||||
|
||||
let seed = options.seed.map(|s| s as u64);
|
||||
let clump_size = options.clump_size.map(|c| c as u64);
|
||||
|
||||
self.modify(|builder| {
|
||||
builder.with_split_strategy(
|
||||
SplitStrategy::Random { seed, sizes },
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
},
|
||||
options.split_names.clone(),
|
||||
)
|
||||
})
|
||||
|
||||
+92
-2
@@ -8,8 +8,8 @@ use chrono::{DateTime, Utc};
|
||||
use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema};
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration,
|
||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, NewColumnTransform, OptimizeAction,
|
||||
OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
|
||||
@@ -411,6 +411,16 @@ impl Table {
|
||||
.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn get_lsm_write_spec(&self) -> napi::Result<Option<LsmWriteSpec>> {
|
||||
let spec = self
|
||||
.inner_ref()?
|
||||
.get_lsm_write_spec()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn close_lsm_writers(&self) -> napi::Result<()> {
|
||||
self.inner_ref()?.close_lsm_writers().await.default_error()
|
||||
@@ -564,6 +574,27 @@ impl Table {
|
||||
.collect::<Vec<_>>())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn tokenize(
|
||||
&self,
|
||||
query: String,
|
||||
column: Option<String>,
|
||||
index_name: Option<String>,
|
||||
) -> napi::Result<Vec<FtsToken>> {
|
||||
let table = self.inner_ref()?;
|
||||
let tokens = match (column.as_deref(), index_name.as_deref()) {
|
||||
(Some(_), Some(_)) | (None, None) => {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Specify exactly one of 'column' or 'indexName'",
|
||||
));
|
||||
}
|
||||
(Some(column), None) => table.tokenize_with_column(&query, column).await,
|
||||
(None, Some(index_name)) => table.tokenize(&query, index_name).await,
|
||||
}
|
||||
.default_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect())
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn index_stats(&self, index_name: String) -> napi::Result<Option<IndexStatistics>> {
|
||||
let tbl = self.inner_ref()?;
|
||||
@@ -671,6 +702,24 @@ impl From<lancedb::index::IndexConfig> for IndexConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
/// A token produced by the tokenizer configured on a full-text search index.
|
||||
pub struct FtsToken {
|
||||
/// The token text after the index tokenizer has applied its filters.
|
||||
pub text: String,
|
||||
/// The token position used by full-text query matching.
|
||||
pub position: u32,
|
||||
}
|
||||
|
||||
impl From<LanceDbFtsToken> for FtsToken {
|
||||
fn from(token: LanceDbFtsToken) -> Self {
|
||||
Self {
|
||||
text: token.text,
|
||||
position: token.position,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Specification selecting Lance's MemWAL LSM-style write path for
|
||||
/// `mergeInsert`.
|
||||
///
|
||||
@@ -728,6 +777,47 @@ impl TryFrom<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(spec: lancedb::table::LsmWriteSpec) -> Self {
|
||||
use lancedb::table::LsmWriteSpec as Native;
|
||||
match spec {
|
||||
Native::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "bucket".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: Some(num_buckets),
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Identity {
|
||||
column,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "identity".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Unsharded {
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "unsharded".to_string(),
|
||||
column: None,
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.34.0-beta.5"
|
||||
current_version = "0.35.0-beta.2"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.34.0-beta.5"
|
||||
version = "0.35.0-beta.2"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -47,6 +47,6 @@ pyo3-build-config = { version = "0.28", features = [
|
||||
] }
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Benchmark for StreamingDataset throughput.
|
||||
|
||||
Sweeps read_batch_size from 1 to 16384 to show how amortising the per-request
|
||||
overhead scales. Each row at each chunk size is timed via the real
|
||||
StreamingDataset so the numbers reflect production code.
|
||||
|
||||
Run with:
|
||||
cd python
|
||||
uv run --extra tests benchmarks/bench_streaming_dataloader.py
|
||||
|
||||
Optional env vars:
|
||||
BENCH_NUM_ROWS — total rows in the table (default 49152 = 24 × 2048)
|
||||
BENCH_NUM_SPLITS — number of splits (default 24)
|
||||
BENCH_STEPS — round-robin cycles to time per chunk size (default 100)
|
||||
BENCH_ROW_BYTES — approximate bytes per row padded with a binary column
|
||||
(default 4096, mimics a small embedding/image patch)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
|
||||
from lancedb.streaming import StreamingDataset
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NUM_SPLITS = int(os.environ.get("BENCH_NUM_SPLITS", 24))
|
||||
# Default: 2048 rows per split so every chunk size up to 16Ki has ≥1 full
|
||||
# chunk (except 16Ki itself which gets a single full-split fetch — still valid).
|
||||
NUM_ROWS = int(os.environ.get("BENCH_NUM_ROWS", NUM_SPLITS * 2048))
|
||||
STEPS = int(os.environ.get("BENCH_STEPS", 100))
|
||||
ROW_BYTES = int(os.environ.get("BENCH_ROW_BYTES", 4096))
|
||||
|
||||
assert NUM_ROWS % NUM_SPLITS == 0, "NUM_ROWS must be divisible by NUM_SPLITS"
|
||||
|
||||
CHUNK_SIZES = [1, 4, 16, 64, 256, 1024, 4096, 16384]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Table helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_table(db_path: str) -> lancedb.table.Table:
|
||||
db = lancedb.connect(db_path)
|
||||
payload = b"x" * ROW_BYTES
|
||||
data = pa.table(
|
||||
{
|
||||
"id": pa.array(range(NUM_ROWS), type=pa.int32()),
|
||||
"payload": pa.array([payload] * NUM_ROWS, type=pa.large_binary()),
|
||||
}
|
||||
)
|
||||
return db.create_table("bench", data, mode="overwrite")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_chunk(table, chunk_size: int, steps: int) -> tuple[int, float]:
|
||||
"""Return (rows_drained, elapsed_seconds) for one timed run."""
|
||||
total_rows = steps * NUM_SPLITS
|
||||
ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk_size
|
||||
)
|
||||
count = 0
|
||||
t0 = time.perf_counter()
|
||||
for _ in ds:
|
||||
count += 1
|
||||
if count >= total_rows:
|
||||
break
|
||||
return count, time.perf_counter() - t0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows_per_split = NUM_ROWS // NUM_SPLITS
|
||||
print("Benchmark config:")
|
||||
print(
|
||||
f" NUM_ROWS={NUM_ROWS} NUM_SPLITS={NUM_SPLITS} "
|
||||
f"rows/split={rows_per_split} STEPS={STEPS} ROW_BYTES={ROW_BYTES}"
|
||||
)
|
||||
print(f" ~{NUM_ROWS * ROW_BYTES / 1024 / 1024:.1f} MB total table size")
|
||||
print()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
print("Creating table...", flush=True)
|
||||
table = make_table(tmp)
|
||||
|
||||
cols = (
|
||||
f"{'chunk':>6} {'rows':>6} {'elapsed':>8} {'rows/s':>10} {'ms/step':>9}"
|
||||
)
|
||||
print(f"\n{cols}")
|
||||
print("-" * 52)
|
||||
|
||||
for chunk in CHUNK_SIZES:
|
||||
# Warm-up pass (one step's worth of rows)
|
||||
warmup_ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk
|
||||
)
|
||||
warmup_count = 0
|
||||
for _ in warmup_ds:
|
||||
warmup_count += 1
|
||||
if warmup_count >= NUM_SPLITS:
|
||||
break
|
||||
|
||||
drained, elapsed = bench_chunk(table, chunk, STEPS)
|
||||
rows_per_sec = drained / elapsed if elapsed > 0 else float("inf")
|
||||
ms_per_step = elapsed / STEPS * 1000
|
||||
|
||||
print(
|
||||
f"{chunk:>6} {drained:>6} {elapsed:>7.3f}s "
|
||||
f"{rows_per_sec:>10.0f} {ms_per_step:>8.1f}ms"
|
||||
)
|
||||
|
||||
print()
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -47,6 +47,10 @@ repository = "https://github.com/lancedb/lancedb"
|
||||
pylance = [
|
||||
"pylance>=5.0.0b5",
|
||||
]
|
||||
# A library only needs the OpenTelemetry API; the application supplies and
|
||||
# configures the SDK (the actual exporter/reader). See
|
||||
# https://opentelemetry.io/docs/languages/python/instrumentation/
|
||||
otel = ["opentelemetry-api"]
|
||||
tests = [
|
||||
"aiohttp>=3.9.0",
|
||||
"boto3>=1.28.57",
|
||||
@@ -61,6 +65,7 @@ tests = [
|
||||
"pylance>=5.0.0b5",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=52,<53",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
]
|
||||
dev = [
|
||||
"ruff>=0.3.0",
|
||||
|
||||
@@ -6,19 +6,22 @@ import importlib.metadata
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional, Union, Any, List
|
||||
from typing import Dict, Optional, Union, Any, List, Iterable
|
||||
|
||||
__version__ = importlib.metadata.version("lancedb")
|
||||
|
||||
from ._lancedb import connect as lancedb_connect
|
||||
from ._lancedb import FtsToken
|
||||
from ._lancedb import tokenize as _tokenize
|
||||
from .common import URI, sanitize_uri
|
||||
from urllib.parse import urlparse
|
||||
from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||
from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .schema import vector
|
||||
from .schema import blob, vector, BlobType
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
from .namespace import (
|
||||
connect_namespace,
|
||||
@@ -149,8 +152,14 @@ def connect(
|
||||
|
||||
For object storage, use a URI prefix:
|
||||
|
||||
>>> db = lancedb.connect("s3://my-bucket/lancedb",
|
||||
... storage_options={"aws_access_key_id": "***"})
|
||||
>>> db = lancedb.connect( # doctest: +SKIP
|
||||
... "s3://my-bucket/lancedb",
|
||||
... storage_options={
|
||||
... "aws_access_key_id": "***",
|
||||
... "aws_secret_access_key": "***",
|
||||
... "aws_region": "us-east-1",
|
||||
... },
|
||||
... )
|
||||
|
||||
For tests and temporary data, use an in-memory database:
|
||||
|
||||
@@ -240,6 +249,40 @@ def connect(
|
||||
)
|
||||
|
||||
|
||||
def tokenize(
|
||||
query: str,
|
||||
*,
|
||||
base_tokenizer: BaseTokenizerType = "simple",
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
) -> Iterable[FtsToken]:
|
||||
"""Tokenize a full-text search query using an explicit tokenizer.
|
||||
|
||||
This does not require a table or FTS index. The tokenizer options match
|
||||
:class:`lancedb.index.FTS`.
|
||||
"""
|
||||
return _tokenize(
|
||||
query,
|
||||
base_tokenizer=base_tokenizer,
|
||||
language=language,
|
||||
max_token_length=max_token_length,
|
||||
lower_case=lower_case,
|
||||
stem=stem,
|
||||
remove_stop_words=remove_stop_words,
|
||||
ascii_folding=ascii_folding,
|
||||
ngram_min_length=ngram_min_length,
|
||||
ngram_max_length=ngram_max_length,
|
||||
prefix_only=prefix_only,
|
||||
)
|
||||
|
||||
|
||||
WORKER_PROPERTY_PREFIX = "_lancedb_worker_"
|
||||
|
||||
|
||||
@@ -450,17 +493,21 @@ async def connect_async(
|
||||
__all__ = [
|
||||
"connect",
|
||||
"connect_async",
|
||||
"tokenize",
|
||||
"connect_namespace",
|
||||
"connect_namespace_async",
|
||||
"AsyncConnection",
|
||||
"AsyncLanceNamespaceDBConnection",
|
||||
"AsyncTable",
|
||||
"FtsToken",
|
||||
"col",
|
||||
"Expr",
|
||||
"func",
|
||||
"lit",
|
||||
"URI",
|
||||
"sanitize_uri",
|
||||
"blob",
|
||||
"BlobType",
|
||||
"vector",
|
||||
"DBConnection",
|
||||
"LanceDBConnection",
|
||||
|
||||
@@ -0,0 +1,420 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Blob fetch API and v2 projection helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from typing import TYPE_CHECKING, Optional, Union
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from .expr import Expr
|
||||
from .schema import blob_v2_column_paths
|
||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||
from .util import get_uri_scheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from _typeshed import WriteableBuffer
|
||||
|
||||
from .remote.table import RemoteTable
|
||||
from .table import AsyncTable, Table
|
||||
|
||||
BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"bytes": "all_binary",
|
||||
"descriptions": "blobs_descriptions",
|
||||
}
|
||||
|
||||
ROW_ID_FIELD_NAME = "_lance_row_id"
|
||||
|
||||
FetchBlobsSync = Callable[[str, pa.Table], pa.Array | pa.ChunkedArray]
|
||||
FetchBlobsAsync = Callable[[str, pa.Table], Awaitable[pa.Array | pa.ChunkedArray]]
|
||||
|
||||
|
||||
class BlobFile(io.RawIOBase):
|
||||
"""Seekable lazy handle from :meth:`~lancedb.table.Table.fetch_blob_files`.
|
||||
|
||||
Bytes load on ``read`` or ``read_range``, not when the handle is opened.
|
||||
Use :meth:`aread` from async code.
|
||||
"""
|
||||
|
||||
def __init__(self, inner) -> None:
|
||||
self._inner = inner
|
||||
|
||||
async def aread(self) -> bytes:
|
||||
return await self._inner.read()
|
||||
|
||||
def close(self) -> None:
|
||||
self._inner.close()
|
||||
|
||||
@property
|
||||
def closed(self) -> bool:
|
||||
return self._inner.is_closed()
|
||||
|
||||
def readable(self) -> bool:
|
||||
return True
|
||||
|
||||
def seekable(self) -> bool:
|
||||
return True
|
||||
|
||||
def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
|
||||
if whence == io.SEEK_SET:
|
||||
self._inner.seek(offset)
|
||||
elif whence == io.SEEK_CUR:
|
||||
self._inner.seek(self._inner.tell() + offset)
|
||||
elif whence == io.SEEK_END:
|
||||
self._inner.seek(self._inner.size() + offset)
|
||||
else:
|
||||
raise ValueError(f"invalid whence: {whence}")
|
||||
return self._inner.tell()
|
||||
|
||||
def tell(self) -> int:
|
||||
return self._inner.tell()
|
||||
|
||||
def size(self) -> int:
|
||||
return self._inner.size()
|
||||
|
||||
def readall(self) -> bytes:
|
||||
return self._inner.read_bytes()
|
||||
|
||||
def read(self, size: int = -1) -> bytes:
|
||||
if size == -1:
|
||||
return self._inner.read_bytes()
|
||||
return super().read(size)
|
||||
|
||||
def read_range(self, offset: int, length: int) -> bytes:
|
||||
return self._inner.read_range(offset, length)
|
||||
|
||||
def readinto(self, b: WriteableBuffer) -> int:
|
||||
view = memoryview(b).cast("B")
|
||||
chunk = self._inner.read_up_to(len(view))
|
||||
view[: len(chunk)] = chunk
|
||||
return len(chunk)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<BlobFile size={self.size()}>"
|
||||
|
||||
|
||||
def validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in BLOB_MODE_TO_HANDLING:
|
||||
modes = ", ".join(repr(mode) for mode in BLOB_MODE_TO_HANDLING)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
|
||||
|
||||
def supports_blob_auto_row_id(table: Table | AsyncTable | RemoteTable) -> bool:
|
||||
"""Blob auto row-id applies to native tables, not LanceDB Cloud."""
|
||||
from .remote.table import RemoteTable
|
||||
|
||||
if isinstance(table, RemoteTable):
|
||||
return False
|
||||
|
||||
inner = getattr(table, "_inner", None)
|
||||
if inner is not None:
|
||||
uri = inner.database().uri
|
||||
if isinstance(uri, str) and get_uri_scheme(uri) == "db":
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def projection_includes_blob_column(
|
||||
projection: QueryProjection,
|
||||
blob_columns: Iterable[str],
|
||||
) -> bool:
|
||||
columns = set(blob_columns)
|
||||
if not columns:
|
||||
return False
|
||||
if projection is None:
|
||||
return True
|
||||
for output, source in _iter_projection_pairs(projection):
|
||||
if output in columns or source in columns:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def blob_v2_projection_sources(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
) -> dict[str, str]:
|
||||
blob_columns = blob_v2_column_paths(schema)
|
||||
if not blob_columns:
|
||||
return {}
|
||||
columns = set(blob_columns)
|
||||
if projection is None:
|
||||
return {column: column for column in blob_columns}
|
||||
return {
|
||||
output: source
|
||||
for output, source in _iter_projection_pairs(projection)
|
||||
if source in columns
|
||||
}
|
||||
|
||||
|
||||
def v2_projection_needs_row_id(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
*,
|
||||
with_row_id: bool,
|
||||
) -> bool:
|
||||
if with_row_id:
|
||||
return False
|
||||
return projection_includes_blob_column(projection, blob_v2_column_paths(schema))
|
||||
|
||||
|
||||
def blob_auto_row_id_for_scan(
|
||||
table: Table | AsyncTable | RemoteTable,
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
*,
|
||||
with_row_id: bool | None,
|
||||
) -> bool:
|
||||
if with_row_id is not None:
|
||||
return False
|
||||
if not supports_blob_auto_row_id(table):
|
||||
return False
|
||||
return v2_projection_needs_row_id(schema, projection, with_row_id=False)
|
||||
|
||||
|
||||
def finalize_blob_query_table(
|
||||
tbl: pa.Table,
|
||||
*,
|
||||
user_requested_row_id: bool,
|
||||
blob_auto_row_id: bool,
|
||||
blob_paths: Iterable[str] = (),
|
||||
) -> pa.Table:
|
||||
if user_requested_row_id or not blob_auto_row_id:
|
||||
return tbl
|
||||
return stash_auto_row_ids(tbl, blob_paths)
|
||||
|
||||
|
||||
async def replace_v2_blob_columns_with_bytes(
|
||||
tbl: pa.Table,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsAsync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if output_name not in tbl.column_names:
|
||||
continue
|
||||
blobs = await fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
return tbl
|
||||
|
||||
|
||||
def replace_v2_blob_columns_with_bytes_sync(
|
||||
tbl: pa.Table,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsSync,
|
||||
) -> pa.Table:
|
||||
for output_name, source_name in blob_sources.items():
|
||||
if output_name not in tbl.column_names:
|
||||
continue
|
||||
blobs = fetch_blobs(source_name, tbl)
|
||||
tbl = _set_blob_column(tbl, output_name, blobs)
|
||||
return tbl
|
||||
|
||||
|
||||
def stash_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
if "_rowid" not in tbl.column_names:
|
||||
raise ValueError("query result has no '_rowid' column to hide")
|
||||
|
||||
present_paths = [p for p in blob_paths if p.split(".")[0] in tbl.column_names]
|
||||
if not present_paths:
|
||||
raise ValueError("query result has no blob v2 column to carry a row id")
|
||||
|
||||
row_ids = tbl["_rowid"]
|
||||
if isinstance(row_ids, pa.ChunkedArray):
|
||||
row_ids = row_ids.combine_chunks()
|
||||
row_ids = row_ids.cast(pa.uint64())
|
||||
|
||||
for path in present_paths:
|
||||
tbl = _embed_row_id_in_column(tbl, path, row_ids)
|
||||
return tbl.drop_columns(["_rowid"])
|
||||
|
||||
|
||||
def read_row_ids_from_hits(hits: pa.Table, blob_column: str) -> list[int]:
|
||||
if "_rowid" in hits.column_names:
|
||||
return hits["_rowid"].to_pylist()
|
||||
|
||||
try:
|
||||
leaf = _leaf_struct_column(hits, blob_column)
|
||||
if ROW_ID_FIELD_NAME in leaf.type.names:
|
||||
return leaf.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
# blob_column is the source name; aliased projections use the output name in hits.
|
||||
row_ids = _find_row_id_in_any_column(hits)
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
|
||||
raise ValueError(
|
||||
f"query result has no '_rowid' column and no '{ROW_ID_FIELD_NAME}' "
|
||||
f"field on blob column '{blob_column}'. Pass fresh blob query "
|
||||
"results, call .with_row_id(True), or pass a list of row ids."
|
||||
)
|
||||
|
||||
|
||||
def _find_row_id_in_any_column(tbl: pa.Table) -> Optional[list[int]]:
|
||||
for name in tbl.column_names:
|
||||
column = tbl.column(name)
|
||||
if isinstance(column, pa.ChunkedArray):
|
||||
column = column.combine_chunks()
|
||||
row_ids = _find_row_id_in_struct(column)
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
return None
|
||||
|
||||
|
||||
def _find_row_id_in_struct(array: pa.Array) -> Optional[list[int]]:
|
||||
if not pa.types.is_struct(array.type):
|
||||
return None
|
||||
if ROW_ID_FIELD_NAME in array.type.names:
|
||||
return array.field(ROW_ID_FIELD_NAME).to_pylist()
|
||||
for i in range(array.type.num_fields):
|
||||
row_ids = _find_row_id_in_struct(array.field(i))
|
||||
if row_ids is not None:
|
||||
return row_ids
|
||||
return None
|
||||
|
||||
|
||||
def _iter_projection_pairs(
|
||||
projection: QueryProjectionSpec,
|
||||
) -> Iterable[tuple[str, str]]:
|
||||
if isinstance(projection, dict):
|
||||
for name, expr in projection.items():
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
return
|
||||
for column in projection:
|
||||
if isinstance(column, str):
|
||||
yield column, column
|
||||
elif isinstance(column, tuple) and len(column) == 2:
|
||||
name, expr = column
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
|
||||
|
||||
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
|
||||
index = tbl.schema.get_field_index(output_name)
|
||||
return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs])
|
||||
|
||||
|
||||
def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table:
|
||||
def add_row_id(children: list, child_fields: list) -> None:
|
||||
children.append(row_ids)
|
||||
child_fields.append(pa.field(ROW_ID_FIELD_NAME, pa.uint64(), nullable=False))
|
||||
|
||||
return _transform_struct_column(tbl, path, add_row_id)
|
||||
|
||||
|
||||
def strip_auto_row_ids(tbl: pa.Table, blob_paths: Iterable[str]) -> pa.Table:
|
||||
"""Remove any `_lance_row_id` field embedded in blob descriptor structs.
|
||||
|
||||
For read-only descriptor views (`blob_mode="descriptions"`) that never
|
||||
fetch bytes, so have no use for the row id.
|
||||
"""
|
||||
|
||||
def drop_row_id(children: list, child_fields: list) -> None:
|
||||
for i, field in enumerate(child_fields):
|
||||
if field.name == ROW_ID_FIELD_NAME:
|
||||
del children[i], child_fields[i]
|
||||
return
|
||||
|
||||
for path in blob_paths:
|
||||
if path.split(".")[0] not in tbl.column_names:
|
||||
continue
|
||||
tbl = _transform_struct_column(tbl, path, drop_row_id)
|
||||
return tbl
|
||||
|
||||
|
||||
def _transform_struct_column(
|
||||
tbl: pa.Table, path: str, leaf_transform: Callable[[list, list], None]
|
||||
) -> pa.Table:
|
||||
top_name, *rest = path.split(".")
|
||||
top_index = tbl.schema.get_field_index(top_name)
|
||||
top_field = tbl.schema.field(top_index)
|
||||
top_array = tbl.column(top_name)
|
||||
if isinstance(top_array, pa.ChunkedArray):
|
||||
top_array = top_array.combine_chunks()
|
||||
|
||||
new_array, new_field = _rebuild_struct(top_array, top_field, rest, leaf_transform)
|
||||
return tbl.set_column(top_index, new_field, new_array)
|
||||
|
||||
|
||||
def _rebuild_struct(
|
||||
struct_array: pa.StructArray,
|
||||
struct_field: pa.Field,
|
||||
remaining_path: list[str],
|
||||
leaf_transform: Callable[[list, list], None],
|
||||
) -> tuple[pa.StructArray, pa.Field]:
|
||||
null_mask = struct_array.is_null()
|
||||
if not remaining_path:
|
||||
children = [struct_array.field(i) for i in range(struct_array.type.num_fields)]
|
||||
child_fields = list(struct_array.type)
|
||||
leaf_transform(children, child_fields)
|
||||
new_array = pa.StructArray.from_arrays(
|
||||
children, fields=child_fields, mask=null_mask
|
||||
)
|
||||
else:
|
||||
child_name = remaining_path[0]
|
||||
child_index = struct_array.type.get_field_index(child_name)
|
||||
child_array = struct_array.field(child_index)
|
||||
child_field = struct_array.type.field(child_index)
|
||||
new_child_array, new_child_field = _rebuild_struct(
|
||||
child_array, child_field, remaining_path[1:], leaf_transform
|
||||
)
|
||||
|
||||
children = []
|
||||
child_fields = []
|
||||
for i in range(struct_array.type.num_fields):
|
||||
field = struct_array.type.field(i)
|
||||
if field.name == child_name:
|
||||
children.append(new_child_array)
|
||||
child_fields.append(new_child_field)
|
||||
else:
|
||||
children.append(struct_array.field(i))
|
||||
child_fields.append(field)
|
||||
new_array = pa.StructArray.from_arrays(
|
||||
children, fields=child_fields, mask=null_mask
|
||||
)
|
||||
|
||||
new_field = pa.field(
|
||||
struct_field.name,
|
||||
new_array.type,
|
||||
nullable=struct_field.nullable,
|
||||
metadata=struct_field.metadata,
|
||||
)
|
||||
return new_array, new_field
|
||||
|
||||
|
||||
def _leaf_struct_column(tbl: pa.Table, path: str) -> pa.StructArray:
|
||||
parts = path.split(".")
|
||||
column = tbl.column(parts[0])
|
||||
if isinstance(column, pa.ChunkedArray):
|
||||
column = column.combine_chunks()
|
||||
for part in parts[1:]:
|
||||
column = column.field(part)
|
||||
return column
|
||||
|
||||
|
||||
def _normalize_blob_row_ids(
|
||||
row_ids: Union[list[int], pa.Table], blob_column: str
|
||||
) -> list[int]:
|
||||
if isinstance(row_ids, pa.Table):
|
||||
return read_row_ids_from_hits(row_ids, blob_column)
|
||||
if isinstance(row_ids, (pa.Array, pa.ChunkedArray)):
|
||||
raise ValueError(
|
||||
"pass a query table with _rowid, not a column array "
|
||||
"(use fetch_blobs('image', hits), not fetch_blobs('image', hits['image']))"
|
||||
)
|
||||
return list(row_ids)
|
||||
|
||||
|
||||
def _wrap_blob_files(handles: Iterable[object]) -> list[Optional[BlobFile]]:
|
||||
return [BlobFile(handle) if handle is not None else None for handle in handles]
|
||||
@@ -1,4 +1,5 @@
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -24,11 +25,45 @@ from lance_namespace import (
|
||||
ListTablesResponse,
|
||||
)
|
||||
from .remote import ClientConfig
|
||||
from .types import BaseTokenizerType
|
||||
|
||||
IvfHnswPq: type[HnswPq] = HnswPq
|
||||
IvfHnswSq: type[HnswSq] = HnswSq
|
||||
IvfHnswFlat: type[HnswFlat] = HnswFlat
|
||||
|
||||
class MetricPoint:
|
||||
name: str
|
||||
kind: str
|
||||
attributes: Dict[str, str]
|
||||
value: Optional[float]
|
||||
buckets: Optional[List[Tuple[str, int]]]
|
||||
count: Optional[int]
|
||||
sum: Optional[float]
|
||||
|
||||
class MetricDescription:
|
||||
name: str
|
||||
kind: str
|
||||
unit: Optional[str]
|
||||
description: str
|
||||
|
||||
def register_lancedb_metrics_recorder() -> bool: ...
|
||||
def lancedb_metrics_catalog() -> List[MetricDescription]: ...
|
||||
def snapshot_lancedb_metrics() -> List[MetricPoint]: ...
|
||||
def tokenize(
|
||||
query: str,
|
||||
*,
|
||||
base_tokenizer: BaseTokenizerType = "simple",
|
||||
language: str = "English",
|
||||
max_token_length: Optional[int] = 40,
|
||||
lower_case: bool = True,
|
||||
stem: bool = True,
|
||||
remove_stop_words: bool = True,
|
||||
ascii_folding: bool = True,
|
||||
ngram_min_length: int = 3,
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
) -> List["FtsToken"]: ...
|
||||
|
||||
class PyExpr:
|
||||
"""A type-safe DataFusion expression node (Rust-side handle)."""
|
||||
|
||||
@@ -53,7 +88,9 @@ class PyExpr:
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ...
|
||||
def expr_lit(
|
||||
value: Union[bool, int, float, str, bytes, date, datetime, Decimal],
|
||||
) -> PyExpr: ...
|
||||
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
||||
|
||||
class Session:
|
||||
@@ -159,6 +196,17 @@ class Connection(object):
|
||||
self,
|
||||
) -> Dict[str, Any]: ...
|
||||
|
||||
class BlobFile:
|
||||
async def read(self) -> bytes: ...
|
||||
def read_bytes(self) -> bytes: ...
|
||||
def close(self) -> None: ...
|
||||
def is_closed(self) -> bool: ...
|
||||
def seek(self, position: int) -> None: ...
|
||||
def tell(self) -> int: ...
|
||||
def size(self) -> int: ...
|
||||
def read_range(self, offset: int, length: int) -> bytes: ...
|
||||
def read_up_to(self, length: int) -> bytes: ...
|
||||
|
||||
class Table:
|
||||
def name(self) -> str: ...
|
||||
def __repr__(self) -> str: ...
|
||||
@@ -205,6 +253,13 @@ class Table:
|
||||
async def prewarm_index(self, index_name: str) -> None: ...
|
||||
async def prewarm_data(self, columns: Optional[List[str]] = None) -> None: ...
|
||||
async def list_indices(self) -> list[IndexConfig]: ...
|
||||
async def tokenize(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
column: Optional[str] = None,
|
||||
index_name: Optional[str] = None,
|
||||
) -> list[FtsToken]: ...
|
||||
async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
|
||||
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
|
||||
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
|
||||
@@ -226,6 +281,7 @@ class Table:
|
||||
async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
|
||||
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
||||
async def unset_lsm_write_spec(self) -> None: ...
|
||||
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
|
||||
async def close_lsm_writers(self) -> None: ...
|
||||
@property
|
||||
def tags(self) -> Tags: ...
|
||||
@@ -235,6 +291,13 @@ class Table:
|
||||
def query(self) -> Query: ...
|
||||
def take_offsets(self, offsets: list[int]) -> TakeQuery: ...
|
||||
def take_row_ids(self, row_ids: list[int]) -> TakeQuery: ...
|
||||
async def blob_columns(self) -> list[str]: ...
|
||||
async def fetch_blobs(
|
||||
self, column: str, row_ids: list[int]
|
||||
) -> pa.LargeBinaryArray: ...
|
||||
async def fetch_blob_files(
|
||||
self, column: str, row_ids: list[int]
|
||||
) -> list[Optional[BlobFile]]: ...
|
||||
def vector_search(self) -> VectorQuery: ...
|
||||
|
||||
class Tags:
|
||||
@@ -470,6 +533,10 @@ class MergeResult:
|
||||
num_attempts: int
|
||||
num_rows: int
|
||||
|
||||
class FtsToken:
|
||||
text: str
|
||||
position: int
|
||||
|
||||
class LsmWriteSpec:
|
||||
"""Specification selecting Lance's MemWAL LSM-style write path for
|
||||
`merge_insert`."""
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
import os
|
||||
from functools import cached_property
|
||||
from typing import List, Union
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -15,6 +15,8 @@ from .base import TextEmbeddingFunction
|
||||
from .registry import register
|
||||
from .utils import TEXT, api_key_not_found_help
|
||||
|
||||
EMBEDDING_BATCH_SIZE = 100
|
||||
|
||||
|
||||
@register("gemini-text")
|
||||
class GeminiText(TextEmbeddingFunction):
|
||||
@@ -81,6 +83,7 @@ class GeminiText(TextEmbeddingFunction):
|
||||
"""
|
||||
|
||||
name: str = "gemini-embedding-001"
|
||||
dim: Optional[int] = None
|
||||
query_task_type: str = "retrieval_query"
|
||||
source_task_type: str = "retrieval_document"
|
||||
|
||||
@@ -93,6 +96,8 @@ class GeminiText(TextEmbeddingFunction):
|
||||
model_config["ignored_types"] = (cached_property,)
|
||||
|
||||
def ndims(self):
|
||||
if self.dim:
|
||||
return self.dim
|
||||
# TODO: fix hardcoding
|
||||
return 768
|
||||
|
||||
@@ -133,22 +138,22 @@ class GeminiText(TextEmbeddingFunction):
|
||||
contents.append({"parts": [{"text": text}]})
|
||||
|
||||
# Build config
|
||||
config_kwargs = {}
|
||||
config_kwargs = {"output_dimensionality": self.ndims()}
|
||||
if task_type:
|
||||
config_kwargs["task_type"] = task_type.upper() # API expects uppercase
|
||||
|
||||
# Call embed_content for each content
|
||||
config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
# Call embed_content in groups of at most EMBEDDING_BATCH_SIZE docs at a time
|
||||
embeddings = []
|
||||
for content in contents:
|
||||
config = (
|
||||
types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
)
|
||||
for i in range(0, len(contents), EMBEDDING_BATCH_SIZE):
|
||||
chunk = contents[i : i + EMBEDDING_BATCH_SIZE]
|
||||
response = self.client.models.embed_content(
|
||||
model=self.name,
|
||||
contents=content,
|
||||
contents=chunk,
|
||||
config=config,
|
||||
)
|
||||
embeddings.append(response.embeddings[0].values)
|
||||
embeddings.extend([np.array(e.values) for e in response.embeddings])
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -160,5 +165,13 @@ class GeminiText(TextEmbeddingFunction):
|
||||
api_key_not_found_help("google")
|
||||
|
||||
from google import genai as genai_module
|
||||
from lancedb import __version__
|
||||
|
||||
return genai_module.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
|
||||
return genai_module.Client(
|
||||
api_key=os.environ.get("GOOGLE_API_KEY"),
|
||||
http_options={
|
||||
"headers": {
|
||||
"x-goog-api-client": f"lancedb/{__version__}",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -56,6 +56,16 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
embeddings = self._compute_embedding(texts)
|
||||
return list(embeddings)
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
state = super().__getstate__()
|
||||
state["__dict__"] = {
|
||||
k: v for k, v in state["__dict__"].items() if k != "_ollama_client"
|
||||
}
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: dict[str, Any]) -> None:
|
||||
super().__setstate__(state)
|
||||
|
||||
@cached_property
|
||||
def _ollama_client(self) -> "ollama.Client":
|
||||
ollama = attempt_import_or_raise("ollama")
|
||||
|
||||
@@ -19,6 +19,8 @@ operators::
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Iterable, Union
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -63,7 +65,7 @@ def _coerce(value: "ExprLike") -> "Expr":
|
||||
|
||||
|
||||
# Type alias used in annotations.
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes]
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes, date, datetime, Decimal]
|
||||
|
||||
|
||||
class Expr:
|
||||
@@ -118,10 +120,18 @@ class Expr:
|
||||
"""Logical AND (``expr_a & expr_b``)."""
|
||||
return Expr(self._inner.and_(_coerce(other)._inner))
|
||||
|
||||
def __rand__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical AND (``True & expr``)."""
|
||||
return Expr(_coerce(other)._inner.and_(self._inner))
|
||||
|
||||
def __or__(self, other: "Expr") -> "Expr":
|
||||
"""Logical OR (``expr_a | expr_b``)."""
|
||||
return Expr(self._inner.or_(_coerce(other)._inner))
|
||||
|
||||
def __ror__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical OR (``False | expr``)."""
|
||||
return Expr(_coerce(other)._inner.or_(self._inner))
|
||||
|
||||
def __invert__(self) -> "Expr":
|
||||
"""Logical NOT (``~expr``)."""
|
||||
return Expr(self._inner.not_())
|
||||
@@ -266,13 +276,14 @@ def col(name: str) -> Expr:
|
||||
return Expr(expr_col(name))
|
||||
|
||||
|
||||
def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr:
|
||||
"""Create a literal (constant) value expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value:
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``.
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``,
|
||||
``datetime``, or ``Decimal``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -280,6 +291,9 @@ def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
>>> col("price") * lit(1.1)
|
||||
Expr((price * 1.1))
|
||||
"""
|
||||
if not isinstance(value, (bool, int, float, str, bytes, date, datetime, Decimal)):
|
||||
raise TypeError(f"Unsupported literal type: {type(value).__name__}")
|
||||
|
||||
return Expr(expr_lit(value))
|
||||
|
||||
|
||||
|
||||
@@ -127,6 +127,8 @@ class FTS:
|
||||
- "whitespace": Split text by whitespace, but not punctuation.
|
||||
- "raw": No tokenization. The entire text is treated as a single token.
|
||||
- "ngram": N-gram tokenizer for substring-style matching.
|
||||
- "icu": ICU dictionary-based word segmentation.
|
||||
- "icu/split": ICU segmentation with simple-style delimiter splitting.
|
||||
- "jieba/*": Jieba tokenizer loaded from Lance's language model home.
|
||||
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
||||
language : str, default "English"
|
||||
|
||||
@@ -51,6 +51,15 @@ class LanceMergeInsertBuilder(object):
|
||||
If there are multiple matches then the behavior is undefined.
|
||||
Currently this causes multiple copies of the row to be created
|
||||
but that behavior is subject to change.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: Optional[str], default None
|
||||
An optional filter to limit which rows are updated. Column
|
||||
references in this expression must be prefixed with "target."
|
||||
to refer to the existing table data. For example, to only
|
||||
update rows where the existing color is red, use:
|
||||
``where="target.color = 'red'"``
|
||||
"""
|
||||
self._when_matched_update_all = True
|
||||
self._when_matched_update_all_condition = where
|
||||
|
||||
+109
-107
@@ -45,11 +45,6 @@ from lancedb._lancedb import (
|
||||
)
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
_normalize_drop_namespace_mode,
|
||||
_normalize_drop_namespace_behavior,
|
||||
)
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
connect as namespace_connect,
|
||||
@@ -58,13 +53,6 @@ from lance_namespace import (
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
ListTablesRequest,
|
||||
DescribeNamespaceRequest,
|
||||
DropTableRequest,
|
||||
RenameTableRequest,
|
||||
ListNamespacesRequest,
|
||||
CreateNamespaceRequest,
|
||||
DropNamespaceRequest,
|
||||
)
|
||||
from lancedb.table import AsyncTable, LanceTable, Table
|
||||
from lancedb.util import validate_table_name
|
||||
@@ -389,7 +377,7 @@ def _builds_namespace_natively(
|
||||
return namespace_client_impl == "rest" and bool(namespace_client_properties)
|
||||
|
||||
|
||||
def _supports_native_sync_namespace(namespace_client_impl: str) -> bool:
|
||||
def _supports_native_namespace(namespace_client_impl: str) -> bool:
|
||||
return namespace_client_impl in {"dir", "rest"}
|
||||
|
||||
|
||||
@@ -412,7 +400,6 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
_route_pushdown_to_rust: Optional[bool] = None,
|
||||
):
|
||||
"""
|
||||
Initialize a namespace-based LanceDB connection.
|
||||
@@ -454,16 +441,12 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# When the namespace client is built natively (see Rust
|
||||
# ``build_namespace_natively``), the underlying Rust table performs
|
||||
# QueryTable pushdown through the read-freshness context provider, which
|
||||
# the pure-Python ``query_table`` path bypasses.
|
||||
self._route_pushdown_to_rust = (
|
||||
_route_pushdown_to_rust
|
||||
if _route_pushdown_to_rust is not None
|
||||
else _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
# When the namespace connection or client is built natively in Rust, the
|
||||
# underlying Rust table performs QueryTable pushdown through the
|
||||
# read-freshness context provider, which the pure-Python ``query_table``
|
||||
# path bypasses.
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
@@ -909,7 +892,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace_client: LanceNamespace,
|
||||
namespace_client: Optional[LanceNamespace] = None,
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
@@ -917,6 +900,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
):
|
||||
"""
|
||||
Initialize an async namespace-based LanceDB connection.
|
||||
@@ -958,29 +942,35 @@ class AsyncLanceNamespaceDBConnection:
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# See LanceNamespaceDBConnection: when built natively the Rust table runs
|
||||
# QueryTable pushdown through the read-freshness provider, so defer to it
|
||||
# rather than the urllib3 client (which omits x-lancedb-min-timestamp).
|
||||
self._route_pushdown_to_rust = _builds_namespace_natively(
|
||||
# See LanceNamespaceDBConnection: when Rust owns the namespace
|
||||
# connection/client, its table performs QueryTable pushdown through the
|
||||
# read-freshness provider, so defer to it rather than the urllib3 client
|
||||
# path (which omits x-lancedb-min-timestamp).
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
else:
|
||||
if namespace_client is None:
|
||||
raise ValueError("namespace_client is required without a native _inner")
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def table_names(
|
||||
self,
|
||||
@@ -1004,11 +994,9 @@ class AsyncLanceNamespaceDBConnection:
|
||||
)
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
return await self._inner.table_names(
|
||||
namespace_path=namespace_path, start_after=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return response.tables if response.tables else []
|
||||
|
||||
async def create_table(
|
||||
self,
|
||||
@@ -1071,8 +1059,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Table not found" in str(e):
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if "Table not found" in str(e) or "was not found" in str(e):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
|
||||
raise
|
||||
@@ -1093,9 +1081,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""Drop a table from the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
table_id = namespace_path + [name]
|
||||
request = DropTableRequest(id=table_id)
|
||||
self._namespace_client.drop_table(request)
|
||||
await self._inner.drop_table(name, namespace_path=namespace_path)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
@@ -1109,14 +1095,17 @@ class AsyncLanceNamespaceDBConnection:
|
||||
cur_namespace_path = []
|
||||
if new_namespace_path is None:
|
||||
new_namespace_path = []
|
||||
cur_table_id = cur_namespace_path + [cur_name]
|
||||
new_namespace_id = new_namespace_path if new_namespace_path else None
|
||||
request = RenameTableRequest(
|
||||
id=cur_table_id,
|
||||
new_table_name=new_name,
|
||||
new_namespace_id=new_namespace_id,
|
||||
)
|
||||
self._namespace_client.rename_table(request)
|
||||
try:
|
||||
await self._inner.rename_table(
|
||||
cur_name,
|
||||
new_name,
|
||||
cur_namespace_path=cur_namespace_path,
|
||||
new_namespace_path=new_namespace_path,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "rename_table not implemented" in str(e):
|
||||
raise NotImplementedError("rename_table not implemented") from e
|
||||
raise
|
||||
|
||||
async def drop_database(self):
|
||||
"""Deprecated method."""
|
||||
@@ -1128,9 +1117,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""Drop all tables in the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
table_names = await self.table_names(namespace_path=namespace_path)
|
||||
for table_name in table_names:
|
||||
await self.drop_table(table_name, namespace_path=namespace_path)
|
||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||
|
||||
async def list_namespaces(
|
||||
self,
|
||||
@@ -1159,13 +1146,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListNamespacesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_namespaces(request)
|
||||
return ListNamespacesResponse(
|
||||
namespaces=response.namespaces if response.namespaces else [],
|
||||
page_token=response.page_token,
|
||||
return await self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
async def create_namespace(
|
||||
@@ -1192,15 +1174,11 @@ class AsyncLanceNamespaceDBConnection:
|
||||
CreateNamespaceResponse
|
||||
Response containing the properties of the created namespace.
|
||||
"""
|
||||
request = CreateNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_create_namespace_mode(mode),
|
||||
return await self._inner.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
)
|
||||
response = self._namespace_client.create_namespace(request)
|
||||
return CreateNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
async def drop_namespace(
|
||||
self,
|
||||
@@ -1226,20 +1204,16 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DropNamespaceResponse
|
||||
Response containing properties and transaction_id if applicable.
|
||||
"""
|
||||
request = DropNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_drop_namespace_mode(mode),
|
||||
behavior=_normalize_drop_namespace_behavior(behavior),
|
||||
)
|
||||
response = self._namespace_client.drop_namespace(request)
|
||||
return DropNamespaceResponse(
|
||||
properties=(
|
||||
response.properties if hasattr(response, "properties") else None
|
||||
),
|
||||
transaction_id=(
|
||||
response.transaction_id if hasattr(response, "transaction_id") else None
|
||||
),
|
||||
)
|
||||
try:
|
||||
return await self._inner.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
|
||||
async def describe_namespace(
|
||||
self, namespace_path: List[str]
|
||||
@@ -1257,11 +1231,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DescribeNamespaceResponse
|
||||
Response containing the namespace properties.
|
||||
"""
|
||||
request = DescribeNamespaceRequest(id=namespace_path)
|
||||
response = self._namespace_client.describe_namespace(request)
|
||||
return DescribeNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
return await self._inner.describe_namespace(namespace_path)
|
||||
|
||||
async def list_tables(
|
||||
self,
|
||||
@@ -1290,13 +1260,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return ListTablesResponse(
|
||||
tables=response.tables if response.tables else [],
|
||||
page_token=response.page_token,
|
||||
return await self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
|
||||
async def namespace_client(self) -> LanceNamespace:
|
||||
@@ -1310,6 +1275,18 @@ class AsyncLanceNamespaceDBConnection:
|
||||
LanceNamespace
|
||||
The namespace client for this connection.
|
||||
"""
|
||||
if self._namespace_client is None:
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot construct a Python namespace client without "
|
||||
"namespace implementation properties"
|
||||
)
|
||||
self._namespace_client = namespace_connect(
|
||||
self._namespace_client_impl, self._namespace_client_properties
|
||||
)
|
||||
return self._namespace_client
|
||||
|
||||
|
||||
@@ -1360,7 +1337,7 @@ def connect_namespace(
|
||||
LanceNamespaceDBConnection
|
||||
A namespace-based connection to LanceDB
|
||||
"""
|
||||
if _supports_native_sync_namespace(namespace_client_impl):
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
@@ -1384,7 +1361,6 @@ def connect_namespace(
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
_route_pushdown_to_rust=True,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
@@ -1462,6 +1438,32 @@ def connect_namespace_async(
|
||||
... tables = await db.table_names()
|
||||
... table = await db.create_table("my_table", schema=schema)
|
||||
"""
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
)
|
||||
)
|
||||
return AsyncLanceNamespaceDBConnection(
|
||||
namespace_client=None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Bridge LanceDB's internal metrics into OpenTelemetry.
|
||||
|
||||
LanceDB (through Lance core) publishes metrics (currently object store request
|
||||
counts, bytes, latency, errors, and throttles) through the Rust ``metrics``
|
||||
facade. This module installs a process-global recorder that aggregates them and
|
||||
registers OpenTelemetry observable instruments that report the aggregated values
|
||||
into the user's ``MeterProvider``.
|
||||
|
||||
The bridge is generic: every metric LanceDB describes is surfaced automatically,
|
||||
with no per-metric Python code. Histograms have no asynchronous OpenTelemetry
|
||||
instrument, so each is exported Prometheus-style as cumulative ``le`` buckets
|
||||
plus ``_count`` and ``_sum`` observable counters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from ._lancedb import (
|
||||
lancedb_metrics_catalog,
|
||||
register_lancedb_metrics_recorder,
|
||||
snapshot_lancedb_metrics,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.metrics import MeterProvider
|
||||
|
||||
_INSTRUMENTED = False
|
||||
|
||||
|
||||
def instrument_lancedb_metrics(
|
||||
meter_provider: Optional["MeterProvider"] = None,
|
||||
) -> bool:
|
||||
"""Register LanceDB metrics as OpenTelemetry observable instruments.
|
||||
|
||||
Installs a process-global metrics recorder and creates one observable
|
||||
instrument per LanceDB metric on the given (or global) ``MeterProvider``. The
|
||||
user's configured ``MetricReader`` then collects them on its own schedule.
|
||||
|
||||
Counters and gauges map directly to observable counters/gauges. Each
|
||||
histogram is exported as cumulative ``le`` bucket counts (``<name>_bucket``,
|
||||
with an ``le`` attribute) plus ``<name>_count`` and ``<name>_sum``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
meter_provider : opentelemetry.metrics.MeterProvider, optional
|
||||
The provider to register instruments on. Defaults to the global provider
|
||||
from ``opentelemetry.metrics.get_meter_provider()``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
``True`` if the recorder is installed and instruments are registered.
|
||||
``False`` if a different ``metrics`` recorder is already installed in
|
||||
this process (``metrics`` permits only one global recorder), in which
|
||||
case a warning is emitted and no instruments are created.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Requires the OpenTelemetry API (``pip install lancedb[otel]``) and, to
|
||||
actually export, an OpenTelemetry SDK (``pip install opentelemetry-sdk``)
|
||||
configured by the application. Calling this more than once is safe;
|
||||
instruments are created only on the first successful call.
|
||||
"""
|
||||
global _INSTRUMENTED
|
||||
|
||||
try:
|
||||
from opentelemetry.metrics import Observation, get_meter_provider
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"instrument_lancedb_metrics requires the OpenTelemetry API/SDK. "
|
||||
"Install it with `pip install lancedb[otel]` or "
|
||||
"`pip install opentelemetry-sdk`."
|
||||
) from exc
|
||||
|
||||
if not register_lancedb_metrics_recorder():
|
||||
warnings.warn(
|
||||
"Could not install the LanceDB metrics recorder: another `metrics` "
|
||||
"recorder is already installed in this process. LanceDB metrics will "
|
||||
"not be exported via OpenTelemetry.",
|
||||
stacklevel=2,
|
||||
)
|
||||
return False
|
||||
|
||||
if _INSTRUMENTED:
|
||||
return True
|
||||
|
||||
provider = meter_provider or get_meter_provider()
|
||||
meter = provider.get_meter("lancedb")
|
||||
|
||||
def scalar_callback(metric_name: str):
|
||||
def callback(_options):
|
||||
return [
|
||||
Observation(point.value, point.attributes)
|
||||
for point in snapshot_lancedb_metrics()
|
||||
if point.name == metric_name and point.value is not None
|
||||
]
|
||||
|
||||
return callback
|
||||
|
||||
def bucket_callback(metric_name: str):
|
||||
def callback(_options):
|
||||
observations = []
|
||||
for point in snapshot_lancedb_metrics():
|
||||
if point.name != metric_name or point.buckets is None:
|
||||
continue
|
||||
for le, cumulative in point.buckets:
|
||||
attributes = dict(point.attributes)
|
||||
attributes["le"] = le
|
||||
observations.append(Observation(cumulative, attributes))
|
||||
return observations
|
||||
|
||||
return callback
|
||||
|
||||
def field_callback(metric_name: str, field: str):
|
||||
def callback(_options):
|
||||
observations = []
|
||||
for point in snapshot_lancedb_metrics():
|
||||
if point.name != metric_name:
|
||||
continue
|
||||
value = getattr(point, field)
|
||||
if value is not None:
|
||||
observations.append(Observation(value, point.attributes))
|
||||
return observations
|
||||
|
||||
return callback
|
||||
|
||||
for desc in lancedb_metrics_catalog():
|
||||
unit = desc.unit or ""
|
||||
if desc.kind == "counter":
|
||||
meter.create_observable_counter(
|
||||
desc.name,
|
||||
callbacks=[scalar_callback(desc.name)],
|
||||
unit=unit,
|
||||
description=desc.description,
|
||||
)
|
||||
elif desc.kind == "gauge":
|
||||
meter.create_observable_gauge(
|
||||
desc.name,
|
||||
callbacks=[scalar_callback(desc.name)],
|
||||
unit=unit,
|
||||
description=desc.description,
|
||||
)
|
||||
elif desc.kind == "histogram":
|
||||
# `_bucket` and `_count` observe cumulative sample counts, not the
|
||||
# histogram's measured quantity, so they are unitless; only `_sum`
|
||||
# carries the histogram's unit.
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_bucket",
|
||||
callbacks=[bucket_callback(desc.name)],
|
||||
description=f"{desc.description} (cumulative buckets)",
|
||||
)
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_count",
|
||||
callbacks=[field_callback(desc.name, "count")],
|
||||
description=f"{desc.description} (count)",
|
||||
)
|
||||
meter.create_observable_counter(
|
||||
f"{desc.name}_sum",
|
||||
callbacks=[field_callback(desc.name, "sum")],
|
||||
unit=unit,
|
||||
description=f"{desc.description} (sum)",
|
||||
)
|
||||
|
||||
_INSTRUMENTED = True
|
||||
return True
|
||||
@@ -11,7 +11,7 @@ import pyarrow as pa
|
||||
from ._lancedb import async_permutation_builder, PermutationReader
|
||||
from .table import LanceTable, Table
|
||||
from .background_loop import LOOP
|
||||
from .util import batch_to_tensor, batch_to_tensor_rows
|
||||
from .util import batch_to_tensor, batch_to_tensor_dict, batch_to_tensor_rows
|
||||
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -65,6 +65,7 @@ class PermutationBuilder:
|
||||
counts: Optional[list[int]] = None,
|
||||
fixed: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
clump_size: Optional[int] = None,
|
||||
split_names: Optional[list[str]] = None,
|
||||
) -> "PermutationBuilder":
|
||||
"""
|
||||
@@ -87,6 +88,9 @@ class PermutationBuilder:
|
||||
Rows will be randomly assigned to splits. The optional seed can be provided to
|
||||
make the assignment deterministic.
|
||||
|
||||
If clump_size is provided, rows are shuffled as contiguous groups of that size,
|
||||
preserving I/O locality while still randomising the split assignment.
|
||||
|
||||
The optional split_names can be provided to name the splits. If not provided,
|
||||
the splits can only be referenced by their index.
|
||||
"""
|
||||
@@ -95,6 +99,7 @@ class PermutationBuilder:
|
||||
counts=counts,
|
||||
fixed=fixed,
|
||||
seed=seed,
|
||||
clump_size=clump_size,
|
||||
split_names=split_names,
|
||||
)
|
||||
return self
|
||||
@@ -941,6 +946,7 @@ class Permutation:
|
||||
"pandas",
|
||||
"arrow",
|
||||
"torch",
|
||||
"torch_row",
|
||||
"torch_col",
|
||||
"polars",
|
||||
],
|
||||
@@ -956,15 +962,19 @@ class Permutation:
|
||||
- "python_col" - the batch will be a dict of lists (one entry per column)
|
||||
- "pandas" - the batch will be a pandas DataFrame
|
||||
- "arrow" - the batch will be a pyarrow RecordBatch
|
||||
- "torch" - the batch will be a list of tensors, one per row
|
||||
- "torch" - the batch will be a list of per-row dicts mapping column
|
||||
name to a 0-D torch tensor. Works with the default
|
||||
``torch.utils.data.DataLoader`` collate, which stacks the per-row
|
||||
dicts back into a dict of batched tensors.
|
||||
- "torch_row" - the batch will be a list of tensors, one per row
|
||||
- "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns)
|
||||
- "polars" - the batch will be a polars DataFrame
|
||||
|
||||
Conversion may or may not involve a data copy. Lance uses Arrow internally
|
||||
and so it is able to zero-copy to the arrow and polars formats.
|
||||
|
||||
Conversion to torch_col will be zero-copy but will only support a subset of data
|
||||
types (numeric types).
|
||||
Conversion to torch and torch_col will be zero-copy but will only support a
|
||||
subset of data types (numeric types).
|
||||
|
||||
Conversion to numpy and/or pandas will typically be zero-copy for numeric
|
||||
types. Conversion of strings, lists, and structs will require creating python
|
||||
@@ -985,6 +995,8 @@ class Permutation:
|
||||
elif format == "arrow":
|
||||
return self.with_transform(Transforms.arrow2arrow)
|
||||
elif format == "torch":
|
||||
return self.with_transform(batch_to_tensor_dict)
|
||||
elif format == "torch_row":
|
||||
return self.with_transform(batch_to_tensor_rows)
|
||||
elif format == "torch_col":
|
||||
return self.with_transform(batch_to_tensor)
|
||||
|
||||
+336
-81
@@ -15,10 +15,12 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
import deprecation
|
||||
@@ -39,15 +41,21 @@ from .expr import Expr
|
||||
from .rerankers.base import Reranker
|
||||
from .rerankers.rrf import RRFReranker
|
||||
from .rerankers.util import check_reranker_result
|
||||
from .schema import is_blob_like_field, schema_has_blob_field
|
||||
from .util import flatten_columns
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
_BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"bytes": "all_binary",
|
||||
"descriptions": "blobs_descriptions",
|
||||
}
|
||||
from ._blob import (
|
||||
BLOB_MODE_TO_HANDLING,
|
||||
FetchBlobsAsync,
|
||||
FetchBlobsSync,
|
||||
blob_auto_row_id_for_scan,
|
||||
blob_v2_projection_sources,
|
||||
finalize_blob_query_table,
|
||||
replace_v2_blob_columns_with_bytes,
|
||||
replace_v2_blob_columns_with_bytes_sync,
|
||||
supports_blob_auto_row_id,
|
||||
validate_blob_mode,
|
||||
)
|
||||
from .types import BlobMode, QueryProjection
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import sys
|
||||
@@ -73,25 +81,22 @@ if TYPE_CHECKING:
|
||||
T = TypeVar("T", bound="LanceModel")
|
||||
|
||||
|
||||
def _validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in _BLOB_MODE_TO_HANDLING:
|
||||
modes = ", ".join(repr(mode) for mode in _BLOB_MODE_TO_HANDLING)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
@runtime_checkable
|
||||
class _LanceScanner(Protocol):
|
||||
projected_schema: pa.Schema | None
|
||||
schema: pa.Schema | None
|
||||
|
||||
def to_pandas(self, blob_mode: BlobMode | None = ..., **kwargs) -> pd.DataFrame: ...
|
||||
|
||||
def _field_is_blob(field: pa.Field) -> bool:
|
||||
metadata = field.metadata or {}
|
||||
return metadata.get(b"lance-encoding:blob") == b"true" or (
|
||||
metadata.get("lance-encoding:blob") == "true"
|
||||
)
|
||||
def to_pyarrow(self): ...
|
||||
|
||||
def to_table(self) -> pa.Table: ...
|
||||
|
||||
def _schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return any(_field_is_blob(field) for field in schema)
|
||||
def to_reader(self): ...
|
||||
|
||||
|
||||
def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool:
|
||||
return blob_mode in _BLOB_MODE_TO_HANDLING and _schema_has_blob_field(schema)
|
||||
return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema)
|
||||
|
||||
|
||||
def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
|
||||
@@ -119,13 +124,28 @@ def _filter_to_sql(filter: Optional[Union[str, Expr]]) -> Optional[str]:
|
||||
return filter
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(
|
||||
columns: Optional[
|
||||
Union[
|
||||
List[str], List[Tuple[str, Union[str, Expr]]], Dict[str, Union[str, Expr]]
|
||||
]
|
||||
],
|
||||
) -> Dict[str, Any]:
|
||||
def _combine_where(
|
||||
existing: Optional[Union[str, Expr]], new: Union[str, Expr]
|
||||
) -> Union[str, Expr]:
|
||||
"""Combine a new filter with an existing one using a logical AND.
|
||||
|
||||
Calling ``where`` more than once composes the filters with AND instead of
|
||||
replacing the previous filter. Two :class:`~lancedb.expr.Expr` filters are
|
||||
combined as an expression; otherwise both filters are lowered to SQL strings
|
||||
and combined as SQL.
|
||||
"""
|
||||
if existing is None:
|
||||
return new
|
||||
existing_is_expr = isinstance(existing, Expr)
|
||||
new_is_expr = isinstance(new, Expr)
|
||||
if existing_is_expr and new_is_expr:
|
||||
return existing & new
|
||||
existing_sql = existing.to_sql() if existing_is_expr else existing
|
||||
new_sql = new.to_sql() if new_is_expr else new
|
||||
return f"({existing_sql}) AND ({new_sql})"
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
if columns is None:
|
||||
return {}
|
||||
if isinstance(columns, list):
|
||||
@@ -150,7 +170,11 @@ def _projection_to_scanner_kwargs(
|
||||
|
||||
|
||||
def _scanner_kwargs_for_query(
|
||||
query: Query, blob_mode: BlobMode, dataset: Optional[Any] = None
|
||||
query: Query,
|
||||
blob_mode: BlobMode,
|
||||
dataset: Optional[Any] = None,
|
||||
*,
|
||||
with_row_id: Optional[bool] = None,
|
||||
) -> Dict[str, Any]:
|
||||
fragments = _scanner_fragments_for_query(query, dataset)
|
||||
kwargs = {
|
||||
@@ -158,10 +182,10 @@ def _scanner_kwargs_for_query(
|
||||
"filter": _filter_to_sql(query.filter),
|
||||
"limit": query.limit,
|
||||
"offset": query.offset,
|
||||
"with_row_id": query.with_row_id,
|
||||
"with_row_id": with_row_id if with_row_id is not None else query.with_row_id,
|
||||
"with_row_address": query.with_row_address,
|
||||
"fast_search": query.fast_search,
|
||||
"blob_handling": _BLOB_MODE_TO_HANDLING[blob_mode],
|
||||
"blob_handling": BLOB_MODE_TO_HANDLING[blob_mode],
|
||||
"fragments": fragments,
|
||||
}
|
||||
return {key: value for key, value in kwargs.items() if value is not None}
|
||||
@@ -194,11 +218,11 @@ def _scanner_fragments_for_query(query: Query, dataset: Optional[Any]) -> Option
|
||||
def _ensure_lazy_blob_frame(
|
||||
df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode
|
||||
) -> "pd.DataFrame":
|
||||
if blob_mode != "lazy" or not _schema_has_blob_field(schema) or len(df) == 0:
|
||||
if blob_mode != "lazy" or not schema_has_blob_field(schema) or len(df) == 0:
|
||||
return df
|
||||
|
||||
for field in schema:
|
||||
if not _field_is_blob(field) or field.name not in df.columns:
|
||||
if not is_blob_like_field(field) or field.name not in df.columns:
|
||||
continue
|
||||
value = df[field.name].iloc[0]
|
||||
if value is not None and not hasattr(value, "readall"):
|
||||
@@ -208,7 +232,7 @@ def _ensure_lazy_blob_frame(
|
||||
return df
|
||||
|
||||
|
||||
def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||
def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
||||
if hasattr(scanner, "to_pyarrow"):
|
||||
reader = scanner.to_pyarrow()
|
||||
return reader.read_all()
|
||||
@@ -218,7 +242,9 @@ def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||
return reader.read_all()
|
||||
|
||||
|
||||
def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataFrame":
|
||||
def _scanner_to_pandas(
|
||||
scanner: _LanceScanner, blob_mode: BlobMode, **kwargs
|
||||
) -> pd.DataFrame:
|
||||
schema = getattr(scanner, "projected_schema", None)
|
||||
if schema is None:
|
||||
schema = getattr(scanner, "schema", None)
|
||||
@@ -239,13 +265,71 @@ def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataF
|
||||
return df
|
||||
|
||||
tbl = _scanner_to_table(scanner)
|
||||
if blob_mode == "lazy" and _schema_has_blob_field(tbl.schema):
|
||||
if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
|
||||
raise _unsupported_blob_pandas_error(
|
||||
"the Lance scanner does not expose to_pandas"
|
||||
)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
|
||||
|
||||
def _finish_plain_scan_pandas(
|
||||
scanner: _LanceScanner,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsSync,
|
||||
strip_auto_row_id: bool,
|
||||
flatten: Optional[Union[int, bool]],
|
||||
**kwargs,
|
||||
) -> pd.DataFrame:
|
||||
if blob_sources:
|
||||
tbl = _scanner_to_table(scanner)
|
||||
tbl = replace_v2_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(tbl, flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
return tbl.to_pandas(**kwargs)
|
||||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
if strip_auto_row_id and "_rowid" in df.columns:
|
||||
return df.drop(columns=["_rowid"])
|
||||
return df
|
||||
|
||||
|
||||
async def _finish_plain_scan_pandas_async(
|
||||
scanner: _LanceScanner,
|
||||
*,
|
||||
blob_mode: BlobMode,
|
||||
blob_sources: dict[str, str],
|
||||
fetch_blobs: FetchBlobsAsync,
|
||||
strip_auto_row_id: bool,
|
||||
flatten: Optional[Union[int, bool]],
|
||||
**kwargs,
|
||||
) -> pd.DataFrame:
|
||||
if blob_sources:
|
||||
tbl = _scanner_to_table(scanner)
|
||||
tbl = await replace_v2_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(tbl, flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||||
tbl = tbl.drop_columns(["_rowid"])
|
||||
return tbl.to_pandas(**kwargs)
|
||||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
if strip_auto_row_id and "_rowid" in df.columns:
|
||||
return df.drop(columns=["_rowid"])
|
||||
return df
|
||||
|
||||
|
||||
# Pydantic validation function for vector queries
|
||||
def ensure_vector_query(
|
||||
val: Any,
|
||||
@@ -653,7 +737,7 @@ class Query(pydantic.BaseModel):
|
||||
distance_type: Optional[str] = None
|
||||
|
||||
# which columns to return in the results (dict values may be str or Expr)
|
||||
columns: Optional[Union[List[str], Dict[str, Union[str, Expr]]]] = None
|
||||
columns: QueryProjection = None
|
||||
|
||||
# minimum number of IVF partitions to search
|
||||
#
|
||||
@@ -937,7 +1021,7 @@ class LanceQueryBuilder(ABC):
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
validate_blob_mode(blob_mode)
|
||||
output_schema = getattr(self, "output_schema", None)
|
||||
if output_schema is not None:
|
||||
schema = output_schema()
|
||||
@@ -996,6 +1080,11 @@ class LanceQueryBuilder(ABC):
|
||||
Execute the query and return the results as a pyarrow
|
||||
[RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html)
|
||||
|
||||
For v2 blob projections, ``to_batches`` keeps the auto ``_rowid``
|
||||
column visible so batch consumers can call ``fetch_blobs``. Use
|
||||
``to_arrow``, ``to_list``, or ``to_pandas`` if you want LanceDB to hide
|
||||
auto row ids in the final collected result.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
batch_size: int
|
||||
@@ -1148,8 +1237,13 @@ class LanceQueryBuilder(ABC):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = where
|
||||
self._where = _combine_where(self._where, where)
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
|
||||
@@ -1169,6 +1263,42 @@ class LanceQueryBuilder(ABC):
|
||||
self._with_row_id = with_row_id
|
||||
return self
|
||||
|
||||
def _user_requested_row_id(self) -> bool:
|
||||
return self._with_row_id is True
|
||||
|
||||
def _blob_auto_row_id_enabled(self) -> bool:
|
||||
if not supports_blob_auto_row_id(self._table):
|
||||
return False
|
||||
return blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
self._table.schema,
|
||||
self._columns,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
|
||||
def _scan_needs_row_id(self) -> bool:
|
||||
return self._user_requested_row_id() or self._blob_auto_row_id_enabled()
|
||||
|
||||
def _query_for_scan(self) -> Query:
|
||||
query = self.to_query_object()
|
||||
if self._scan_needs_row_id():
|
||||
query.with_row_id = True
|
||||
return query
|
||||
|
||||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
blob_paths = (
|
||||
blob_v2_projection_sources(self._table.schema, self._columns).keys()
|
||||
if blob_auto_row_id
|
||||
else ()
|
||||
)
|
||||
return finalize_blob_query_table(
|
||||
tbl,
|
||||
user_requested_row_id=self._user_requested_row_id(),
|
||||
blob_auto_row_id=blob_auto_row_id,
|
||||
blob_paths=blob_paths,
|
||||
)
|
||||
|
||||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||||
"""Set whether to return row addresses.
|
||||
|
||||
@@ -1345,13 +1475,29 @@ class LanceQueryBuilder(ABC):
|
||||
return None
|
||||
|
||||
dataset = self._table.to_lance()
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||||
blob_sources = (
|
||||
blob_v2_projection_sources(self._table.schema, query.columns)
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(
|
||||
query,
|
||||
"descriptions" if blob_sources else blob_mode,
|
||||
dataset,
|
||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||||
)
|
||||
)
|
||||
return _finish_plain_scan_pandas(
|
||||
scanner,
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def to_query_object(self) -> Query:
|
||||
@@ -1599,7 +1745,9 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
"""
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
"""
|
||||
@@ -1659,7 +1807,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
vector = self._query if isinstance(self._query, list) else self._query.tolist()
|
||||
if isinstance(vector[0], np.ndarray):
|
||||
vector = [v.tolist() for v in vector]
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
result_set = self._table._execute_query(
|
||||
query, batch_size=batch_size, timeout=timeout
|
||||
)
|
||||
@@ -1693,8 +1841,13 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = where
|
||||
self._where = _combine_where(self._where, where)
|
||||
if prefilter is not None:
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
@@ -1798,8 +1951,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1809,6 +1961,21 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
self._phrase_query = phrase_query
|
||||
return self
|
||||
|
||||
def _query_with_phrase_semantics(self) -> str | FullTextQuery:
|
||||
query = self._query
|
||||
if not self._phrase_query:
|
||||
return query
|
||||
if isinstance(query, str):
|
||||
if not query.startswith('"') or not query.endswith('"'):
|
||||
return f'"{query}"'
|
||||
return query
|
||||
if isinstance(query, PhraseQuery):
|
||||
return query
|
||||
raise TypeError(
|
||||
"phrase_query() requires a string or PhraseQuery, "
|
||||
f"got {type(query).__name__}"
|
||||
)
|
||||
|
||||
def fast_search(self) -> LanceFtsQueryBuilder:
|
||||
"""
|
||||
Skip a flat search of unindexed data. This will improve
|
||||
@@ -1833,7 +2000,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
full_text_query=FullTextSearchQuery(
|
||||
query=self._query, columns=self._fts_columns
|
||||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
||||
),
|
||||
offset=self._offset,
|
||||
fast_search=self._fast_search,
|
||||
@@ -1851,22 +2018,13 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
self._table._ensure_no_legacy_fts_index()
|
||||
|
||||
query = self._query
|
||||
if self._phrase_query:
|
||||
if isinstance(query, str):
|
||||
if not query.startswith('"') or not query.endswith('"'):
|
||||
self._query = f'"{query}"'
|
||||
elif isinstance(query, FullTextQuery) and not isinstance(
|
||||
query, PhraseQuery
|
||||
):
|
||||
raise TypeError("Please use PhraseQuery for phrase queries.")
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
results = self._table._execute_query(query, timeout=timeout)
|
||||
results = results.read_all()
|
||||
if self._reranker is not None:
|
||||
results = self._reranker.rerank_fts(self._query, results)
|
||||
check_reranker_result(results)
|
||||
return results
|
||||
return self._finalize_blob_query_table(results)
|
||||
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
@@ -1894,7 +2052,9 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
return Query(
|
||||
@@ -1916,7 +2076,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
) -> pa.RecordBatchReader:
|
||||
query = self.to_query_object()
|
||||
query = self._query_for_scan()
|
||||
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
||||
|
||||
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
||||
@@ -1988,14 +2148,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
return vector_query, text_query
|
||||
|
||||
def phrase_query(self, phrase_query: bool = None) -> LanceHybridQueryBuilder:
|
||||
def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
|
||||
"""Set whether to use phrase query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -2020,15 +2179,25 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
fts_results = fts_future.result()
|
||||
vector_results = vector_future.result()
|
||||
|
||||
return self._combine_hybrid_results(
|
||||
results = self._combine_hybrid_results(
|
||||
fts_results=fts_results,
|
||||
vector_results=vector_results,
|
||||
norm=self._norm,
|
||||
fts_query=self._fts_query._query,
|
||||
reranker=self._reranker,
|
||||
limit=self._limit,
|
||||
with_row_ids=self._with_row_id,
|
||||
with_row_ids=True,
|
||||
)
|
||||
return self._finish_hybrid_results(results)
|
||||
|
||||
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
|
||||
if self._user_requested_row_id():
|
||||
return results
|
||||
if self._blob_auto_row_id_enabled():
|
||||
return self._finalize_blob_query_table(results)
|
||||
if "_rowid" in results.column_names:
|
||||
return results.drop(["_rowid"])
|
||||
return results
|
||||
|
||||
@staticmethod
|
||||
def _combine_hybrid_results(
|
||||
@@ -2469,7 +2638,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
self._vector_query.ef(self._ef)
|
||||
if self._bypass_vector_index:
|
||||
self._vector_query.bypass_vector_index()
|
||||
if self._lower_bound or self._upper_bound:
|
||||
if self._lower_bound is not None or self._upper_bound is not None:
|
||||
self._vector_query.distance_range(
|
||||
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
||||
)
|
||||
@@ -2499,6 +2668,9 @@ class AsyncQueryBase(object):
|
||||
self._with_row_address = None
|
||||
self._fragments = None
|
||||
self._fragment_ids = None
|
||||
self._with_row_id = None
|
||||
self._blob_auto_row_id = False
|
||||
self._blob_paths: tuple[str, ...] = ()
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
"""
|
||||
@@ -2508,11 +2680,46 @@ class AsyncQueryBase(object):
|
||||
python and more easily serializable.
|
||||
"""
|
||||
query = Query.from_inner(self._inner.to_query_request())
|
||||
query.with_row_id = self._user_requested_row_id()
|
||||
query.with_row_address = self._with_row_address
|
||||
query.fragments = self._fragments
|
||||
query.fragment_ids = self._fragment_ids
|
||||
return query
|
||||
|
||||
def _user_requested_row_id(self) -> bool:
|
||||
return self._with_row_id is True
|
||||
|
||||
def _blob_auto_row_id_enabled(self) -> bool:
|
||||
return self._blob_auto_row_id
|
||||
|
||||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||||
return finalize_blob_query_table(
|
||||
tbl,
|
||||
user_requested_row_id=self._user_requested_row_id(),
|
||||
blob_auto_row_id=self._blob_auto_row_id_enabled(),
|
||||
blob_paths=self._blob_paths,
|
||||
)
|
||||
|
||||
async def _maybe_add_blob_row_id(self) -> None:
|
||||
if self._table is None or not supports_blob_auto_row_id(self._table):
|
||||
self._blob_auto_row_id = False
|
||||
self._blob_paths = ()
|
||||
return
|
||||
|
||||
req = self._inner.to_query_request()
|
||||
schema = await self._table.schema()
|
||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if not self._blob_auto_row_id:
|
||||
self._blob_paths = ()
|
||||
return
|
||||
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
|
||||
self._inner.with_row_id()
|
||||
|
||||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||||
"""
|
||||
Return only the specified columns.
|
||||
@@ -2565,6 +2772,7 @@ class AsyncQueryBase(object):
|
||||
"""
|
||||
Include the _rowid column in the results.
|
||||
"""
|
||||
self._with_row_id = True
|
||||
self._inner.with_row_id()
|
||||
return self
|
||||
|
||||
@@ -2611,6 +2819,7 @@ class AsyncQueryBase(object):
|
||||
If not specified, no timeout is applied. If the query does not
|
||||
complete within the specified time, an error will be raised.
|
||||
"""
|
||||
await self._maybe_add_blob_row_id()
|
||||
return AsyncRecordBatchReader(
|
||||
await self._inner.execute(
|
||||
max_batch_length=max_batch_length, timeout=timeout
|
||||
@@ -2641,8 +2850,8 @@ class AsyncQueryBase(object):
|
||||
complete within the specified time, an error will be raised.
|
||||
"""
|
||||
batch_iter = await self.to_batches(timeout=timeout)
|
||||
return pa.Table.from_batches(
|
||||
await batch_iter.read_all(), schema=batch_iter.schema
|
||||
return self._finalize_blob_query_table(
|
||||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
||||
)
|
||||
|
||||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
@@ -2709,7 +2918,7 @@ class AsyncQueryBase(object):
|
||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||
optional flattening.
|
||||
"""
|
||||
_validate_blob_mode(blob_mode)
|
||||
validate_blob_mode(blob_mode)
|
||||
if hasattr(self._inner, "output_schema"):
|
||||
schema = await self.output_schema()
|
||||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||||
@@ -2750,14 +2959,36 @@ class AsyncQueryBase(object):
|
||||
if not _query_is_plain_scan(query):
|
||||
return None
|
||||
|
||||
schema = await self._table.schema()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
query.columns,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
blob_sources = (
|
||||
blob_v2_projection_sources(schema, query.columns)
|
||||
if blob_mode == "bytes"
|
||||
else {}
|
||||
)
|
||||
dataset = await self._table._to_lance()
|
||||
scanner = dataset.scanner(
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
**_scanner_kwargs_for_query(
|
||||
query,
|
||||
"descriptions" if blob_sources else blob_mode,
|
||||
dataset,
|
||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||||
)
|
||||
)
|
||||
return await _finish_plain_scan_pandas_async(
|
||||
scanner,
|
||||
blob_mode=blob_mode,
|
||||
blob_sources=blob_sources,
|
||||
fetch_blobs=self._table.fetch_blobs,
|
||||
strip_auto_row_id=blob_auto_row_id,
|
||||
flatten=flatten,
|
||||
**kwargs,
|
||||
)
|
||||
if flatten is not None:
|
||||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||
return tbl.to_pandas(**kwargs)
|
||||
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||
|
||||
async def to_polars(
|
||||
self,
|
||||
@@ -2894,6 +3125,9 @@ class AsyncStandardQuery(AsyncQueryBase):
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
if isinstance(predicate, Expr):
|
||||
self._inner.where_expr(predicate._inner)
|
||||
@@ -3539,9 +3773,24 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
||||
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||||
|
||||
# save the row ID choice that was made on the query builder and force it
|
||||
# to actually fetch the row ids because we need this for reranking
|
||||
with_row_ids = self._inner.get_with_row_id()
|
||||
req = fts_query._inner.to_query_request()
|
||||
blob_auto_row_id = False
|
||||
blob_paths: tuple[str, ...] = ()
|
||||
if self._table is not None and supports_blob_auto_row_id(self._table):
|
||||
schema = await self._table.schema()
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
self._table,
|
||||
schema,
|
||||
req.select,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if blob_auto_row_id:
|
||||
blob_paths = tuple(
|
||||
blob_v2_projection_sources(schema, req.select).keys()
|
||||
)
|
||||
self._blob_auto_row_id = blob_auto_row_id
|
||||
self._blob_paths = blob_paths
|
||||
|
||||
fts_query.with_row_id()
|
||||
vec_query.with_row_id()
|
||||
|
||||
@@ -3557,8 +3806,14 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query=fts_query.get_query(),
|
||||
reranker=self._reranker,
|
||||
limit=self._inner.get_limit(),
|
||||
with_row_ids=with_row_ids,
|
||||
with_row_ids=True,
|
||||
)
|
||||
if (
|
||||
not self._user_requested_row_id()
|
||||
and not blob_auto_row_id
|
||||
and "_rowid" in result.column_names
|
||||
):
|
||||
result = result.drop(["_rowid"])
|
||||
|
||||
return AsyncRecordBatchReader(result, max_batch_length=max_batch_length)
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ from lancedb._lancedb import (
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
IndexConfig,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
@@ -244,6 +245,23 @@ class RemoteTable(Table):
|
||||
"""List all the indices on the table"""
|
||||
return LOOP.run(self._table.list_indices())
|
||||
|
||||
def tokenize(
|
||||
self,
|
||||
query: str,
|
||||
*,
|
||||
column: Optional[str] = None,
|
||||
index_name: Optional[str] = None,
|
||||
) -> Iterable[FtsToken]:
|
||||
"""Tokenize a query using the tokenizer configured on an FTS index.
|
||||
|
||||
Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are
|
||||
rebuilt in the client process from index metadata, so the same tokenizer
|
||||
model files must exist locally.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._table.tokenize(query, column=column, index_name=index_name)
|
||||
)
|
||||
|
||||
def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]:
|
||||
"""List all the stats of a specified index"""
|
||||
return LOOP.run(self._table.index_stats(index_uuid))
|
||||
@@ -912,6 +930,10 @@ class RemoteTable(Table):
|
||||
"""Not supported on LanceDB Cloud."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""No-op on LanceDB Cloud (no local shard writers)."""
|
||||
return LOOP.run(self._table.close_lsm_writers())
|
||||
@@ -990,6 +1012,19 @@ class RemoteTable(Table):
|
||||
"migrate_v2_manifest_paths() is not supported on the LanceDB Cloud"
|
||||
)
|
||||
|
||||
def blob_columns(self) -> list[str]:
|
||||
raise NotImplementedError(
|
||||
"blob_columns() is not yet supported on the LanceDB Cloud"
|
||||
)
|
||||
|
||||
def fetch_blobs(self, column: str, row_ids) -> pa.LargeBinaryArray:
|
||||
raise NotImplementedError("fetch_blobs() is not supported on LanceDB Cloud")
|
||||
|
||||
def fetch_blob_files(self, column: str, row_ids):
|
||||
raise NotImplementedError(
|
||||
"fetch_blob_files() is not supported on LanceDB Cloud"
|
||||
)
|
||||
|
||||
def head(self, n=5) -> pa.Table:
|
||||
"""
|
||||
Return the first `n` rows of the table.
|
||||
|
||||
@@ -12,6 +12,7 @@ from .rrf import RRFReranker
|
||||
from .mrr import MRRReranker
|
||||
from .answerdotai import AnswerdotaiRerankers
|
||||
from .voyageai import VoyageAIReranker
|
||||
from .watsonx import WatsonxReranker
|
||||
|
||||
__all__ = [
|
||||
"Reranker",
|
||||
@@ -25,4 +26,5 @@ __all__ = [
|
||||
"AnswerdotaiRerankers",
|
||||
"VoyageAIReranker",
|
||||
"MRRReranker",
|
||||
"WatsonxReranker",
|
||||
]
|
||||
|
||||
@@ -156,9 +156,16 @@ class MRRReranker(Reranker):
|
||||
reciprocal_rank = 1.0 / rank
|
||||
mrr_score_map[result_id].append(reciprocal_rank)
|
||||
|
||||
# MRR averages the reciprocal rank across *all* ranking systems, treating
|
||||
# a system in which a document does not appear as a reciprocal rank of 0.
|
||||
# We therefore divide by the total number of systems, not by the number of
|
||||
# systems the document happens to appear in -- otherwise a document found
|
||||
# by a single ranking would outrank one ranked highly by every system,
|
||||
# defeating the purpose of fusing the rankings.
|
||||
num_systems = len(vector_results)
|
||||
final_mrr_scores = {}
|
||||
for result_id, reciprocal_ranks in mrr_score_map.items():
|
||||
mean_rr = np.mean(reciprocal_ranks)
|
||||
mean_rr = float(np.sum(reciprocal_ranks)) / num_systems
|
||||
final_mrr_scores[result_id] = mean_rr
|
||||
|
||||
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
import os
|
||||
from functools import cached_property
|
||||
from typing import Dict, Optional
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
from ..util import attempt_import_or_raise
|
||||
from .base import Reranker
|
||||
|
||||
DEFAULT_WATSONX_URL = "https://us-south.ml.cloud.ibm.com"
|
||||
|
||||
|
||||
class WatsonxReranker(Reranker):
|
||||
"""
|
||||
Reranks the results using the IBM watsonx.ai Rerank API.
|
||||
|
||||
Uses the ``ibm_watsonx_ai`` SDK (``Rerank.generate``) under the hood.
|
||||
|
||||
API Docs:
|
||||
https://cloud.ibm.com/docs/apis/watsonx-ai#text-rerank
|
||||
|
||||
Supported rerank models:
|
||||
https://dataplatform.cloud.ibm.com/docs/content/wsj/analyze-data/fm-models-embed.html?context=wx#rerank
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model_name : str, default "cross-encoder/ms-marco-minilm-l-12-v2"
|
||||
The ID of the rerank model to use.
|
||||
column : str, default "text"
|
||||
The name of the column to use as input to the reranker.
|
||||
top_n : int, optional
|
||||
Return only the top-n results. If ``None``, all results are returned.
|
||||
return_score : str, default "relevance"
|
||||
Options are ``"relevance"`` or ``"all"``.
|
||||
api_key : str, optional
|
||||
IBM Cloud API key. Falls back to the ``WATSONX_API_KEY`` environment
|
||||
variable when not provided.
|
||||
project_id : str, optional
|
||||
watsonx.ai project ID. Falls back to the ``WATSONX_PROJECT_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
``space_id`` — exactly one must be supplied.
|
||||
space_id : str, optional
|
||||
watsonx.ai deployment space ID. Falls back to the ``WATSONX_SPACE_ID``
|
||||
environment variable when not provided. Mutually exclusive with
|
||||
``project_id`` — exactly one must be supplied.
|
||||
url : str, optional
|
||||
watsonx.ai service URL. Defaults to
|
||||
``"https://us-south.ml.cloud.ibm.com"``.
|
||||
truncate_input_tokens : int, optional
|
||||
Truncate each input to this many tokens before scoring. Passed
|
||||
directly to the ``parameters`` dict of ``Rerank.generate``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str = "cross-encoder/ms-marco-minilm-l-12-v2",
|
||||
column: str = "text",
|
||||
top_n: Optional[int] = None,
|
||||
return_score: str = "relevance",
|
||||
api_key: Optional[str] = None,
|
||||
project_id: Optional[str] = None,
|
||||
space_id: Optional[str] = None,
|
||||
url: Optional[str] = None,
|
||||
truncate_input_tokens: Optional[int] = None,
|
||||
):
|
||||
super().__init__(return_score)
|
||||
self.model_name = model_name
|
||||
self.column = column
|
||||
self.top_n = top_n
|
||||
self.api_key = api_key
|
||||
self.project_id = project_id
|
||||
self.space_id = space_id
|
||||
self.url = url
|
||||
self.truncate_input_tokens = truncate_input_tokens
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"WatsonxReranker(model_name={self.model_name})"
|
||||
|
||||
@cached_property
|
||||
def _client(self):
|
||||
ibm_watsonx_ai = attempt_import_or_raise("ibm_watsonx_ai")
|
||||
ibm_watsonx_ai_foundation_models = attempt_import_or_raise(
|
||||
"ibm_watsonx_ai.foundation_models"
|
||||
)
|
||||
|
||||
# --- credentials ---
|
||||
api_key = self.api_key or os.environ.get("WATSONX_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"WATSONX_API_KEY not set. Either set it in your environment or "
|
||||
"pass it as `api_key` argument to WatsonxReranker."
|
||||
)
|
||||
credentials = ibm_watsonx_ai.Credentials(
|
||||
api_key=api_key,
|
||||
url=self.url or DEFAULT_WATSONX_URL,
|
||||
)
|
||||
|
||||
# --- project_id / space_id (exactly one required) ---
|
||||
project_id = self.project_id or os.environ.get("WATSONX_PROJECT_ID")
|
||||
space_id = self.space_id or os.environ.get("WATSONX_SPACE_ID")
|
||||
|
||||
if project_id and space_id:
|
||||
raise ValueError("Provide either `project_id` or `space_id`, not both.")
|
||||
if not project_id and not space_id:
|
||||
raise ValueError(
|
||||
"Either WATSONX_PROJECT_ID or WATSONX_SPACE_ID must be set. "
|
||||
"Pass one as an argument to WatsonxReranker or set the corresponding "
|
||||
"environment variable."
|
||||
)
|
||||
|
||||
kwargs: Dict = dict(model_id=self.model_name, credentials=credentials)
|
||||
if project_id:
|
||||
kwargs["project_id"] = project_id
|
||||
else:
|
||||
kwargs["space_id"] = space_id
|
||||
|
||||
return ibm_watsonx_ai_foundation_models.Rerank(**kwargs)
|
||||
|
||||
def _build_params(self) -> Dict:
|
||||
"""Build the ``parameters`` dict forwarded to ``Rerank.generate``."""
|
||||
return_options: Dict = {"inputs": True}
|
||||
if self.top_n is not None:
|
||||
return_options["top_n"] = self.top_n
|
||||
params: Dict = {"return_options": return_options}
|
||||
if self.truncate_input_tokens is not None:
|
||||
params["truncate_input_tokens"] = self.truncate_input_tokens
|
||||
return params
|
||||
|
||||
def _rerank(self, result_set: pa.Table, query: str) -> pa.Table:
|
||||
result_set = self._handle_empty_results(result_set)
|
||||
if len(result_set) == 0:
|
||||
return result_set
|
||||
|
||||
docs = result_set[self.column].to_pylist()
|
||||
response = self._client.generate(
|
||||
query=query,
|
||||
inputs=docs,
|
||||
params=self._build_params(),
|
||||
)
|
||||
results = response["results"]
|
||||
|
||||
indices, scores = zip(
|
||||
*[(result["index"], result["score"]) for result in results]
|
||||
)
|
||||
result_set = result_set.take(list(indices))
|
||||
result_set = result_set.append_column(
|
||||
"_relevance_score", pa.array(scores, type=pa.float32())
|
||||
)
|
||||
return result_set
|
||||
|
||||
def rerank_hybrid(
|
||||
self,
|
||||
query: str,
|
||||
vector_results: pa.Table,
|
||||
fts_results: pa.Table,
|
||||
) -> pa.Table:
|
||||
if self.score == "all":
|
||||
combined_results = self._merge_and_keep_scores(vector_results, fts_results)
|
||||
else:
|
||||
combined_results = self.merge_results(vector_results, fts_results)
|
||||
combined_results = self._rerank(combined_results, query)
|
||||
if self.score == "relevance":
|
||||
combined_results = self._keep_relevance_score(combined_results)
|
||||
return combined_results
|
||||
|
||||
def rerank_vector(self, query: str, vector_results: pa.Table) -> pa.Table:
|
||||
vector_results = self._rerank(vector_results, query)
|
||||
if self.score == "relevance":
|
||||
vector_results = vector_results.drop_columns(["_distance"])
|
||||
return vector_results
|
||||
|
||||
def rerank_fts(self, query: str, fts_results: pa.Table) -> pa.Table:
|
||||
fts_results = self._rerank(fts_results, query)
|
||||
if self.score == "relevance":
|
||||
fts_results = fts_results.drop_columns(["_score"])
|
||||
return fts_results
|
||||
@@ -2,10 +2,134 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
"""Schema related utilities."""
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
|
||||
import pyarrow as pa
|
||||
|
||||
_BLOB_EXTENSION_NAME = "lance.blob.v2"
|
||||
_BLOB_V1_KEY = "lance-encoding:blob"
|
||||
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
|
||||
|
||||
|
||||
class BlobType(pa.ExtensionType):
|
||||
"""PyArrow extension type for a Lance blob v2 column.
|
||||
|
||||
Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files`
|
||||
for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
storage_type = pa.struct(
|
||||
[
|
||||
pa.field("data", pa.large_binary(), nullable=True),
|
||||
pa.field("uri", pa.utf8(), nullable=True),
|
||||
pa.field("position", pa.uint64(), nullable=True),
|
||||
pa.field("size", pa.uint64(), nullable=True),
|
||||
]
|
||||
)
|
||||
super().__init__(storage_type, _BLOB_EXTENSION_NAME)
|
||||
|
||||
def __arrow_ext_serialize__(self) -> bytes:
|
||||
return b""
|
||||
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(
|
||||
cls, storage_type: pa.DataType, serialized: bytes
|
||||
) -> "BlobType":
|
||||
return cls()
|
||||
|
||||
def __reduce__(self):
|
||||
# Ensure pickle round-trips on older pyarrow (apache/arrow#35599).
|
||||
return type(self).__arrow_ext_deserialize__, (
|
||||
self.storage_type,
|
||||
self.__arrow_ext_serialize__(),
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
pa.register_extension_type(BlobType()) # type: ignore[arg-type]
|
||||
except pa.ArrowKeyError:
|
||||
pass
|
||||
|
||||
|
||||
def _metadata_value(metadata: dict, key: str):
|
||||
return metadata.get(key.encode()) or metadata.get(key)
|
||||
|
||||
|
||||
def _metadata_marks_blob_v2(metadata: dict) -> bool:
|
||||
if not metadata:
|
||||
return False
|
||||
|
||||
extension_name = _metadata_value(metadata, _ARROW_EXT_NAME_KEY)
|
||||
return extension_name in (_BLOB_EXTENSION_NAME, _BLOB_EXTENSION_NAME.encode())
|
||||
|
||||
|
||||
def _metadata_marks_legacy_blob(metadata: dict) -> bool:
|
||||
if not metadata:
|
||||
return False
|
||||
|
||||
return _metadata_value(metadata, _BLOB_V1_KEY) in ("true", b"true")
|
||||
|
||||
|
||||
def is_blob_v2_field(field: pa.Field) -> bool:
|
||||
"""Return True if `field` declares a blob v2 extension column."""
|
||||
field_type = field.type
|
||||
if (
|
||||
isinstance(field_type, pa.ExtensionType)
|
||||
and field_type.extension_name == _BLOB_EXTENSION_NAME
|
||||
):
|
||||
return True
|
||||
return _metadata_marks_blob_v2(field.metadata or {})
|
||||
|
||||
|
||||
def is_blob_like_field(field: pa.Field) -> bool:
|
||||
"""Blob detection for ``to_pandas(blob_mode=...)`` and scanner paths only.
|
||||
|
||||
Matches v2 extension fields on table schema, legacy ``lance-encoding:blob``
|
||||
storage columns, and v2 query descriptor fields (the engine tags those with
|
||||
the same metadata). Not used for fetch or auto ``_rowid``.
|
||||
"""
|
||||
return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {})
|
||||
|
||||
|
||||
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]:
|
||||
paths: list[str] = []
|
||||
|
||||
def walk(fields, prefix: str) -> None:
|
||||
for field in fields:
|
||||
path = f"{prefix}.{field.name}" if prefix else field.name
|
||||
if is_blob(field):
|
||||
paths.append(path)
|
||||
elif pa.types.is_struct(field.type):
|
||||
walk(field.type, path)
|
||||
elif (
|
||||
pa.types.is_list(field.type)
|
||||
or pa.types.is_large_list(field.type)
|
||||
or pa.types.is_fixed_size_list(field.type)
|
||||
):
|
||||
walk([field.type.value_field], path)
|
||||
|
||||
walk(schema, "")
|
||||
return paths
|
||||
|
||||
|
||||
def blob_column_paths(schema: pa.Schema) -> list[str]:
|
||||
"""Dotted paths of blob-like columns (v2 extension or legacy metadata)."""
|
||||
return _collect_blob_paths(schema, is_blob_like_field)
|
||||
|
||||
|
||||
def blob_v2_column_paths(schema: pa.Schema) -> list[str]:
|
||||
return _collect_blob_paths(schema, is_blob_v2_field)
|
||||
|
||||
|
||||
def schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return bool(blob_column_paths(schema))
|
||||
|
||||
|
||||
def blob(name: str, nullable: bool = True) -> pa.Field:
|
||||
"""Create a Lance blob v2 column field."""
|
||||
return pa.field(name, BlobType(), nullable=nullable)
|
||||
|
||||
|
||||
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||
"""A help function to create a vector type.
|
||||
|
||||
@@ -0,0 +1,607 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Elastic streaming dataloader for PyTorch.
|
||||
|
||||
Provides StreamingDataset, a PyTorch IterableDataset that guarantees:
|
||||
|
||||
- **Elastic determinism**: for a fixed (num_splits, shuffle_seed, epoch) the set
|
||||
of samples that forms each global training step is identical regardless of
|
||||
world_size or num_workers.
|
||||
- **Resumability**: state_dict / load_state_dict capture per-split consumption
|
||||
counts so training can resume from an exact mid-epoch position even when the
|
||||
distributed topology changes between runs.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from multiprocessing import RawArray
|
||||
from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
from torch.utils.data import IterableDataset, get_worker_info
|
||||
|
||||
from .permutation import (
|
||||
Permutation,
|
||||
Transforms,
|
||||
permutation_builder,
|
||||
_table_from_pickle_state,
|
||||
_table_to_pickle_state,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Multiplier used to combine shuffle_seed and epoch into a single permutation
|
||||
# seed. Chosen to be a large prime so different (seed, epoch) pairs produce
|
||||
# distinct seeds for any practically encountered epoch count.
|
||||
_EPOCH_PRIME = 100003
|
||||
|
||||
DEFAULT_READ_BATCH_SIZE = 64
|
||||
DEFAULT_PREFETCH_BATCHES = 4
|
||||
|
||||
|
||||
class StreamingDataset(IterableDataset):
|
||||
"""An elastic, resumable PyTorch IterableDataset backed by a LanceDB table.
|
||||
|
||||
The table is partitioned into ``num_splits`` fixed splits using a
|
||||
deterministic random shuffle controlled by ``shuffle_seed`` and ``epoch``.
|
||||
Each rank is assigned a contiguous block of splits, and within a rank each
|
||||
DataLoader worker is assigned a contiguous sub-block. Samples are yielded
|
||||
by round-robining over the assigned splits, one sample per split per cycle.
|
||||
|
||||
Internally ``__iter__`` runs a two-stage pipeline:
|
||||
|
||||
- **Stage 1 (I/O)**: one thread pool with ``num_splits * prefetch_batches``
|
||||
workers fetches raw ``RecordBatch`` objects from LanceDB in parallel
|
||||
across all splits and places them in a per-split raw-batch queue.
|
||||
- **Stage 2 (transform)**: a second thread pool with ``os.cpu_count()``
|
||||
workers picks up raw batches, applies the transform, and places the
|
||||
results in a per-split cooked-row queue.
|
||||
|
||||
The main thread round-robins over the cooked queues, yielding one row per
|
||||
split per cycle.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
table:
|
||||
LanceDB table to stream from.
|
||||
num_splits:
|
||||
Number of fixed splits to partition the table into. Must be divisible
|
||||
by ``world_size``. When used with DataLoader workers it must also be
|
||||
divisible by ``world_size * num_workers``. Defaults to ``world_size``.
|
||||
If the row count (after any ``filter``) is not evenly divisible by
|
||||
``num_splits``, the surplus rows — at most ``num_splits - 1`` per epoch
|
||||
— are silently dropped to keep all splits the same length.
|
||||
shuffle:
|
||||
Whether to randomly assign rows to splits. When ``True`` (the
|
||||
default) rows are shuffled using ``shuffle_seed`` and ``epoch``.
|
||||
When ``False`` rows are divided into splits sequentially in storage
|
||||
order, which can be useful for deterministic debugging or evaluation.
|
||||
shuffle_seed:
|
||||
Base seed for the random permutation. Combined with ``epoch`` so
|
||||
each epoch produces a different ordering. Pass ``None`` to generate
|
||||
a random seed at construction time.
|
||||
epoch:
|
||||
Current training epoch. Combined with ``shuffle_seed`` so that each
|
||||
epoch produces a different sample ordering.
|
||||
rank:
|
||||
This process's rank in the distributed training group.
|
||||
world_size:
|
||||
Total number of processes in the distributed training group.
|
||||
read_batch_size:
|
||||
Number of rows fetched from each split in a single ``take_offsets``
|
||||
call. Larger values amortise per-request overhead (critical on object
|
||||
storage) at the cost of higher memory usage per split buffer. Defaults
|
||||
to ``DEFAULT_READ_BATCH_SIZE`` (64).
|
||||
prefetch_batches:
|
||||
Number of I/O batches to keep in flight per split. Higher values
|
||||
overlap storage latency with transform and training compute at the cost
|
||||
of more memory and threads. Defaults to ``DEFAULT_PREFETCH_BATCHES``
|
||||
(4).
|
||||
columns:
|
||||
Optional list of column names to read. When set, only those columns
|
||||
are fetched from storage; all others are omitted. ``None`` (the
|
||||
default) reads every column.
|
||||
shuffle_clump_size:
|
||||
When set, rows are shuffled in contiguous groups of this size rather
|
||||
than individually. Larger clumps improve I/O locality (important on
|
||||
object storage) at the cost of reduced randomness. ``None`` (the
|
||||
default) shuffles rows individually.
|
||||
filter:
|
||||
Optional SQL filter expression (e.g. ``"label = 'dog'"``). Only rows
|
||||
that satisfy the predicate are included in the permutation. The filter
|
||||
is applied during permutation construction so split sizes reflect the
|
||||
filtered row count.
|
||||
transform:
|
||||
Optional callable applied to each ``pyarrow.RecordBatch`` before rows
|
||||
are yielded. Receives one batch at a time and must return an iterable
|
||||
whose length equals the number of rows in the batch. When ``None``
|
||||
(the default) rows are returned as plain Python dicts.
|
||||
worker_info_override:
|
||||
If set, used in place of ``torch.utils.data.get_worker_info()`` to
|
||||
determine the DataLoader worker assignment. Intended for unit tests
|
||||
that need to simulate multiple workers without spawning real processes.
|
||||
If both this and the real worker info are non-None a warning is logged
|
||||
and the override takes precedence.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
table,
|
||||
*,
|
||||
num_splits: Optional[int] = None,
|
||||
shuffle: bool = True,
|
||||
shuffle_seed: Optional[int] = 0,
|
||||
epoch: int = 0,
|
||||
rank: int = 0,
|
||||
world_size: int = 1,
|
||||
read_batch_size: int = DEFAULT_READ_BATCH_SIZE,
|
||||
prefetch_batches: int = DEFAULT_PREFETCH_BATCHES,
|
||||
columns: Optional[list[str]] = None,
|
||||
shuffle_clump_size: Optional[int] = None,
|
||||
filter: Optional[str] = None,
|
||||
transform: Optional[Callable] = None,
|
||||
connection_factory: Optional[Callable[[str], Any]] = None,
|
||||
worker_info_override=None,
|
||||
):
|
||||
super().__init__()
|
||||
if num_splits is None:
|
||||
num_splits = world_size
|
||||
if shuffle_seed is None:
|
||||
shuffle_seed = random.randrange(2**32)
|
||||
if num_splits % world_size != 0:
|
||||
raise ValueError(
|
||||
f"num_splits ({num_splits}) must be divisible by "
|
||||
f"world_size ({world_size})"
|
||||
)
|
||||
|
||||
self._table = table
|
||||
self._num_splits = num_splits
|
||||
self._shuffle = shuffle
|
||||
self._shuffle_seed = shuffle_seed
|
||||
self._epoch = epoch
|
||||
self._rank = rank
|
||||
self._world_size = world_size
|
||||
self._read_batch_size = read_batch_size
|
||||
self._prefetch_batches = prefetch_batches
|
||||
self._columns = columns
|
||||
self._shuffle_clump_size = shuffle_clump_size
|
||||
self._filter = filter
|
||||
self._transform = transform
|
||||
self._connection_factory = connection_factory
|
||||
self._worker_info_override = worker_info_override
|
||||
|
||||
# Live references to pipeline state, set only while __iter__ is running
|
||||
# in the same process. Used by the observability properties when the
|
||||
# DataLoader runs with num_workers=0.
|
||||
self._raw_batches_ref: Optional[list[deque]] = None
|
||||
self._cooked_ref: Optional[list[deque]] = None
|
||||
self._fetch_head_ref: Optional[list[int]] = None
|
||||
self._split_sizes_ref: Optional[list[int]] = None
|
||||
self._local_consumed_ref: Optional[list[int]] = None
|
||||
|
||||
# Shared-memory counters written by __iter__ (which may run in a
|
||||
# DataLoader worker process) and read by the observability properties
|
||||
# in the main process. RawArray is picklable via the forkserver
|
||||
# reduction protocol so it survives the dataset pickle round-trip.
|
||||
# Layout: [unscanned_rows, raw_rows, cooked_rows, consumed_rows,
|
||||
# bytes_loaded, fetch_time_us, transform_time_us]
|
||||
self._worker_stats: RawArray = RawArray(ctypes.c_int64, 7)
|
||||
|
||||
# Cumulative bytes of Arrow buffer data fetched across all iterations.
|
||||
self._bytes_loaded: int = 0
|
||||
# Cumulative seconds spent in LanceDB I/O and in transform functions.
|
||||
self._fetch_time: float = 0.0
|
||||
self._transform_time: float = 0.0
|
||||
|
||||
# Number of samples each split has already been consumed. At global
|
||||
# step boundaries all splits have consumed this many samples, so a
|
||||
# single scalar captures the topology-independent checkpoint state.
|
||||
self._resume_offset: int = 0
|
||||
|
||||
# Build the permutation table once, deterministically.
|
||||
builder = permutation_builder(table)
|
||||
if filter is not None:
|
||||
builder = builder.filter(filter)
|
||||
if shuffle:
|
||||
perm_seed = shuffle_seed + epoch * _EPOCH_PRIME
|
||||
self._perm_table = builder.split_random(
|
||||
fixed=num_splits, seed=perm_seed, clump_size=shuffle_clump_size
|
||||
).execute()
|
||||
else:
|
||||
self._perm_table = builder.split_sequential(fixed=num_splits).execute()
|
||||
|
||||
# Contiguous block of global split indices assigned to this rank.
|
||||
splits_per_rank = num_splits // world_size
|
||||
rank_start = rank * splits_per_rank
|
||||
self._rank_splits: list[int] = list(
|
||||
range(rank_start, rank_start + splits_per_rank)
|
||||
)
|
||||
|
||||
def _resolve_my_splits(self) -> list[int]:
|
||||
"""Return the split indices this instance should read in __iter__."""
|
||||
torch_worker_info = get_worker_info()
|
||||
if self._worker_info_override is not None:
|
||||
if torch_worker_info is not None:
|
||||
logger.warning(
|
||||
"worker_info_override is set but get_worker_info() also returned a "
|
||||
"non-None value; ignoring the real torch worker info and using the "
|
||||
"override instead. This may lead to duplicated or incorrect data "
|
||||
"from the dataset."
|
||||
)
|
||||
worker_info = self._worker_info_override
|
||||
else:
|
||||
worker_info = torch_worker_info
|
||||
|
||||
if worker_info is None:
|
||||
return self._rank_splits
|
||||
|
||||
num_workers: int = worker_info.num_workers
|
||||
worker_id: int = worker_info.id
|
||||
n_rank_splits = len(self._rank_splits)
|
||||
if n_rank_splits % num_workers != 0:
|
||||
raise ValueError(
|
||||
f"Number of rank splits ({n_rank_splits}) must be divisible by "
|
||||
f"num_workers ({num_workers})"
|
||||
)
|
||||
splits_per_worker = n_rank_splits // num_workers
|
||||
start = worker_id * splits_per_worker
|
||||
return self._rank_splits[start : start + splits_per_worker]
|
||||
|
||||
def __iter__(self) -> Iterator[dict[str, Any]]:
|
||||
if self._raw_batches_ref is not None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataset does not support concurrent iteration. "
|
||||
"Only one active iterator per dataset instance is allowed."
|
||||
)
|
||||
my_splits = self._resolve_my_splits()
|
||||
if not my_splits:
|
||||
return
|
||||
|
||||
# Set identity transform on each Permutation so __getitems__ returns
|
||||
# the raw RecordBatch. Stage 2 applies the real transform.
|
||||
permutations: list[Permutation] = []
|
||||
for split_idx in my_splits:
|
||||
perm = Permutation.from_tables(
|
||||
self._table, self._perm_table, split=split_idx
|
||||
)
|
||||
if self._columns is not None:
|
||||
perm = perm.select_columns(self._columns)
|
||||
perm = perm.with_transform(lambda batch: batch)
|
||||
if self._resume_offset > 0:
|
||||
perm = perm.with_skip(self._resume_offset)
|
||||
permutations.append(perm)
|
||||
|
||||
n = len(permutations)
|
||||
split_sizes = [perm.num_rows for perm in permutations]
|
||||
initial_offset = self._resume_offset
|
||||
local_consumed = [0] * n
|
||||
|
||||
batch_size = self._read_batch_size
|
||||
max_prefetch = self._prefetch_batches
|
||||
cpu_workers = os.cpu_count() or 1
|
||||
final_transform = (
|
||||
self._transform if self._transform is not None else Transforms.arrow2python
|
||||
)
|
||||
|
||||
# Per-split pipeline state.
|
||||
fetch_head = [0] * n
|
||||
io_pending = [deque() for _ in range(n)] # Future[RecordBatch]
|
||||
raw_batches = [deque() for _ in range(n)] # RecordBatch — fetched, awaiting tx
|
||||
tx_pending = [deque() for _ in range(n)] # Future[list[Any]]
|
||||
cooked = [deque() for _ in range(n)] # rows ready to yield
|
||||
|
||||
# Limit simultaneous transforms to cpu_workers across all splits.
|
||||
tx_semaphore = threading.Semaphore(cpu_workers)
|
||||
|
||||
# ── Stage 1 helpers ───────────────────────────────────────────────────
|
||||
|
||||
def _io_call(perm, indices):
|
||||
t0 = time.perf_counter()
|
||||
batch = perm.__getitems__(indices)
|
||||
self._bytes_loaded += batch.nbytes
|
||||
self._fetch_time += time.perf_counter() - t0
|
||||
return batch
|
||||
|
||||
def _submit_io(i: int) -> None:
|
||||
remaining = split_sizes[i] - fetch_head[i]
|
||||
if remaining <= 0:
|
||||
return
|
||||
fetch = min(batch_size, remaining)
|
||||
start = fetch_head[i]
|
||||
fetch_head[i] += fetch
|
||||
perm_i = permutations[i]
|
||||
indices = list(range(start, start + fetch))
|
||||
io_pending[i].append(io_pool.submit(_io_call, perm_i, indices))
|
||||
|
||||
def _fill_io(i: int) -> None:
|
||||
while len(io_pending[i]) < max_prefetch and fetch_head[i] < split_sizes[i]:
|
||||
_submit_io(i)
|
||||
|
||||
def _drain_io(i: int) -> None:
|
||||
"""Move completed I/O futures into raw_batches non-blockingly."""
|
||||
while io_pending[i] and io_pending[i][0].done():
|
||||
raw_batches[i].append(io_pending[i].popleft().result())
|
||||
|
||||
# ── Stage 2 helpers ───────────────────────────────────────────────────
|
||||
|
||||
def _tx_call_guarded(batch):
|
||||
try:
|
||||
t0 = time.perf_counter()
|
||||
result = final_transform(batch)
|
||||
self._transform_time += time.perf_counter() - t0
|
||||
return result
|
||||
finally:
|
||||
tx_semaphore.release()
|
||||
|
||||
def _try_submit_tx(i: int) -> None:
|
||||
"""Submit transforms for raw_batches[i] up to available capacity."""
|
||||
while raw_batches[i] and tx_semaphore.acquire(blocking=False):
|
||||
batch = raw_batches[i].popleft()
|
||||
tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch))
|
||||
|
||||
def _drain_tx(i: int) -> None:
|
||||
"""Move completed transform futures into cooked non-blockingly."""
|
||||
while tx_pending[i] and tx_pending[i][0].done():
|
||||
cooked[i].extend(tx_pending[i].popleft().result())
|
||||
|
||||
# ── Combined advance ──────────────────────────────────────────────────
|
||||
|
||||
def _advance(i: int) -> None:
|
||||
"""Non-blocking pipeline pump for split i."""
|
||||
_drain_io(i)
|
||||
_drain_tx(i)
|
||||
_try_submit_tx(i)
|
||||
_fill_io(i)
|
||||
|
||||
def _ensure_cooked(i: int) -> None:
|
||||
"""Ensure cooked[i] has at least one row, blocking if necessary."""
|
||||
_advance(i)
|
||||
while not cooked[i]:
|
||||
if tx_pending[i]:
|
||||
# Wait for the oldest in-flight transform.
|
||||
cooked[i].extend(tx_pending[i].popleft().result())
|
||||
_advance(i)
|
||||
elif raw_batches[i]:
|
||||
# Acquire a transform slot (may block briefly if all
|
||||
# cpu_workers are busy with other splits).
|
||||
tx_semaphore.acquire()
|
||||
batch = raw_batches[i].popleft()
|
||||
tx_pending[i].append(tx_pool.submit(_tx_call_guarded, batch))
|
||||
elif io_pending[i]:
|
||||
# Block on the oldest in-flight I/O fetch.
|
||||
raw_batches[i].append(io_pending[i].popleft().result())
|
||||
_advance(i)
|
||||
else:
|
||||
break # split exhausted
|
||||
|
||||
# ── Main loop ─────────────────────────────────────────────────────────
|
||||
|
||||
with ThreadPoolExecutor(max_workers=n * max_prefetch) as io_pool:
|
||||
with ThreadPoolExecutor(max_workers=cpu_workers) as tx_pool:
|
||||
self._raw_batches_ref = raw_batches
|
||||
self._cooked_ref = cooked
|
||||
self._fetch_head_ref = fetch_head
|
||||
self._split_sizes_ref = split_sizes
|
||||
self._local_consumed_ref = local_consumed
|
||||
try:
|
||||
for i in range(n):
|
||||
_fill_io(i)
|
||||
|
||||
while True:
|
||||
# Stop when any split is exhausted (all exhaust
|
||||
# simultaneously: equal split sizes + round-robin).
|
||||
if any(local_consumed[i] >= split_sizes[i] for i in range(n)):
|
||||
break
|
||||
|
||||
for i in range(n):
|
||||
_ensure_cooked(i)
|
||||
row = cooked[i].popleft()
|
||||
local_consumed[i] += 1
|
||||
_advance(i)
|
||||
|
||||
# After the last split in each cycle: update the
|
||||
# global offset and refresh the shared-memory stats
|
||||
# so the main process can observe pipeline depth
|
||||
# even when __iter__ runs in a worker process.
|
||||
if i == n - 1:
|
||||
self._resume_offset = initial_offset + local_consumed[i]
|
||||
ws = self._worker_stats
|
||||
ws[0] = sum(
|
||||
split_sizes[j] - fetch_head[j] for j in range(n)
|
||||
)
|
||||
ws[1] = sum(
|
||||
batch.num_rows for q in raw_batches for batch in q
|
||||
)
|
||||
ws[2] = sum(len(q) for q in cooked)
|
||||
ws[3] = sum(local_consumed)
|
||||
ws[4] = self._bytes_loaded
|
||||
ws[5] = int(self._fetch_time * 1_000_000)
|
||||
ws[6] = int(self._transform_time * 1_000_000)
|
||||
|
||||
yield row
|
||||
finally:
|
||||
self._raw_batches_ref = None
|
||||
self._cooked_ref = None
|
||||
self._fetch_head_ref = None
|
||||
self._split_sizes_ref = None
|
||||
self._local_consumed_ref = None
|
||||
|
||||
@property
|
||||
def bytes_loaded(self) -> int:
|
||||
"""Cumulative bytes of raw Arrow buffer data fetched from storage.
|
||||
|
||||
Measured on the ``RecordBatch`` before any transform is applied, so
|
||||
the value reflects actual I/O rather than the size of transformed
|
||||
output. Accumulates across multiple iterations of the same dataset
|
||||
instance and is never reset automatically.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._bytes_loaded
|
||||
return int(self._worker_stats[4])
|
||||
|
||||
@property
|
||||
def fetch_time(self) -> float:
|
||||
"""Cumulative seconds spent waiting for data from LanceDB.
|
||||
|
||||
Measured per batch in the Stage 1 I/O threads as the total elapsed
|
||||
time of the ``take_offsets`` call. Accumulates across all splits and
|
||||
all iterations.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._fetch_time
|
||||
return self._worker_stats[5] / 1_000_000
|
||||
|
||||
@property
|
||||
def transform_time(self) -> float:
|
||||
"""Cumulative seconds spent applying the transform.
|
||||
|
||||
Measured per batch in the Stage 2 transform threads as the elapsed
|
||||
time inside the transform callable (or the default ``arrow2python``
|
||||
conversion when no transform is set). Accumulates across all splits
|
||||
and all iterations.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return self._transform_time
|
||||
return self._worker_stats[6] / 1_000_000
|
||||
|
||||
@property
|
||||
def raw_queue_depth(self) -> int:
|
||||
"""Number of raw rows waiting for a transform thread across all splits.
|
||||
|
||||
A persistently non-zero value means Stage 2 (transform) is the
|
||||
bottleneck: I/O is completing faster than transforms can consume
|
||||
batches. Returns 0 when not iterating.
|
||||
"""
|
||||
if self._raw_batches_ref is not None:
|
||||
return sum(batch.num_rows for q in self._raw_batches_ref for batch in q)
|
||||
return int(self._worker_stats[1])
|
||||
|
||||
@property
|
||||
def prefetch_queue_depth(self) -> int:
|
||||
"""Number of rows transformed and ready to yield across all splits.
|
||||
|
||||
Counts rows whose transform has completed and are sitting in memory
|
||||
waiting for the main thread — rows that can be handed off with no
|
||||
I/O or CPU wait. Returns 0 when not iterating.
|
||||
"""
|
||||
if self._cooked_ref is not None:
|
||||
return sum(len(q) for q in self._cooked_ref)
|
||||
return int(self._worker_stats[2])
|
||||
|
||||
@property
|
||||
def unscanned_rows(self) -> int:
|
||||
"""Number of rows not yet submitted to the I/O stage across all splits.
|
||||
|
||||
Decreases as the I/O stage submits fetch requests. When this reaches
|
||||
zero all data has been requested from storage (though it may not have
|
||||
arrived yet). Returns 0 when not iterating.
|
||||
"""
|
||||
if self._fetch_head_ref is not None:
|
||||
return sum(
|
||||
size - head
|
||||
for size, head in zip(self._split_sizes_ref, self._fetch_head_ref)
|
||||
)
|
||||
return int(self._worker_stats[0])
|
||||
|
||||
@property
|
||||
def consumed_rows(self) -> int:
|
||||
"""Number of rows already yielded to the caller across all splits.
|
||||
|
||||
Monotonically increases throughout iteration. Returns 0 when not
|
||||
iterating.
|
||||
"""
|
||||
if self._local_consumed_ref is not None:
|
||||
return sum(self._local_consumed_ref)
|
||||
return int(self._worker_stats[3])
|
||||
|
||||
def __getstate__(self):
|
||||
"""Support pickling for multi-worker DataLoader (forkserver / spawn).
|
||||
|
||||
The live LanceDB table object contains non-picklable connection state
|
||||
(sockets, Rust-backed PyO3 objects). If a ``connection_factory`` was
|
||||
supplied only the table name is serialised; the factory is called in
|
||||
the worker to reopen the connection without embedding any credentials.
|
||||
Without a factory the table's own picklable reopen state is captured
|
||||
via ``_table_to_pickle_state`` (mirrors the ``Permutation`` approach).
|
||||
"""
|
||||
state = self.__dict__.copy()
|
||||
# _table: replace with reconnect info (credentials must not be embedded).
|
||||
state["_table_name"] = self._table.name
|
||||
if self._connection_factory is not None:
|
||||
state["_table"] = None
|
||||
else:
|
||||
state["_table"] = _table_to_pickle_state(self._table)
|
||||
# _perm_table: always in-memory; serialise as Arrow data (mirrors
|
||||
# how Permutation.__getstate__ handles its permutation_table).
|
||||
state["_perm_table"] = (
|
||||
self._perm_table.name,
|
||||
self._perm_table.to_arrow(),
|
||||
)
|
||||
for key in (
|
||||
"_raw_batches_ref",
|
||||
"_cooked_ref",
|
||||
"_fetch_head_ref",
|
||||
"_split_sizes_ref",
|
||||
"_local_consumed_ref",
|
||||
):
|
||||
state[key] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""Reconnect to LanceDB after unpickling in a worker process."""
|
||||
from . import connect as _connect
|
||||
|
||||
table_name = state.pop("_table_name")
|
||||
table_state = state.pop("_table")
|
||||
perm_name, perm_data = state.pop("_perm_table")
|
||||
self.__dict__.update(state)
|
||||
if self._connection_factory is not None:
|
||||
self._table = self._connection_factory(table_name)
|
||||
else:
|
||||
self._table = _table_from_pickle_state(table_state)
|
||||
self._perm_table = _connect("memory://").create_table(perm_name, perm_data)
|
||||
|
||||
def state_dict(self) -> dict:
|
||||
"""Snapshot the dataset's consumption state.
|
||||
|
||||
The returned dict is topology-independent: at global step boundaries
|
||||
every split has been consumed the same number of times (by the
|
||||
round-robin design), so the per-split count is a single uniform value
|
||||
that is identical across all ranks and DataLoader workers.
|
||||
"""
|
||||
return {
|
||||
"shuffle_seed": self._shuffle_seed,
|
||||
"num_splits": self._num_splits,
|
||||
"epoch": self._epoch,
|
||||
"samples_consumed_per_split": [self._resume_offset] * self._num_splits,
|
||||
}
|
||||
|
||||
def load_state_dict(self, state: dict) -> None:
|
||||
"""Resume from a previously snapshotted state.
|
||||
|
||||
Raises ``ValueError`` if ``num_splits`` or ``shuffle_seed`` differ
|
||||
from the checkpoint, since a different split structure or shuffle order
|
||||
makes mid-epoch resumption meaningless.
|
||||
"""
|
||||
if state["num_splits"] != self._num_splits:
|
||||
raise ValueError(
|
||||
f"num_splits mismatch: checkpoint has {state['num_splits']}, "
|
||||
f"current dataset has {self._num_splits}"
|
||||
)
|
||||
if state["shuffle_seed"] != self._shuffle_seed:
|
||||
raise ValueError(
|
||||
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
|
||||
f"current dataset has {self._shuffle_seed}"
|
||||
)
|
||||
consumed = state["samples_consumed_per_split"]
|
||||
# All entries are equal at step boundaries; use the first.
|
||||
if isinstance(consumed, list):
|
||||
self._resume_offset = consumed[0] if consumed else 0
|
||||
else:
|
||||
self._resume_offset = int(consumed)
|
||||
+226
-35
@@ -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 (
|
||||
@@ -88,10 +96,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 +105,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",
|
||||
@@ -185,6 +173,7 @@ if TYPE_CHECKING:
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
UpdateResult,
|
||||
@@ -651,6 +640,16 @@ def _append_vector_columns(
|
||||
col_data = func.compute_source_embeddings_with_retry(
|
||||
batch[conf.source_column]
|
||||
)
|
||||
# Replace vectors with wrong length (including empty lists
|
||||
# returned for inputs like empty strings) with None so that
|
||||
# _handle_bad_vectors can process them according to the
|
||||
# on_bad_vectors policy instead of crashing when PyArrow
|
||||
# tries to cast them into a fixed-size list array.
|
||||
expected_ndims = conf.function.ndims()
|
||||
col_data = [
|
||||
v if v is not None and len(v) == expected_ndims else None
|
||||
for v in col_data
|
||||
]
|
||||
if no_vector_column:
|
||||
batch = batch.append_column(
|
||||
schema.field(vector_column),
|
||||
@@ -1149,6 +1148,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"
|
||||
@@ -1513,6 +1514,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,
|
||||
@@ -1776,6 +1802,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]:
|
||||
"""
|
||||
@@ -2194,6 +2238,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.
|
||||
@@ -2389,9 +2446,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"
|
||||
@@ -2400,6 +2462,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:
|
||||
@@ -3700,6 +3765,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
|
||||
@@ -3749,6 +3834,11 @@ class LanceTable(Table):
|
||||
[`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``. See
|
||||
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""Close cached MemWAL shard writers. See
|
||||
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
|
||||
@@ -4020,7 +4110,16 @@ def _handle_bad_vector_column(
|
||||
dim = _infer_vector_dim(vec_arr)
|
||||
if dim is None:
|
||||
return data
|
||||
has_wrong_dim = pc.not_equal(pc.list_value_length(vec_arr), dim)
|
||||
|
||||
is_null = pc.is_null(vec_arr)
|
||||
# pc.list_value_length returns null for null list entries, so
|
||||
# pc.not_equal(null, dim) also returns null. Use or_kleene so that
|
||||
# True OR null = True (Kleene three-valued logic), ensuring null vectors
|
||||
# are counted as wrong-dim.
|
||||
has_wrong_dim = pc.or_kleene(
|
||||
is_null,
|
||||
pc.not_equal(pc.list_value_length(vec_arr), dim),
|
||||
)
|
||||
|
||||
has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
|
||||
|
||||
@@ -4055,17 +4154,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])
|
||||
@@ -4398,6 +4538,17 @@ class AsyncTable:
|
||||
"""
|
||||
await self._inner.unset_lsm_write_spec()
|
||||
|
||||
async def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the LsmWriteSpec currently installed on this table.
|
||||
|
||||
Returns ``None`` when the MemWAL LSM write path is not enabled (no
|
||||
spec has been set, or it was removed with `unset_lsm_write_spec`).
|
||||
The returned spec — including its ``maintained_indexes`` and
|
||||
``writer_config_defaults`` — mirrors what was passed to
|
||||
`set_lsm_write_spec`.
|
||||
"""
|
||||
return await self._inner.get_lsm_write_spec()
|
||||
|
||||
async def close_lsm_writers(self) -> None:
|
||||
"""Drain and close any cached MemWAL shard writers for this table.
|
||||
|
||||
@@ -4504,14 +4655,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:
|
||||
@@ -5612,6 +5767,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.
|
||||
@@ -5713,6 +5886,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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -177,7 +177,10 @@ def flatten_columns(tbl: pa.Table, flatten: Optional[Union[int, bool]] = None):
|
||||
continue
|
||||
else:
|
||||
break
|
||||
elif isinstance(flatten, int):
|
||||
# `bool` is a subclass of `int`, so guard against it explicitly: `flatten=False`
|
||||
# (and `None`) must mean "do not flatten" rather than falling into the integer
|
||||
# branch and raising on the `flatten <= 0` check.
|
||||
elif isinstance(flatten, int) and not isinstance(flatten, bool):
|
||||
if flatten <= 0:
|
||||
raise ValueError(
|
||||
"Please specify a positive integer for flatten or the boolean "
|
||||
@@ -515,3 +518,38 @@ def batch_to_tensor_rows(batch: pa.RecordBatch):
|
||||
stacked = torch.tensor(numpy.column_stack(columns))
|
||||
rows = list(stacked.unbind(dim=0))
|
||||
return rows
|
||||
|
||||
|
||||
def batch_to_tensor_dict(batch: pa.RecordBatch):
|
||||
"""
|
||||
Convert a PyArrow RecordBatch to a list of per-row dicts of PyTorch Tensors.
|
||||
|
||||
Each column is first converted to a 1-D tensor (zero-copy via DLPack), then
|
||||
sliced per row. The result is a list whose length is ``batch.num_rows`` and
|
||||
whose items are dicts keyed by column name. This shape composes directly
|
||||
with PyTorch's default ``DataLoader`` collate, which stacks the per-row
|
||||
dicts back into a dict of batched tensors.
|
||||
|
||||
Fails if torch is not installed.
|
||||
Fails if a column's data type is not supported by PyTorch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
batch : pa.RecordBatch
|
||||
The record batch to convert.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict[str, torch.Tensor]]
|
||||
One per-row dict per row in the batch. Each dict maps column name to a
|
||||
0-D tensor view into the column.
|
||||
"""
|
||||
torch = attempt_import_or_raise("torch", "torch")
|
||||
tensors = {
|
||||
name: torch.from_dlpack(col)
|
||||
for name, col in zip(batch.schema.names, batch.columns)
|
||||
}
|
||||
return [
|
||||
{name: tensor[i] for name, tensor in tensors.items()}
|
||||
for i in range(batch.num_rows)
|
||||
]
|
||||
|
||||
@@ -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))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import List, Optional, Union
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -242,6 +243,49 @@ def test_embedding_with_bad_results(tmp_path):
|
||||
assert tbl["vector"].null_count == 1
|
||||
|
||||
|
||||
def test_embedding_with_empty_output_vectors(tmp_path):
|
||||
"""Regression test for issue #1672.
|
||||
|
||||
When an embedding function returns an empty list (e.g. for empty-string
|
||||
inputs), _append_vector_columns used to crash because PyArrow cannot cast
|
||||
[] into a fixed-size list element. The fix replaces wrong-length vectors
|
||||
with None before building the Arrow array so that _handle_bad_vectors can
|
||||
process them normally.
|
||||
"""
|
||||
|
||||
@register("empty-vec-embedding")
|
||||
class EmptyVecEmbeddingFunction(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 128
|
||||
|
||||
def generate_embeddings(self, texts: Union[List[str], np.ndarray]) -> list:
|
||||
# Simulate a model that returns an empty list for blank inputs
|
||||
return [
|
||||
[] if text.strip() == "" else np.random.randn(self.ndims()).tolist()
|
||||
for text in texts
|
||||
]
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
model = registry.get("empty-vec-embedding").create()
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = model.SourceField()
|
||||
vector: Vector(model.ndims()) = model.VectorField()
|
||||
|
||||
table = db.create_table("test_empty_vec", schema=Schema, mode="overwrite")
|
||||
|
||||
# Should not crash; the row with the empty string should be dropped
|
||||
table.add(
|
||||
[{"text": "hello world"}, {"text": ""}, {"text": "foo"}],
|
||||
on_bad_vectors="drop",
|
||||
)
|
||||
|
||||
assert len(table) == 2
|
||||
texts = table.to_arrow()["text"].to_pylist()
|
||||
assert "" not in texts
|
||||
|
||||
|
||||
def test_with_existing_vectors(tmp_path):
|
||||
@register("mock-embedding")
|
||||
class MockEmbeddingFunction(TextEmbeddingFunction):
|
||||
@@ -548,6 +592,22 @@ def test_openai_no_retry_on_401(mock_sleep):
|
||||
assert mock_sleep.call_count == 0
|
||||
|
||||
|
||||
def test_ollama_embeddings_pickle():
|
||||
"""OllamaEmbeddings must pickle even after the cached client is created."""
|
||||
registry = get_registry()
|
||||
model = registry.get("ollama").create(name="nomic-embed-text")
|
||||
|
||||
# Simulate accessing the cached client, which stores it on the instance.
|
||||
model.__dict__["_ollama_client"] = MagicMock()
|
||||
|
||||
pickled = pickle.dumps(model)
|
||||
restored = pickle.loads(pickled)
|
||||
|
||||
assert restored.name == "nomic-embed-text"
|
||||
assert restored.host == "http://localhost:11434"
|
||||
assert "_ollama_client" not in restored.__dict__
|
||||
|
||||
|
||||
def test_url_retrieve_downloads_image():
|
||||
"""
|
||||
Embedding functions like open-clip, siglip, and jinaai use url_retrieve()
|
||||
|
||||
@@ -3,10 +3,14 @@
|
||||
|
||||
"""Tests for the type-safe expression builder API."""
|
||||
|
||||
import pytest
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from decimal import Decimal
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import lancedb
|
||||
from lancedb.expr import Expr, col, lit, func
|
||||
from lancedb.expr import Expr, col, func, lit
|
||||
|
||||
|
||||
# ── unit tests for Expr construction ─────────────────────────────────────────
|
||||
@@ -54,6 +58,28 @@ class TestExprConstruction:
|
||||
with pytest.raises(Exception):
|
||||
func("not_a_real_function", col("x"))
|
||||
|
||||
def test_lit_date(self):
|
||||
e = lit(date(2024, 1, 1))
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_datetime(self):
|
||||
# Naive datetime
|
||||
e = lit(datetime(2024, 1, 1, 10, 0))
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_datetime_tz(self):
|
||||
# Timezone-aware datetime
|
||||
tz = timezone(timedelta(hours=5))
|
||||
dt = datetime(2024, 1, 1, 10, 0, tzinfo=tz)
|
||||
e = lit(dt)
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_decimal_precision(self):
|
||||
# High precision Decimal that would be rounded if converted to float
|
||||
d = Decimal("1.234567890123456789")
|
||||
e = lit(d)
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
|
||||
class TestExprOperators:
|
||||
def test_eq_operator(self):
|
||||
@@ -142,6 +168,20 @@ class TestExprOperators:
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(name = 'alice')"
|
||||
|
||||
def test_reflexive_comparisons(self):
|
||||
# 10 < col("age") swaps to col("age") > 10
|
||||
assert (10 < col("age")).to_sql() == "(age > 10)"
|
||||
assert (10 <= col("age")).to_sql() == "(age >= 10)"
|
||||
assert (10 > col("age")).to_sql() == "(age < 10)"
|
||||
assert (10 >= col("age")).to_sql() == "(age <= 10)"
|
||||
assert (10 == col("age")).to_sql() == "(age = 10)"
|
||||
assert (10 != col("age")).to_sql() == "(age <> 10)"
|
||||
|
||||
def test_reflexive_logical(self):
|
||||
# True & Expr calls Expr.__rand__(True)
|
||||
assert (True & (col("age") > 18)).to_sql() == "(true AND (age > 18))"
|
||||
assert (False | (col("age") > 18)).to_sql() == "(false OR (age > 18))"
|
||||
|
||||
|
||||
class TestExprBytesLiteral:
|
||||
def test_bytes_to_sql(self):
|
||||
@@ -282,6 +322,40 @@ class TestExprRepr:
|
||||
{e: 1}
|
||||
|
||||
|
||||
class TestExprReflexive:
|
||||
def test_reflexive_eq(self):
|
||||
e = 1 == col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x = 1)"
|
||||
|
||||
def test_reflexive_ne(self):
|
||||
e = 1 != col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x <> 1)"
|
||||
|
||||
def test_reflexive_lt(self):
|
||||
# 1 < x => (x > 1)
|
||||
e = 1 < col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x > 1)"
|
||||
|
||||
def test_reflexive_gt(self):
|
||||
# 1 > x => (x < 1)
|
||||
e = 1 > col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x < 1)"
|
||||
|
||||
def test_reflexive_and(self):
|
||||
e = True & col("active")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(true AND active)"
|
||||
|
||||
def test_reflexive_or(self):
|
||||
e = False | col("inactive")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(false OR inactive)"
|
||||
|
||||
|
||||
# ── integration tests: end-to-end query against a real table ─────────────────
|
||||
|
||||
|
||||
@@ -432,6 +506,72 @@ class TestColNamingIntegration:
|
||||
assert sorted(result["upper_name"].to_pylist()) == ["ALICE", "BOB", "CHARLIE"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def type_check_table(tmp_path):
|
||||
"""Fixture that creates a table with Date32 and Decimal128 columns."""
|
||||
db = lancedb.connect(str(tmp_path))
|
||||
schema = pa.schema(
|
||||
[
|
||||
("date", pa.date32()),
|
||||
("decimal", pa.decimal128(10, 2)),
|
||||
("binary", pa.binary()),
|
||||
]
|
||||
)
|
||||
data = pa.table(
|
||||
{
|
||||
"date": [date(2024, 1, 1), date(2024, 1, 2)],
|
||||
"decimal": [Decimal("10.50"), Decimal("20.75")],
|
||||
"binary": [b"\x01", b"\x02"],
|
||||
},
|
||||
schema=schema,
|
||||
)
|
||||
return db.create_table("extended_types", data)
|
||||
|
||||
|
||||
class TestExtendedTypeIntegration:
|
||||
"""Integration tests verifying that typed literals work correctly in filters."""
|
||||
|
||||
def test_date_integration(self, type_check_table):
|
||||
"""Verify that Date32 literals are correctly parsed and filtered."""
|
||||
result = (
|
||||
type_check_table.search()
|
||||
.where(col("date") == lit(date(2024, 1, 1)))
|
||||
.to_arrow()
|
||||
)
|
||||
assert result.num_rows == 1
|
||||
assert result["date"][0].as_py() == date(2024, 1, 1)
|
||||
|
||||
def test_decimal_integration(self, tmp_path):
|
||||
"""A Decimal literal must retain full 128-bit precision in a filter.
|
||||
|
||||
1.234567890123456789 and 1.234567890123456790 differ only in the last
|
||||
digit and are indistinguishable once truncated to f64. The filter
|
||||
therefore returns the single expected row only if ``lit(Decimal)``
|
||||
produces a true ``Decimal128`` scalar rather than being coerced to f64.
|
||||
"""
|
||||
low = Decimal("1.234567890123456789")
|
||||
high = Decimal("1.234567890123456790")
|
||||
|
||||
db = lancedb.connect(str(tmp_path / "decimal_precision"))
|
||||
schema = pa.schema([("val", pa.decimal128(19, 18))])
|
||||
table = db.create_table(
|
||||
"decimal_precision",
|
||||
pa.table({"val": [low, high]}, schema=schema),
|
||||
)
|
||||
|
||||
result = table.search().where(col("val") < lit(high)).to_arrow()
|
||||
assert result.num_rows == 1
|
||||
assert result["val"][0].as_py() == low
|
||||
|
||||
def test_binary_integration(self, type_check_table):
|
||||
"""Verify that Binary literals are correctly filtered."""
|
||||
result = (
|
||||
type_check_table.search().where(col("binary") == lit(b"\x01")).to_arrow()
|
||||
)
|
||||
assert result.num_rows == 1
|
||||
assert result["binary"][0].as_py() == b"\x01"
|
||||
|
||||
|
||||
# ── bytes / binary column integration tests ───────────────────────────────────
|
||||
|
||||
|
||||
@@ -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,23 @@ 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 = []
|
||||
table._analyze_plan.side_effect = lambda query: analyzed_queries.append(query) or ""
|
||||
|
||||
(
|
||||
LanceHybridQueryBuilder(table)
|
||||
.vector([0.1, 0.2])
|
||||
.text("puppy runs")
|
||||
.phrase_query()
|
||||
.analyze_plan()
|
||||
)
|
||||
|
||||
assert len(analyzed_queries) == 2
|
||||
assert analyzed_queries[1].full_text_query.query == '"puppy runs"'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def table_with_id(tmpdir_factory) -> Table:
|
||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user