mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-27 08:28:28 +00:00
Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 1545bd4b4f | |||
| 589331d86b | |||
| 56f856569c |
@@ -1,14 +1,6 @@
|
||||
---
|
||||
name: lancedb-branch-ops
|
||||
description: >-
|
||||
Manage LanceDB table branches through the REST API: list, create, and delete
|
||||
branches; target schema reads, field-metadata updates, and index creation to a
|
||||
named branch; and verify that branch changes remain isolated from main. Use
|
||||
when a task involves branch lifecycle, an experimental or isolated table
|
||||
version, directing an operation to a non-main branch, or confirming that a
|
||||
mutation did not affect main. This skill also explains that LanceDB has no
|
||||
checkout operation; each request selects its target branch in the request
|
||||
body.
|
||||
description: Branch management for LanceDB tables via the REST API. Use this skill whenever someone wants to create, delete, list, or switch branches on a LanceDB table — or needs to make sure a write (metadata update, index build, etc.) lands on a specific branch instead of main. Invoke it even without the word "branch" if context makes clear they want an experimental copy of a table, want to isolate changes, or want to confirm a mutation didn't touch main. Covers: branches/list, branches/create, branches/delete, and passing "branch" in describe/update_field_metadata/create_index to target a non-main version.
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
@@ -1,79 +0,0 @@
|
||||
---
|
||||
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.
|
||||
@@ -1,105 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,173 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,131 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,78 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,100 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,78 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,135 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Scan Python and TypeScript for likely unsafe LanceDB materialization."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PY_FULL_TABLE = re.compile(r"\b\w+\.(to_pandas|to_arrow|to_polars)\s*\(")
|
||||
TS_TABLE_TO_ARROW = re.compile(r"\b\w+\.toArrow\s*\(")
|
||||
TS_QUERY_COLLECTOR = re.compile(r"\.query\s*\(\s*\)[\s\S]*?\.to(Array|Arrow)\s*\(")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Finding:
|
||||
path: Path
|
||||
line: int
|
||||
message: str
|
||||
text: str
|
||||
|
||||
|
||||
def iter_files(paths: list[Path]) -> list[Path]:
|
||||
files: list[Path] = []
|
||||
for path in paths:
|
||||
if path.is_dir():
|
||||
files.extend(
|
||||
p
|
||||
for p in path.rglob("*")
|
||||
if p.suffix in {".py", ".ts", ".tsx"} and "node_modules" not in p.parts
|
||||
)
|
||||
elif path.suffix in {".py", ".ts", ".tsx"}:
|
||||
files.append(path)
|
||||
return sorted(set(files))
|
||||
|
||||
|
||||
def line_number(text: str, offset: int) -> int:
|
||||
return text.count("\n", 0, offset) + 1
|
||||
|
||||
|
||||
def scan_python(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in PY_FULL_TABLE.finditer(text):
|
||||
line_start = text.rfind("\n", 0, match.start()) + 1
|
||||
line_end = text.find("\n", match.start())
|
||||
if line_end == -1:
|
||||
line_end = len(text)
|
||||
line = text[line_start:line_end].strip()
|
||||
if ".search(" in line or ".query(" in line:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
f"Review Python `{match.group(1)}()` call; table-level materialization is not portable to remote tables.",
|
||||
line,
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def statement_around(text: str, start: int, end: int) -> str:
|
||||
before = max(text.rfind(";", 0, start), text.rfind("\n\n", 0, start))
|
||||
after_candidates = [pos for pos in (text.find(";", end), text.find("\n\n", end)) if pos != -1]
|
||||
after = min(after_candidates) if after_candidates else len(text)
|
||||
return text[before + 1 : after].strip()
|
||||
|
||||
|
||||
def scan_typescript(path: Path, text: str) -> list[Finding]:
|
||||
findings: list[Finding] = []
|
||||
for match in TS_TABLE_TO_ARROW.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".query(" in stmt or ".search(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review TypeScript `table.toArrow()`-style call; table-level materialization is not portable for large/remote tables.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
|
||||
for match in TS_QUERY_COLLECTOR.finditer(text):
|
||||
stmt = statement_around(text, match.start(), match.end())
|
||||
if ".limit(" in stmt:
|
||||
continue
|
||||
findings.append(
|
||||
Finding(
|
||||
path,
|
||||
line_number(text, match.start()),
|
||||
"Review unbounded TypeScript query collection; add `limit()` or stream batches.",
|
||||
stmt.splitlines()[0].strip(),
|
||||
)
|
||||
)
|
||||
return findings
|
||||
|
||||
|
||||
def scan_file(path: Path) -> list[Finding]:
|
||||
text = path.read_text(encoding="utf-8", errors="replace")
|
||||
if path.suffix == ".py":
|
||||
return scan_python(path, text)
|
||||
if path.suffix in {".ts", ".tsx"}:
|
||||
return scan_typescript(path, text)
|
||||
return []
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("paths", nargs="+", type=Path)
|
||||
parser.add_argument(
|
||||
"--no-fail", action="store_true", help="Always exit 0 after reporting findings."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
findings: list[Finding] = []
|
||||
for path in iter_files(args.paths):
|
||||
findings.extend(scan_file(path))
|
||||
|
||||
for finding in findings:
|
||||
print(f"{finding.path}:{finding.line}: {finding.message}")
|
||||
print(f" {finding.text}")
|
||||
|
||||
if findings:
|
||||
print(
|
||||
f"\n{len(findings)} finding(s). Review manually; bounded query result conversion may be OK."
|
||||
)
|
||||
return 0 if args.no_fail or not findings else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.32.0-beta.2"
|
||||
current_version = "0.31.0-beta.6"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -34,16 +34,15 @@ 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/' -e PROTOC=/usr/local/bin/protoc"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
target: x86_64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
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
|
||||
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
|
||||
- name: Build Arm Manylinux Wheel
|
||||
if: ${{ inputs.arm-build == 'true' }}
|
||||
uses: PyO3/maturin-action@v1
|
||||
@@ -51,14 +50,13 @@ 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/' -e PROTOC=/usr/local/bin/protoc"
|
||||
docker-options: "-e PIP_EXTRA_INDEX_URL='https://pypi.fury.io/lance-format/ https://pypi.fury.io/lancedb/'"
|
||||
target: aarch64-unknown-linux-gnu
|
||||
manylinux: ${{ inputs.manylinux }}
|
||||
args: ${{ inputs.args }}
|
||||
before-script-linux: |
|
||||
set -e
|
||||
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
|
||||
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
|
||||
|
||||
@@ -103,7 +103,7 @@ jobs:
|
||||
features: fp16kernels
|
||||
pre_build: brew install protobuf
|
||||
- target: x86_64-pc-windows-msvc
|
||||
host: windows-2025-8x-x64
|
||||
host: windows-latest
|
||||
features: ","
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
@@ -111,21 +111,12 @@ 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-2025-8x-x64
|
||||
host: windows-latest
|
||||
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
|
||||
|
||||
@@ -125,26 +125,10 @@ jobs:
|
||||
- uses: rui314/setup-mold@v1
|
||||
- name: Make Swap
|
||||
run: |
|
||||
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"
|
||||
sudo fallocate -l 16G /swapfile
|
||||
sudo chmod 600 /swapfile
|
||||
sudo mkswap /swapfile
|
||||
sudo swapon /swapfile
|
||||
- name: Build
|
||||
run: cargo build --profile ci --all-features --tests --locked --examples
|
||||
- name: Run feature tests
|
||||
|
||||
Generated
+256
-483
File diff suppressed because it is too large
Load Diff
+23
-25
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.0.0-beta.23", default-features = false, "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.23", default-features = false, "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.23", default-features = false, "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.23", "rev" = "2ebc588f809b2937e9dbf1bae48ffdfd844cef61", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=8.0.0", default-features = false }
|
||||
lance-core = "=8.0.0"
|
||||
lance-datagen = "=8.0.0"
|
||||
lance-file = "=8.0.0"
|
||||
lance-io = { "version" = "=8.0.0", default-features = false }
|
||||
lance-index = "=8.0.0"
|
||||
lance-linalg = "=8.0.0"
|
||||
lance-namespace = "=8.0.0"
|
||||
lance-namespace-impls = { "version" = "=8.0.0", default-features = false }
|
||||
lance-table = "=8.0.0"
|
||||
lance-testing = "=8.0.0"
|
||||
lance-datafusion = "=8.0.0"
|
||||
lance-encoding = "=8.0.0"
|
||||
lance-arrow = "=8.0.0"
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
@@ -39,23 +39,21 @@ arrow-schema = "58.0.0"
|
||||
arrow-select = "58.0.0"
|
||||
arrow-cast = "58.0.0"
|
||||
async-trait = "0"
|
||||
datafusion = { version = "54.0.0", default-features = false }
|
||||
datafusion-catalog = "54.0.0"
|
||||
datafusion-common = { version = "54.0.0", default-features = false }
|
||||
datafusion-execution = "54.0.0"
|
||||
datafusion-expr = "54.0.0"
|
||||
datafusion-functions = "54.0.0"
|
||||
datafusion-physical-plan = "54.0.0"
|
||||
datafusion-physical-expr = "54.0.0"
|
||||
datafusion-sql = "54.0.0"
|
||||
datafusion = { version = "53.0.0", default-features = false }
|
||||
datafusion-catalog = "53.0.0"
|
||||
datafusion-common = { version = "53.0.0", default-features = false }
|
||||
datafusion-execution = "53.0.0"
|
||||
datafusion-expr = "53.0.0"
|
||||
datafusion-functions = "53.0.0"
|
||||
datafusion-physical-plan = "53.0.0"
|
||||
datafusion-physical-expr = "53.0.0"
|
||||
datafusion-sql = "53.0.0"
|
||||
env_logger = "0.11"
|
||||
half = { "version" = "2.7.1", default-features = false, features = [
|
||||
"num-traits",
|
||||
] }
|
||||
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,6 +51,18 @@ 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.
|
||||
@@ -68,6 +80,18 @@ 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.
|
||||
@@ -84,23 +108,17 @@ 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.32.0-beta.2</version>
|
||||
<version>0.31.0-beta.6</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -398,26 +398,6 @@ 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
|
||||
@@ -934,32 +914,6 @@ 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
|
||||
|
||||
@@ -1,42 +0,0 @@
|
||||
[**@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.
|
||||
@@ -1,26 +0,0 @@
|
||||
[**@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,7 +72,6 @@
|
||||
- [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)
|
||||
@@ -108,7 +107,6 @@
|
||||
- [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)
|
||||
@@ -118,7 +116,6 @@
|
||||
|
||||
## Type Aliases
|
||||
|
||||
- [BaseTokenizer](type-aliases/BaseTokenizer.md)
|
||||
- [Data](type-aliases/Data.md)
|
||||
- [DataLike](type-aliases/DataLike.md)
|
||||
- [FieldLike](type-aliases/FieldLike.md)
|
||||
@@ -128,15 +125,12 @@
|
||||
- [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: BaseTokenizer;
|
||||
optional baseTokenizer: "raw" | "simple" | "whitespace" | "ngram";
|
||||
```
|
||||
|
||||
The tokenizer to use when building the index.
|
||||
@@ -37,10 +37,6 @@ 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?
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
[**@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,14 +8,6 @@
|
||||
|
||||
## Properties
|
||||
|
||||
### clumpSize?
|
||||
|
||||
```ts
|
||||
optional clumpSize: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### counts?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -1,109 +0,0 @@
|
||||
[**@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.
|
||||
@@ -1,19 +0,0 @@
|
||||
[**@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}`;
|
||||
```
|
||||
@@ -1,11 +0,0 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TokenizeTableOptions
|
||||
|
||||
# Type Alias: TokenizeTableOptions
|
||||
|
||||
```ts
|
||||
type TokenizeTableOptions: object | object;
|
||||
```
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.31.0-beta.6</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.32.0-beta.2</version>
|
||||
<version>0.31.0-beta.6</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.23</lance-core.version>
|
||||
<lance-core.version>8.0.0</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.32.0-beta.2"
|
||||
version = "0.31.0-beta.6"
|
||||
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", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
// 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();
|
||||
});
|
||||
});
|
||||
@@ -16,7 +16,6 @@ import {
|
||||
PhraseQuery,
|
||||
Table,
|
||||
connect,
|
||||
tokenize,
|
||||
} from "../lancedb";
|
||||
import {
|
||||
Table as ArrowTable,
|
||||
@@ -2308,75 +2307,6 @@ 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 }];
|
||||
@@ -3062,56 +2992,6 @@ 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,21 +13,13 @@ 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,
|
||||
@@ -117,7 +109,6 @@ export {
|
||||
HnswPqOptions,
|
||||
HnswSqOptions,
|
||||
FtsOptions,
|
||||
BaseTokenizer,
|
||||
} from "./indices";
|
||||
|
||||
export {
|
||||
@@ -128,8 +119,6 @@ export {
|
||||
OptimizeOptions,
|
||||
Version,
|
||||
WriteProgress,
|
||||
FtsToken,
|
||||
TokenizeTableOptions,
|
||||
LsmWriteSpec,
|
||||
ColumnAlteration,
|
||||
FieldMetadataUpdate,
|
||||
@@ -161,68 +150,6 @@ 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,16 +486,6 @@ 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
|
||||
*/
|
||||
@@ -519,12 +509,8 @@ 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?: BaseTokenizer;
|
||||
baseTokenizer?: "simple" | "whitespace" | "raw" | "ngram";
|
||||
|
||||
/**
|
||||
* language for stemming and stop words
|
||||
|
||||
@@ -1,137 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -158,26 +158,6 @@ 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`.
|
||||
@@ -605,17 +585,6 @@ 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.
|
||||
*
|
||||
@@ -736,19 +705,6 @@ 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>;
|
||||
|
||||
@@ -1135,15 +1091,6 @@ 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();
|
||||
}
|
||||
@@ -1206,17 +1153,6 @@ export class LocalTable extends Table {
|
||||
return await this.inner.listIndices();
|
||||
}
|
||||
|
||||
async tokenize(
|
||||
query: string,
|
||||
options: TokenizeTableOptions,
|
||||
): Promise<FtsToken[]> {
|
||||
return await this.inner.tokenize(
|
||||
query,
|
||||
options?.column,
|
||||
options?.indexName,
|
||||
);
|
||||
}
|
||||
|
||||
async toArrow(): Promise<ArrowTable> {
|
||||
return await this.query().toArrow();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+2
-73
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
@@ -18,7 +18,6 @@
|
||||
"win32"
|
||||
],
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"devDependencies": {
|
||||
@@ -28,7 +27,6 @@
|
||||
"@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",
|
||||
@@ -4150,75 +4148,6 @@
|
||||
"@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",
|
||||
|
||||
+1
-3
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.32.0-beta.2",
|
||||
"version": "0.31.0-beta.6",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
@@ -44,7 +44,6 @@
|
||||
"@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",
|
||||
@@ -93,7 +92,6 @@
|
||||
"version": "napi version"
|
||||
},
|
||||
"dependencies": {
|
||||
"@opentelemetry/api": "^1.9.0",
|
||||
"reflect-metadata": "^0.2.2"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
|
||||
Generated
-53
@@ -8,9 +8,6 @@ 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
|
||||
@@ -36,9 +33,6 @@ 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
|
||||
@@ -1313,32 +1307,6 @@ 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==}
|
||||
|
||||
@@ -4957,27 +4925,6 @@ 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,11 +9,8 @@ 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]
|
||||
@@ -33,65 +30,6 @@ 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,7 +12,6 @@ mod header;
|
||||
mod index;
|
||||
mod iterator;
|
||||
pub mod merge;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
mod query;
|
||||
pub mod remote;
|
||||
|
||||
@@ -1,119 +0,0 @@
|
||||
// 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,7 +16,6 @@ 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>>,
|
||||
}
|
||||
|
||||
@@ -126,15 +125,10 @@ 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,
|
||||
clump_size,
|
||||
},
|
||||
SplitStrategy::Random { seed, sizes },
|
||||
options.split_names.clone(),
|
||||
)
|
||||
})
|
||||
|
||||
+2
-92
@@ -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, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, NewColumnTransform, OptimizeAction,
|
||||
OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
|
||||
@@ -411,16 +411,6 @@ 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()
|
||||
@@ -574,27 +564,6 @@ 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()?;
|
||||
@@ -702,24 +671,6 @@ 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`.
|
||||
///
|
||||
@@ -777,47 +728,6 @@ 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.35.0-beta.2"
|
||||
current_version = "0.34.0"
|
||||
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.35.0-beta.2"
|
||||
version = "0.34.0"
|
||||
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", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
#!/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,10 +47,6 @@ 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",
|
||||
@@ -65,7 +61,6 @@ tests = [
|
||||
"pylance>=5.0.0b5",
|
||||
"requests>=2.31.0",
|
||||
"datafusion>=52,<53",
|
||||
"opentelemetry-sdk>=1.30.0",
|
||||
]
|
||||
dev = [
|
||||
"ruff>=0.3.0",
|
||||
|
||||
@@ -6,22 +6,19 @@ import importlib.metadata
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional, Union, Any, List, Iterable
|
||||
from typing import Dict, Optional, Union, Any, List
|
||||
|
||||
__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 blob, vector, BlobType
|
||||
from .schema import vector
|
||||
from .table import AsyncTable, Table
|
||||
from .types import BaseTokenizerType
|
||||
from ._lancedb import Session
|
||||
from .namespace import (
|
||||
connect_namespace,
|
||||
@@ -152,14 +149,8 @@ def connect(
|
||||
|
||||
For object storage, use a URI prefix:
|
||||
|
||||
>>> db = lancedb.connect( # doctest: +SKIP
|
||||
... "s3://my-bucket/lancedb",
|
||||
... storage_options={
|
||||
... "aws_access_key_id": "***",
|
||||
... "aws_secret_access_key": "***",
|
||||
... "aws_region": "us-east-1",
|
||||
... },
|
||||
... )
|
||||
>>> db = lancedb.connect("s3://my-bucket/lancedb",
|
||||
... storage_options={"aws_access_key_id": "***"})
|
||||
|
||||
For tests and temporary data, use an in-memory database:
|
||||
|
||||
@@ -249,40 +240,6 @@ 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_"
|
||||
|
||||
|
||||
@@ -493,21 +450,17 @@ 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",
|
||||
|
||||
@@ -1,420 +0,0 @@
|
||||
# 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,5 +1,4 @@
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -25,45 +24,11 @@ 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)."""
|
||||
|
||||
@@ -88,9 +53,7 @@ class PyExpr:
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
def expr_lit(
|
||||
value: Union[bool, int, float, str, bytes, date, datetime, Decimal],
|
||||
) -> PyExpr: ...
|
||||
def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ...
|
||||
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
||||
|
||||
class Session:
|
||||
@@ -196,17 +159,6 @@ 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: ...
|
||||
@@ -253,13 +205,6 @@ 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: ...
|
||||
@@ -281,7 +226,6 @@ 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: ...
|
||||
@@ -291,13 +235,6 @@ 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:
|
||||
@@ -533,10 +470,6 @@ 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, Optional, Union
|
||||
from typing import List, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -15,8 +15,6 @@ 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):
|
||||
@@ -83,7 +81,6 @@ class GeminiText(TextEmbeddingFunction):
|
||||
"""
|
||||
|
||||
name: str = "gemini-embedding-001"
|
||||
dim: Optional[int] = None
|
||||
query_task_type: str = "retrieval_query"
|
||||
source_task_type: str = "retrieval_document"
|
||||
|
||||
@@ -96,8 +93,6 @@ class GeminiText(TextEmbeddingFunction):
|
||||
model_config["ignored_types"] = (cached_property,)
|
||||
|
||||
def ndims(self):
|
||||
if self.dim:
|
||||
return self.dim
|
||||
# TODO: fix hardcoding
|
||||
return 768
|
||||
|
||||
@@ -138,22 +133,22 @@ class GeminiText(TextEmbeddingFunction):
|
||||
contents.append({"parts": [{"text": text}]})
|
||||
|
||||
# Build config
|
||||
config_kwargs = {"output_dimensionality": self.ndims()}
|
||||
config_kwargs = {}
|
||||
if task_type:
|
||||
config_kwargs["task_type"] = task_type.upper() # API expects uppercase
|
||||
|
||||
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
|
||||
# Call embed_content for each content
|
||||
embeddings = []
|
||||
for i in range(0, len(contents), EMBEDDING_BATCH_SIZE):
|
||||
chunk = contents[i : i + EMBEDDING_BATCH_SIZE]
|
||||
for content in contents:
|
||||
config = (
|
||||
types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||
)
|
||||
response = self.client.models.embed_content(
|
||||
model=self.name,
|
||||
contents=chunk,
|
||||
contents=content,
|
||||
config=config,
|
||||
)
|
||||
embeddings.extend([np.array(e.values) for e in response.embeddings])
|
||||
embeddings.append(response.embeddings[0].values)
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -165,13 +160,5 @@ 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"),
|
||||
http_options={
|
||||
"headers": {
|
||||
"x-goog-api-client": f"lancedb/{__version__}",
|
||||
}
|
||||
},
|
||||
)
|
||||
return genai_module.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
|
||||
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -56,16 +56,6 @@ 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,8 +19,6 @@ operators::
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Iterable, Union
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -65,7 +63,7 @@ def _coerce(value: "ExprLike") -> "Expr":
|
||||
|
||||
|
||||
# Type alias used in annotations.
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes, date, datetime, Decimal]
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes]
|
||||
|
||||
|
||||
class Expr:
|
||||
@@ -120,18 +118,10 @@ 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_())
|
||||
@@ -276,14 +266,13 @@ def col(name: str) -> Expr:
|
||||
return Expr(expr_col(name))
|
||||
|
||||
|
||||
def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr:
|
||||
def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
"""Create a literal (constant) value expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value:
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``,
|
||||
``datetime``, or ``Decimal``.
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -291,9 +280,6 @@ def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) ->
|
||||
>>> 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,8 +127,6 @@ 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"
|
||||
|
||||
@@ -1,170 +0,0 @@
|
||||
# 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_dict, batch_to_tensor_rows
|
||||
from .util import batch_to_tensor, batch_to_tensor_rows
|
||||
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -65,7 +65,6 @@ 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":
|
||||
"""
|
||||
@@ -88,9 +87,6 @@ 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.
|
||||
"""
|
||||
@@ -99,7 +95,6 @@ class PermutationBuilder:
|
||||
counts=counts,
|
||||
fixed=fixed,
|
||||
seed=seed,
|
||||
clump_size=clump_size,
|
||||
split_names=split_names,
|
||||
)
|
||||
return self
|
||||
@@ -946,7 +941,6 @@ class Permutation:
|
||||
"pandas",
|
||||
"arrow",
|
||||
"torch",
|
||||
"torch_row",
|
||||
"torch_col",
|
||||
"polars",
|
||||
],
|
||||
@@ -962,19 +956,15 @@ 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 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" - 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 and torch_col will be zero-copy but will only support a
|
||||
subset of data types (numeric types).
|
||||
Conversion to 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
|
||||
@@ -995,8 +985,6 @@ 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)
|
||||
|
||||
+78
-299
@@ -15,12 +15,10 @@ from typing import (
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
Tuple,
|
||||
Type,
|
||||
TypeVar,
|
||||
Union,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
import deprecation
|
||||
@@ -41,21 +39,15 @@ 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
|
||||
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
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
_BLOB_MODE_TO_HANDLING = {
|
||||
"lazy": "blobs_descriptions",
|
||||
"bytes": "all_binary",
|
||||
"descriptions": "blobs_descriptions",
|
||||
}
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import sys
|
||||
@@ -81,22 +73,25 @@ if TYPE_CHECKING:
|
||||
T = TypeVar("T", bound="LanceModel")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _LanceScanner(Protocol):
|
||||
projected_schema: pa.Schema | None
|
||||
schema: pa.Schema | None
|
||||
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 to_pandas(self, blob_mode: BlobMode | None = ..., **kwargs) -> pd.DataFrame: ...
|
||||
|
||||
def to_pyarrow(self): ...
|
||||
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_table(self) -> pa.Table: ...
|
||||
|
||||
def to_reader(self): ...
|
||||
def _schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return any(_field_is_blob(field) for field in schema)
|
||||
|
||||
|
||||
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:
|
||||
@@ -145,7 +140,13 @@ def _combine_where(
|
||||
return f"({existing_sql}) AND ({new_sql})"
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
def _projection_to_scanner_kwargs(
|
||||
columns: Optional[
|
||||
Union[
|
||||
List[str], List[Tuple[str, Union[str, Expr]]], Dict[str, Union[str, Expr]]
|
||||
]
|
||||
],
|
||||
) -> Dict[str, Any]:
|
||||
if columns is None:
|
||||
return {}
|
||||
if isinstance(columns, list):
|
||||
@@ -170,11 +171,7 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
|
||||
|
||||
def _scanner_kwargs_for_query(
|
||||
query: Query,
|
||||
blob_mode: BlobMode,
|
||||
dataset: Optional[Any] = None,
|
||||
*,
|
||||
with_row_id: Optional[bool] = None,
|
||||
query: Query, blob_mode: BlobMode, dataset: Optional[Any] = None
|
||||
) -> Dict[str, Any]:
|
||||
fragments = _scanner_fragments_for_query(query, dataset)
|
||||
kwargs = {
|
||||
@@ -182,10 +179,10 @@ def _scanner_kwargs_for_query(
|
||||
"filter": _filter_to_sql(query.filter),
|
||||
"limit": query.limit,
|
||||
"offset": query.offset,
|
||||
"with_row_id": with_row_id if with_row_id is not None else query.with_row_id,
|
||||
"with_row_id": 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}
|
||||
@@ -218,11 +215,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 is_blob_like_field(field) or field.name not in df.columns:
|
||||
if not _field_is_blob(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"):
|
||||
@@ -232,7 +229,7 @@ def _ensure_lazy_blob_frame(
|
||||
return df
|
||||
|
||||
|
||||
def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
||||
def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||
if hasattr(scanner, "to_pyarrow"):
|
||||
reader = scanner.to_pyarrow()
|
||||
return reader.read_all()
|
||||
@@ -242,9 +239,7 @@ def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
||||
return reader.read_all()
|
||||
|
||||
|
||||
def _scanner_to_pandas(
|
||||
scanner: _LanceScanner, blob_mode: BlobMode, **kwargs
|
||||
) -> pd.DataFrame:
|
||||
def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataFrame":
|
||||
schema = getattr(scanner, "projected_schema", None)
|
||||
if schema is None:
|
||||
schema = getattr(scanner, "schema", None)
|
||||
@@ -265,71 +260,13 @@ def _scanner_to_pandas(
|
||||
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,
|
||||
@@ -737,7 +674,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: QueryProjection = None
|
||||
columns: Optional[Union[List[str], Dict[str, Union[str, Expr]]]] = None
|
||||
|
||||
# minimum number of IVF partitions to search
|
||||
#
|
||||
@@ -1021,7 +958,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()
|
||||
@@ -1080,11 +1017,6 @@ 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
|
||||
@@ -1263,42 +1195,6 @@ 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.
|
||||
|
||||
@@ -1475,29 +1371,13 @@ class LanceQueryBuilder(ABC):
|
||||
return None
|
||||
|
||||
dataset = self._table.to_lance()
|
||||
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,
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
)
|
||||
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:
|
||||
@@ -1745,9 +1625,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
The maximum time to wait for the query to complete.
|
||||
If None, wait indefinitely.
|
||||
"""
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
"""
|
||||
@@ -1807,7 +1685,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._query_for_scan()
|
||||
query = self.to_query_object()
|
||||
result_set = self._table._execute_query(
|
||||
query, batch_size=batch_size, timeout=timeout
|
||||
)
|
||||
@@ -1951,7 +1829,8 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1961,21 +1840,6 @@ 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
|
||||
@@ -2000,7 +1864,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
fragments=self._fragments,
|
||||
fragment_ids=self._fragment_ids,
|
||||
full_text_query=FullTextSearchQuery(
|
||||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
||||
query=self._query, columns=self._fts_columns
|
||||
),
|
||||
offset=self._offset,
|
||||
fast_search=self._fast_search,
|
||||
@@ -2018,13 +1882,22 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
self._table._ensure_no_legacy_fts_index()
|
||||
|
||||
query = self._query_for_scan()
|
||||
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()
|
||||
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 self._finalize_blob_query_table(results)
|
||||
return results
|
||||
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
@@ -2052,9 +1925,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||
return self._finalize_blob_query_table(
|
||||
self.to_batches(timeout=timeout).read_all()
|
||||
)
|
||||
return self.to_batches(timeout=timeout).read_all()
|
||||
|
||||
def to_query_object(self) -> Query:
|
||||
return Query(
|
||||
@@ -2076,7 +1947,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||
def to_batches(
|
||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||
) -> pa.RecordBatchReader:
|
||||
query = self._query_for_scan()
|
||||
query = self.to_query_object()
|
||||
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
||||
|
||||
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
||||
@@ -2148,13 +2019,14 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
|
||||
return vector_query, text_query
|
||||
|
||||
def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
|
||||
def phrase_query(self, phrase_query: bool = None) -> LanceHybridQueryBuilder:
|
||||
"""Set whether to use phrase query.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
phrase_query: bool, default True
|
||||
If True, then an unquoted string query will be wrapped in quotes.
|
||||
If True, then the query will be wrapped in quotes and
|
||||
double quotes replaced by single quotes.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -2179,25 +2051,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||||
fts_results = fts_future.result()
|
||||
vector_results = vector_future.result()
|
||||
|
||||
results = self._combine_hybrid_results(
|
||||
return 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=True,
|
||||
with_row_ids=self._with_row_id,
|
||||
)
|
||||
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(
|
||||
@@ -2638,7 +2500,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 is not None or self._upper_bound is not None:
|
||||
if self._lower_bound or self._upper_bound:
|
||||
self._vector_query.distance_range(
|
||||
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
||||
)
|
||||
@@ -2668,9 +2530,6 @@ 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:
|
||||
"""
|
||||
@@ -2680,46 +2539,11 @@ 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.
|
||||
@@ -2772,7 +2596,6 @@ class AsyncQueryBase(object):
|
||||
"""
|
||||
Include the _rowid column in the results.
|
||||
"""
|
||||
self._with_row_id = True
|
||||
self._inner.with_row_id()
|
||||
return self
|
||||
|
||||
@@ -2819,7 +2642,6 @@ 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
|
||||
@@ -2850,8 +2672,8 @@ class AsyncQueryBase(object):
|
||||
complete within the specified time, an error will be raised.
|
||||
"""
|
||||
batch_iter = await self.to_batches(timeout=timeout)
|
||||
return self._finalize_blob_query_table(
|
||||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
||||
return pa.Table.from_batches(
|
||||
await batch_iter.read_all(), schema=batch_iter.schema
|
||||
)
|
||||
|
||||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||||
@@ -2918,7 +2740,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):
|
||||
@@ -2959,36 +2781,14 @@ 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,
|
||||
"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,
|
||||
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||
)
|
||||
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,
|
||||
@@ -3773,24 +3573,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
||||
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||||
|
||||
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
|
||||
|
||||
# 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()
|
||||
fts_query.with_row_id()
|
||||
vec_query.with_row_id()
|
||||
|
||||
@@ -3806,14 +3591,8 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
fts_query=fts_query.get_query(),
|
||||
reranker=self._reranker,
|
||||
limit=self._inner.get_limit(),
|
||||
with_row_ids=True,
|
||||
with_row_ids=with_row_ids,
|
||||
)
|
||||
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,7 +28,6 @@ from lancedb._lancedb import (
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
IndexConfig,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
@@ -245,23 +244,6 @@ 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))
|
||||
@@ -930,10 +912,6 @@ 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())
|
||||
@@ -1012,19 +990,6 @@ 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,7 +12,6 @@ from .rrf import RRFReranker
|
||||
from .mrr import MRRReranker
|
||||
from .answerdotai import AnswerdotaiRerankers
|
||||
from .voyageai import VoyageAIReranker
|
||||
from .watsonx import WatsonxReranker
|
||||
|
||||
__all__ = [
|
||||
"Reranker",
|
||||
@@ -26,5 +25,4 @@ __all__ = [
|
||||
"AnswerdotaiRerankers",
|
||||
"VoyageAIReranker",
|
||||
"MRRReranker",
|
||||
"WatsonxReranker",
|
||||
]
|
||||
|
||||
@@ -1,180 +0,0 @@
|
||||
# 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,134 +2,10 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
"""Schema related utilities."""
|
||||
|
||||
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.
|
||||
|
||||
@@ -1,607 +0,0 @@
|
||||
# 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)
|
||||
+35
-226
@@ -29,14 +29,6 @@ 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 (
|
||||
@@ -96,7 +88,10 @@ from .util import (
|
||||
value_to_sql,
|
||||
)
|
||||
from .index import lang_mapping
|
||||
from .schema import blob_v2_column_paths, schema_has_blob_field
|
||||
|
||||
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||
|
||||
_VALID_BLOB_MODES = ("lazy", "bytes", "descriptions")
|
||||
|
||||
|
||||
def _should_push_down_query_table(
|
||||
@@ -105,6 +100,23 @@ def _should_push_down_query_table(
|
||||
return namespace_client is not None and "QueryTable" in pushdown_operations
|
||||
|
||||
|
||||
def _validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||
if blob_mode not in _VALID_BLOB_MODES:
|
||||
modes = ", ".join(repr(mode) for mode in _VALID_BLOB_MODES)
|
||||
raise ValueError(f"blob_mode must be one of {modes}, got {blob_mode!r}")
|
||||
|
||||
|
||||
def _field_is_blob(field: pa.Field) -> bool:
|
||||
metadata = field.metadata or {}
|
||||
return metadata.get(b"lance-encoding:blob") == b"true" or (
|
||||
metadata.get("lance-encoding:blob") == "true"
|
||||
)
|
||||
|
||||
|
||||
def _schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return any(_field_is_blob(field) for field in schema)
|
||||
|
||||
|
||||
_MODEL_BACKED_TOKENIZER_PREFIXES = ("jieba", "lindera")
|
||||
_MODEL_BACKED_TOKENIZER_ERRORS = (
|
||||
"unknown base tokenizer",
|
||||
@@ -173,7 +185,6 @@ if TYPE_CHECKING:
|
||||
UpdateFieldMetadataResult,
|
||||
DeleteResult,
|
||||
DropColumnsResult,
|
||||
FtsToken,
|
||||
LsmWriteSpec,
|
||||
MergeResult,
|
||||
UpdateResult,
|
||||
@@ -640,16 +651,6 @@ 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),
|
||||
@@ -1148,8 +1149,6 @@ 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"
|
||||
@@ -1514,31 +1513,6 @@ 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,
|
||||
@@ -1802,24 +1776,6 @@ 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]:
|
||||
"""
|
||||
@@ -2238,19 +2194,6 @@ 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.
|
||||
@@ -2446,14 +2389,9 @@ class LanceTable(Table):
|
||||
-------
|
||||
pd.DataFrame
|
||||
"""
|
||||
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)
|
||||
_validate_blob_mode(blob_mode)
|
||||
if blob_mode == "descriptions" or not _schema_has_blob_field(self.schema):
|
||||
return self.to_arrow().to_pandas(**kwargs)
|
||||
|
||||
if (
|
||||
blob_mode == "lazy"
|
||||
@@ -2462,9 +2400,6 @@ 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:
|
||||
@@ -3765,26 +3700,6 @@ 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
|
||||
@@ -3834,11 +3749,6 @@ 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]."""
|
||||
@@ -4110,16 +4020,7 @@ def _handle_bad_vector_column(
|
||||
dim = _infer_vector_dim(vec_arr)
|
||||
if dim is None:
|
||||
return data
|
||||
|
||||
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_wrong_dim = 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()
|
||||
|
||||
@@ -4154,58 +4055,17 @@ def _handle_bad_vector_column(
|
||||
raise ValueError(
|
||||
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
||||
)
|
||||
vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value)
|
||||
vec_arr = pc.if_else(
|
||||
is_bad,
|
||||
pa.scalar([fill_value] * dim, type=vec_arr.type),
|
||||
vec_arr,
|
||||
)
|
||||
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])
|
||||
@@ -4538,17 +4398,6 @@ 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.
|
||||
|
||||
@@ -4655,18 +4504,14 @@ class AsyncTable:
|
||||
-------
|
||||
pd.DataFrame
|
||||
"""
|
||||
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)
|
||||
_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)
|
||||
|
||||
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:
|
||||
@@ -5767,24 +5612,6 @@ 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.
|
||||
@@ -5886,24 +5713,6 @@ 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,24 +1,11 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Dict, List, Literal, Optional, Tuple, Union
|
||||
|
||||
from .expr import Expr
|
||||
from typing import Literal
|
||||
|
||||
# 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"]
|
||||
@@ -55,7 +42,5 @@ IndexType = Literal[
|
||||
]
|
||||
|
||||
# Tokenizer literals
|
||||
BuiltinTokenizerType = Literal[
|
||||
"simple", "raw", "whitespace", "ngram", "icu", "icu/split"
|
||||
]
|
||||
BuiltinTokenizerType = Literal["simple", "raw", "whitespace", "ngram"]
|
||||
BaseTokenizerType = BuiltinTokenizerType | str
|
||||
|
||||
@@ -177,10 +177,7 @@ def flatten_columns(tbl: pa.Table, flatten: Optional[Union[int, bool]] = None):
|
||||
continue
|
||||
else:
|
||||
break
|
||||
# `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):
|
||||
elif isinstance(flatten, int):
|
||||
if flatten <= 0:
|
||||
raise ValueError(
|
||||
"Please specify a positive integer for flatten or the boolean "
|
||||
@@ -518,38 +515,3 @@ 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)
|
||||
]
|
||||
|
||||
@@ -1,562 +0,0 @@
|
||||
# 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,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import List, Optional, Union
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -243,49 +242,6 @@ 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):
|
||||
@@ -592,22 +548,6 @@ 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()
|
||||
|
||||
@@ -786,97 +786,6 @@ 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(
|
||||
{
|
||||
@@ -1175,84 +1084,6 @@ 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")
|
||||
|
||||
|
||||
@@ -1,183 +0,0 @@
|
||||
# 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,8 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from unittest import mock
|
||||
|
||||
import lancedb
|
||||
|
||||
from lancedb.query import LanceHybridQueryBuilder
|
||||
@@ -141,20 +139,6 @@ 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")
|
||||
@@ -193,23 +177,6 @@ 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"))
|
||||
|
||||
@@ -11,7 +11,6 @@ import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb._lancedb import LsmWriteSpec
|
||||
from lancedb.index import BTree
|
||||
|
||||
SCHEMA = pa.schema(
|
||||
[
|
||||
@@ -137,76 +136,3 @@ def test_lsm_write_spec_identity_and_writer_config_defaults():
|
||||
s = s.with_writer_config_defaults({"durable_write": "false"})
|
||||
assert s.writer_config_defaults == {"durable_write": "false"}
|
||||
assert "durable_write" in repr(s)
|
||||
|
||||
|
||||
def test_get_lsm_write_spec(tmp_path):
|
||||
_db, table = _make_table(tmp_path)
|
||||
table.set_unenforced_primary_key("id")
|
||||
|
||||
# None when nothing is installed.
|
||||
assert table.get_lsm_write_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
table.create_index("id", config=BTree())
|
||||
idx_name = table.list_indices()[0].name
|
||||
|
||||
# Bucket spec round-trips, including maintained indexes and writer config
|
||||
# defaults.
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 4)
|
||||
.with_maintained_indexes([idx_name])
|
||||
.with_writer_config_defaults({"durable_write": "false"})
|
||||
)
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
assert spec.num_buckets == 4
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
assert spec.writer_config_defaults == {"durable_write": "false"}
|
||||
|
||||
# After unset, None again.
|
||||
table.unset_lsm_write_spec()
|
||||
assert table.get_lsm_write_spec() is None
|
||||
|
||||
# Identity round-trips (column recovered from the schema).
|
||||
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec.spec_type == "identity"
|
||||
assert spec.column == "id"
|
||||
table.unset_lsm_write_spec()
|
||||
|
||||
# Unsharded round-trips (no routing column).
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec.spec_type == "unsharded"
|
||||
assert spec.column is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_lsm_write_spec(tmp_path):
|
||||
db = await lancedb.connect_async(
|
||||
tmp_path, read_consistency_interval=timedelta(seconds=0)
|
||||
)
|
||||
table = await db.create_table(
|
||||
"t",
|
||||
pa.RecordBatchReader.from_batches(SCHEMA, [_batch(["seed"], [0])]),
|
||||
)
|
||||
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
await table.create_index("id", config=BTree())
|
||||
idx_name = (await table.list_indices())[0].name
|
||||
|
||||
await table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
|
||||
)
|
||||
spec = await table.get_lsm_write_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
assert spec.num_buckets == 8
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
await table.unset_lsm_write_spec()
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
|
||||
@@ -134,11 +134,8 @@ def test_split_hash_with_discard(mem_db):
|
||||
)
|
||||
|
||||
permutation_tbl = (
|
||||
# Hash a high-cardinality column: "category" has only two distinct
|
||||
# values, so whether anything is discarded would hinge on where those
|
||||
# two hashes land rather than on the discard weight.
|
||||
permutation_builder(tbl)
|
||||
.split_hash(["id"], [1, 1], discard_weight=2) # Should discard ~50%
|
||||
.split_hash(["category"], [1, 1], discard_weight=2) # Should discard ~50%
|
||||
.execute()
|
||||
)
|
||||
|
||||
@@ -938,41 +935,14 @@ def test_transform_fn(mem_db):
|
||||
try:
|
||||
import torch
|
||||
|
||||
# "torch" returns a list of per-row dicts. Default DataLoader collate
|
||||
# stacks the per-row dicts back into a dict of batched tensors.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
torch_batch = list(torch_perm.iter(10, skip_last_batch=False))[0]
|
||||
assert isinstance(torch_batch, list)
|
||||
assert len(torch_batch) == 10
|
||||
assert isinstance(torch_batch[0], dict)
|
||||
assert set(torch_batch[0].keys()) == {"id", "value"}
|
||||
assert isinstance(torch_batch[0]["id"], torch.Tensor)
|
||||
assert torch_batch[0]["id"].dtype == torch.int64
|
||||
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert set(rows[0].keys()) == {"id", "value"}
|
||||
assert isinstance(rows[0]["id"], torch.Tensor)
|
||||
|
||||
# "torch_row" returns a list of tensors, one per row.
|
||||
torch_rows = list(
|
||||
permutation.with_format("torch_row").iter(10, skip_last_batch=False)
|
||||
torch_result = list(
|
||||
permutation.with_format("torch").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_rows, list)
|
||||
assert len(torch_rows) == 10
|
||||
assert isinstance(torch_rows[0], torch.Tensor)
|
||||
assert torch_rows[0].shape == (2,)
|
||||
assert torch_rows[0].dtype == torch.int64
|
||||
|
||||
# "torch_col" stacks columns into a single 2D tensor.
|
||||
torch_col = list(
|
||||
permutation.with_format("torch_col").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_col, torch.Tensor)
|
||||
assert torch_col.shape == (2, 10)
|
||||
assert torch_col.dtype == torch.int64
|
||||
assert isinstance(torch_result, list)
|
||||
assert len(torch_result) == 10
|
||||
assert isinstance(torch_result[0], torch.Tensor)
|
||||
assert torch_result[0].shape == (2,)
|
||||
assert torch_result[0].dtype == torch.int64
|
||||
except ImportError:
|
||||
# Skip check if torch is not installed
|
||||
pass
|
||||
|
||||
@@ -11,7 +11,6 @@ import lancedb
|
||||
from lancedb.db import AsyncConnection
|
||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||
from lancedb.embeddings.registry import get_registry, register
|
||||
from lancedb.expr import col
|
||||
from lancedb.index import FTS, IvfPq
|
||||
import lancedb.pydantic
|
||||
import numpy as np
|
||||
@@ -64,71 +63,11 @@ def _blob_query_data():
|
||||
)
|
||||
|
||||
|
||||
def _create_blob_v2_query_table(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
async def _create_blob_v2_query_table_async(db, name):
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("tag", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add(
|
||||
[
|
||||
{"id": 1, "tag": "drop", "vector": [1.0, 0.0], "blob": b"one"},
|
||||
{"id": 2, "tag": "keep", "vector": [2.0, 0.0], "blob": b"two"},
|
||||
{"id": 3, "tag": "keep", "vector": [3.0, 0.0], "blob": b"three"},
|
||||
{"id": 4, "tag": "keep", "vector": [4.0, 0.0], "blob": b"four"},
|
||||
]
|
||||
)
|
||||
return table
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
|
||||
|
||||
def _assert_blob_bytes_projection(df):
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
|
||||
|
||||
def _blob_query_table(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, _blob_query_data())
|
||||
return _create_blob_v2_query_table(db, name)
|
||||
|
||||
|
||||
async def _blob_query_table_async(db, name, blob_schema):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, _blob_query_data())
|
||||
return await _create_blob_v2_query_table_async(db, name)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def table(tmpdir_factory) -> lancedb.table.Table:
|
||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||
@@ -296,11 +235,10 @@ def test_plain_scan_query_to_pandas_blob_modes(tmp_db, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
||||
def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||
pytest.importorskip("lance")
|
||||
table = _blob_query_table(
|
||||
tmp_db, f"test_query_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
table = tmp_db.create_table(
|
||||
"test_query_to_pandas_blob_projection", _blob_query_data()
|
||||
)
|
||||
|
||||
df = (
|
||||
@@ -312,8 +250,9 @@ def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
assert df["id_alias"].tolist() == [3, 4]
|
||||
assert df["payload"].tolist() == [b"three", b"four"]
|
||||
assert df["double_id"].tolist() == [6, 8]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
@@ -409,6 +348,18 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert lazy_df["id"].tolist() == [1]
|
||||
_assert_lazy_blob(lazy_df["blob"].iloc[0], b"one")
|
||||
|
||||
bytes_df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
assert bytes_df["id_alias"].tolist() == [3, 4]
|
||||
assert bytes_df["payload"].tolist() == [b"three", b"four"]
|
||||
assert bytes_df["double_id"].tolist() == [6, 8]
|
||||
|
||||
desc_df = await (
|
||||
table.query()
|
||||
.where("id = 1")
|
||||
@@ -420,31 +371,6 @@ async def test_async_plain_scan_query_to_pandas_blob_projection(tmp_db_async):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_bytes_projection(
|
||||
tmp_db_async, blob_schema
|
||||
):
|
||||
pytest.importorskip("lance")
|
||||
table = await _blob_query_table_async(
|
||||
tmp_db_async,
|
||||
f"test_async_query_to_pandas_blob_{blob_schema}_bytes",
|
||||
blob_schema,
|
||||
)
|
||||
|
||||
df = await (
|
||||
table.query()
|
||||
.where("id >= 2")
|
||||
.select({"id_alias": "id", "payload": "blob", "double_id": "id * 2"})
|
||||
.limit(2)
|
||||
.offset(1)
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
_assert_blob_bytes_projection(df)
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||
async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow(
|
||||
@@ -576,18 +502,6 @@ def test_with_row_id(table: lancedb.table.Table):
|
||||
assert rs["_rowid"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_blob_v2_query_omits_auto_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_omits_auto_rowid")
|
||||
|
||||
query_obj = table.search().select(["id", "blob"]).limit(2).to_query_object()
|
||||
assert query_obj.with_row_id is None
|
||||
|
||||
rs = table.search().select(["id", "blob"]).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in rs.column_names
|
||||
assert rs["id"].to_pylist() == [1, 2]
|
||||
|
||||
|
||||
def test_where_repeated_combines_with_and(table: lancedb.table.Table):
|
||||
# Calling where() more than once should AND the filters together instead of
|
||||
# silently replacing the previous one (regression test for #2649).
|
||||
@@ -2032,39 +1946,3 @@ def test_fast_search(tmp_path):
|
||||
# 2. Fast Search -> Should NOT include "LanceScan" (Uses Index)
|
||||
plan = table.search(q).fast_search().explain_plan(True)
|
||||
assert "LanceScan" not in plan
|
||||
|
||||
|
||||
def test_blob_v2_with_row_id_bytes_pandas(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_rowid_bytes_pandas")
|
||||
|
||||
df = (
|
||||
table.search()
|
||||
.with_row_id(True)
|
||||
.select(["id", "blob"])
|
||||
.to_pandas(blob_mode="bytes")
|
||||
)
|
||||
|
||||
assert "_rowid" in df.columns
|
||||
assert df["id"].tolist() == [1, 2, 3, 4]
|
||||
assert df["blob"].tolist() == [b"one", b"two", b"three", b"four"]
|
||||
|
||||
|
||||
def test_blob_v2_expr_projection_stash(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_expr_projection_stash")
|
||||
|
||||
hits = table.search().select({"blob_alias": col("blob")}).limit(2).to_arrow()
|
||||
|
||||
assert "_rowid" not in hits.column_names
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
|
||||
def test_blob_v2_to_batches_row_id(tmp_db):
|
||||
table = _create_blob_v2_query_table(tmp_db, "test_blob_v2_to_batches_rowid")
|
||||
|
||||
hits = table.search().select(["id", "blob"]).limit(2).to_batches().read_all()
|
||||
|
||||
assert "_rowid" in hits.column_names
|
||||
blobs = table.fetch_blobs("blob", hits)
|
||||
assert [blobs[i].as_py() for i in range(len(blobs))] == [b"one", b"two"]
|
||||
|
||||
@@ -412,12 +412,10 @@ def test_remote_permutation_is_picklable():
|
||||
content_len = int(request.headers.get("Content-Length"))
|
||||
body = json.loads(request.rfile.read(content_len))
|
||||
if "filter" in body:
|
||||
match = re.search(
|
||||
r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
|
||||
)
|
||||
offsets = [int(o.strip()) for o in match.group(1).split(",")]
|
||||
match = re.search(r"_rowoffset in \((.*?)\)", body["filter"])
|
||||
offsets = [int(offset.strip()) for offset in match.group(1).split(",")]
|
||||
else:
|
||||
offsets = list(range(len(rows)))
|
||||
offsets = rows
|
||||
table = pa.table({"a": [rows[offset] for offset in offsets]})
|
||||
|
||||
request.send_response(200)
|
||||
|
||||
@@ -23,7 +23,6 @@ from lancedb.rerankers import (
|
||||
AnswerdotaiRerankers,
|
||||
VoyageAIReranker,
|
||||
MRRReranker,
|
||||
WatsonxReranker,
|
||||
)
|
||||
from lancedb.table import LanceTable
|
||||
|
||||
@@ -728,19 +727,3 @@ def test_linear_combination_missing_fts_is_penalised():
|
||||
f"Document with FTS score (rowid 0, {scores[0]:.4f}) should beat "
|
||||
f"document with no FTS match (rowid 1, {scores[1]:.4f})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("WATSONX_API_KEY") is None
|
||||
or (
|
||||
os.environ.get("WATSONX_PROJECT_ID") is None
|
||||
and os.environ.get("WATSONX_SPACE_ID") is None
|
||||
),
|
||||
reason="WATSONX_API_KEY and one of WATSONX_PROJECT_ID / "
|
||||
"WATSONX_SPACE_ID must be set",
|
||||
)
|
||||
def test_watsonx_reranker(tmp_path):
|
||||
pytest.importorskip("ibm_watsonx_ai")
|
||||
table, schema = get_test_table(tmp_path)
|
||||
reranker = WatsonxReranker()
|
||||
_run_test_reranker(reranker, table, "single player experience", None, schema)
|
||||
|
||||
@@ -7,8 +7,7 @@ Tests for S3 bucket names containing dots.
|
||||
Related issue: https://github.com/lancedb/lancedb/issues/1898
|
||||
|
||||
These tests validate the early error checking for S3 bucket names with dots.
|
||||
When validation succeeds, eager namespace initialization may still fail later
|
||||
because these tests intentionally do not provide real S3 credentials.
|
||||
No actual S3 connection is made - validation happens before connection.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -21,13 +20,6 @@ BUCKET_WITH_DOTS_AND_AWS_REGION = ("s3://my.bucket.name", {"aws_region": "us-eas
|
||||
BUCKET_WITHOUT_DOTS = "s3://my-bucket/path"
|
||||
|
||||
|
||||
def assert_not_rejected_for_bucket_dots(connect):
|
||||
try:
|
||||
connect()
|
||||
except ValueError as err:
|
||||
assert "contains dots" not in str(err)
|
||||
|
||||
|
||||
class TestS3BucketWithDotsSync:
|
||||
"""Tests for connect()."""
|
||||
|
||||
@@ -35,22 +27,19 @@ class TestS3BucketWithDotsSync:
|
||||
with pytest.raises(ValueError, match="contains dots"):
|
||||
lancedb.connect(BUCKET_WITH_DOTS)
|
||||
|
||||
def test_bucket_with_dots_and_region_is_not_rejected(self):
|
||||
def test_bucket_with_dots_and_region_passes(self):
|
||||
uri, opts = BUCKET_WITH_DOTS_AND_REGION
|
||||
assert_not_rejected_for_bucket_dots(
|
||||
lambda: lancedb.connect(uri, storage_options=opts)
|
||||
)
|
||||
db = lancedb.connect(uri, storage_options=opts)
|
||||
assert db is not None
|
||||
|
||||
def test_bucket_with_dots_and_aws_region_is_not_rejected(self):
|
||||
def test_bucket_with_dots_and_aws_region_passes(self):
|
||||
uri, opts = BUCKET_WITH_DOTS_AND_AWS_REGION
|
||||
assert_not_rejected_for_bucket_dots(
|
||||
lambda: lancedb.connect(uri, storage_options=opts)
|
||||
)
|
||||
db = lancedb.connect(uri, storage_options=opts)
|
||||
assert db is not None
|
||||
|
||||
def test_bucket_without_dots_is_not_rejected(self):
|
||||
assert_not_rejected_for_bucket_dots(
|
||||
lambda: lancedb.connect(BUCKET_WITHOUT_DOTS)
|
||||
)
|
||||
def test_bucket_without_dots_passes(self):
|
||||
db = lancedb.connect(BUCKET_WITHOUT_DOTS)
|
||||
assert db is not None
|
||||
|
||||
|
||||
class TestS3BucketWithDotsAsync:
|
||||
@@ -62,24 +51,18 @@ class TestS3BucketWithDotsAsync:
|
||||
await lancedb.connect_async(BUCKET_WITH_DOTS)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bucket_with_dots_and_region_is_not_rejected(self):
|
||||
async def test_bucket_with_dots_and_region_passes(self):
|
||||
uri, opts = BUCKET_WITH_DOTS_AND_REGION
|
||||
try:
|
||||
await lancedb.connect_async(uri, storage_options=opts)
|
||||
except ValueError as err:
|
||||
assert "contains dots" not in str(err)
|
||||
db = await lancedb.connect_async(uri, storage_options=opts)
|
||||
assert db is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bucket_with_dots_and_aws_region_is_not_rejected(self):
|
||||
async def test_bucket_with_dots_and_aws_region_passes(self):
|
||||
uri, opts = BUCKET_WITH_DOTS_AND_AWS_REGION
|
||||
try:
|
||||
await lancedb.connect_async(uri, storage_options=opts)
|
||||
except ValueError as err:
|
||||
assert "contains dots" not in str(err)
|
||||
db = await lancedb.connect_async(uri, storage_options=opts)
|
||||
assert db is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bucket_without_dots_is_not_rejected(self):
|
||||
try:
|
||||
await lancedb.connect_async(BUCKET_WITHOUT_DOTS)
|
||||
except ValueError as err:
|
||||
assert "contains dots" not in str(err)
|
||||
async def test_bucket_without_dots_passes(self):
|
||||
db = await lancedb.connect_async(BUCKET_WITHOUT_DOTS)
|
||||
assert db is not None
|
||||
|
||||
@@ -45,32 +45,6 @@ def _blob_test_data():
|
||||
)
|
||||
|
||||
|
||||
def _blob_v2_table(db: DBConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
async def _blob_v2_table_async(db: AsyncConnection, name: str):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = await db.create_table(name, schema=schema)
|
||||
await table.add([{"id": 1, "blob": b"hello"}, {"id": 2, "blob": b"world"}])
|
||||
return table
|
||||
|
||||
|
||||
def _blob_table(db: DBConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return db.create_table(name, data=_blob_test_data())
|
||||
return _blob_v2_table(db, name)
|
||||
|
||||
|
||||
async def _blob_table_async(db: AsyncConnection, name: str, blob_schema: str):
|
||||
if blob_schema == "v1":
|
||||
return await db.create_table(name, data=_blob_test_data())
|
||||
return await _blob_v2_table_async(db, name)
|
||||
|
||||
|
||||
def _assert_lazy_blob(value, expected: bytes):
|
||||
assert hasattr(value, "readall")
|
||||
assert value.readall() == expected
|
||||
@@ -133,18 +107,6 @@ def test_table_to_pandas_blob_modes(tmp_db: DBConnection, blob_mode):
|
||||
assert not hasattr(first, "readall")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
def test_table_to_pandas_blob_bytes(tmp_db: DBConnection, blob_schema):
|
||||
pytest.importorskip("lance")
|
||||
table = _blob_table(tmp_db, f"test_to_pandas_blob_{blob_schema}_bytes", blob_schema)
|
||||
|
||||
df = table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
pd = pytest.importorskip("pandas")
|
||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
||||
@@ -156,20 +118,15 @@ def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
||||
async def test_async_table_to_pandas_blob_bytes(
|
||||
tmp_db_async: AsyncConnection, blob_schema
|
||||
):
|
||||
async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection):
|
||||
pytest.importorskip("lance")
|
||||
table = await _blob_table_async(
|
||||
tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
||||
table = await tmp_db_async.create_table(
|
||||
"test_async_to_pandas_blob_bytes", data=_blob_test_data()
|
||||
)
|
||||
|
||||
df = await table.to_pandas(blob_mode="bytes")
|
||||
|
||||
assert list(df.columns) == ["id", "blob"]
|
||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||
assert "_rowid" not in df.columns
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1611,23 +1568,16 @@ def test_create_with_nans(mem_db: DBConnection):
|
||||
"fill_test",
|
||||
data=[
|
||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
||||
{"vector": [2.1, 4.1], "item": "foo", "price": 9.0},
|
||||
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
||||
{"vector": [5], "item": "bar", "price": 22.0},
|
||||
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||
],
|
||||
on_bad_vectors="fill",
|
||||
fill_value=0.0,
|
||||
)
|
||||
assert len(table) == 5
|
||||
assert len(table) == 3
|
||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||
filled_vectors = {
|
||||
row["price"]: row["vector"]
|
||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
||||
}
|
||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||
v = arrow_tbl["vector"].to_pylist()[0]
|
||||
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||
|
||||
|
||||
def test_add_with_nans(mem_db: DBConnection):
|
||||
@@ -1670,21 +1620,15 @@ def test_add_with_nans(mem_db: DBConnection):
|
||||
data=[
|
||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
||||
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
||||
{"vector": [5], "item": "bar", "price": 22.0},
|
||||
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||
],
|
||||
on_bad_vectors="fill",
|
||||
fill_value=0.0,
|
||||
)
|
||||
assert len(table) == 4
|
||||
assert len(table) == 3
|
||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||
filled_vectors = {
|
||||
row["price"]: row["vector"]
|
||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
||||
}
|
||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
||||
v = arrow_tbl["vector"].to_pylist()[0]
|
||||
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||
|
||||
|
||||
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
|
||||
@@ -1845,9 +1789,7 @@ def test_on_bad_vectors_fill_preserves_arrow_nested_vector_type(mem_db: DBConnec
|
||||
fill_value=0.0,
|
||||
)
|
||||
|
||||
vector = table.to_arrow()["vector"]
|
||||
assert vector.type == pa.list_(pa.float32())
|
||||
assert vector.to_pylist() == [[1.0, 2.0], [0.0, 3.0]]
|
||||
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0], [0.0, 0.0]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -148,47 +148,15 @@ def test_permutation_dataloader(mem_db):
|
||||
for batch in dataloader:
|
||||
assert batch["a"].size(0) == 10
|
||||
|
||||
# "torch" produces a list of per-row dicts per batch. The default
|
||||
# DataLoader collate stacks the per-row dicts back into a batched dict.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
batch = next(torch_perm.iter(10, skip_last_batch=False))
|
||||
assert isinstance(batch, list)
|
||||
assert len(batch) == 10
|
||||
assert isinstance(batch[0], dict)
|
||||
assert isinstance(batch[0]["a"], torch.Tensor)
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert isinstance(rows[0]["a"], torch.Tensor)
|
||||
dataloader = torch.utils.data.DataLoader(torch_perm, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
# Spawn-based workers exercise the pickle round-trip path: the new
|
||||
# transform-as-list shape must survive pickling so workers produce the
|
||||
# same per-row dicts the parent does.
|
||||
spawn_loader = torch.utils.data.DataLoader(
|
||||
torch_perm,
|
||||
batch_size=10,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
for batch in spawn_loader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
|
||||
# "torch_row" returns a list of row tensors. Works with the default
|
||||
# DataLoader collate (stacks rows into 2D).
|
||||
row_perm = permutation.with_format("torch_row")
|
||||
dataloader = torch.utils.data.DataLoader(row_perm, batch_size=10, shuffle=True)
|
||||
permutation = permutation.with_format("torch")
|
||||
dataloader = torch.utils.data.DataLoader(permutation, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 10
|
||||
assert batch.size(1) == 1
|
||||
|
||||
col_perm = permutation.with_format("torch_col")
|
||||
permutation = permutation.with_format("torch_col")
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
col_perm, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
permutation, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 1
|
||||
|
||||
@@ -13,7 +13,6 @@ from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
from lancedb.table import (
|
||||
_append_vector_columns,
|
||||
_cast_to_target_schema,
|
||||
_fill_bad_vector_values,
|
||||
_handle_bad_vectors,
|
||||
_into_pyarrow_reader,
|
||||
_infer_target_schema,
|
||||
@@ -26,42 +25,10 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import get_uri_scheme, join_uri, value_to_sql
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
def _struct_table() -> pa.Table:
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
"nested": pa.array([{"a": 1, "b": 2}, {"a": 3, "b": 4}]),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_flatten_columns():
|
||||
tbl = _struct_table()
|
||||
|
||||
# None / False mean "do not flatten": the struct column is preserved.
|
||||
# `False` is a regression guard: because bool is a subclass of int it used
|
||||
# to fall into the integer branch and raise ValueError (see issue).
|
||||
for no_flatten in (None, False):
|
||||
result = flatten_columns(tbl, no_flatten)
|
||||
assert result.column_names == ["id", "nested"]
|
||||
|
||||
# True flattens all nested levels.
|
||||
flattened = flatten_columns(tbl, True)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# A positive integer flattens up to that depth.
|
||||
flattened = flatten_columns(tbl, 1)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# Non-positive integers are still rejected.
|
||||
with pytest.raises(ValueError):
|
||||
flatten_columns(tbl, 0)
|
||||
|
||||
|
||||
def test_normalize_uri():
|
||||
uris = [
|
||||
"relative/path",
|
||||
@@ -288,9 +255,7 @@ def test_append_vector_columns():
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
vector = pa.array(
|
||||
[[1.0, 2.0], [3.0], [4.0, 5.0], [6.0, 7.0, 8.0], [None, 9.0], None]
|
||||
)
|
||||
vector = pa.array([[1.0, 2.0], [3.0], [4.0, 5.0]])
|
||||
schema = pa.schema({"vector": pa.list_(pa.float64())})
|
||||
data = pa.table({"vector": vector}, schema=schema)
|
||||
|
||||
@@ -316,54 +281,15 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||
).read_all()
|
||||
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0], [None, 9.0]])
|
||||
expected = pa.array([[1.0, 2.0], [4.0, 5.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array(
|
||||
[
|
||||
[1.0, 2.0],
|
||||
[3.0, 42.0],
|
||||
[4.0, 5.0],
|
||||
[6.0, 7.0],
|
||||
[None, 9.0],
|
||||
[42.0, 42.0],
|
||||
]
|
||||
)
|
||||
expected = pa.array([[1.0, 2.0], [42.0, 42.0], [4.0, 5.0]])
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0], None, [None, 9.0], None])
|
||||
expected = pa.array([[1.0, 2.0], None, [4.0, 5.0]])
|
||||
|
||||
assert output["vector"].combine_chunks() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("vector_type", "vectors", "expected"),
|
||||
[
|
||||
(
|
||||
pa.list_(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.large_list(pa.float64()),
|
||||
[[1.0, float("nan")], [2.0], None, [None, 3.0], [4.0, 5.0, 6.0]],
|
||||
[[1.0, 42.0], [2.0, 42.0], [42.0, 42.0], [None, 3.0], [4.0, 5.0]],
|
||||
),
|
||||
(
|
||||
pa.list_(pa.float64(), 2),
|
||||
[[1.0, float("nan")], None, [None, 3.0]],
|
||||
[[1.0, 42.0], [42.0, 42.0], [None, 3.0]],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_fill_bad_vector_values_arrow_types(vector_type, vectors, expected):
|
||||
arr = pa.array([[0.0, 0.0], *vectors, [9.0, 9.0]], type=vector_type)
|
||||
arr = arr.slice(1, len(vectors))
|
||||
|
||||
actual = _fill_bad_vector_values(arr, dim=2, fill_value=42.0)
|
||||
|
||||
assert actual.type == vector_type
|
||||
assert actual.to_pylist() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||
def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
|
||||
@@ -393,7 +319,7 @@ def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||
if on_bad_vectors == "drop":
|
||||
expected = pa.array([[3.0, 4.0]])
|
||||
elif on_bad_vectors == "fill":
|
||||
expected = pa.array([[1.0, 42.0], [3.0, 4.0]])
|
||||
expected = pa.array([[42.0, 42.0], [3.0, 4.0]])
|
||||
elif on_bad_vectors == "null":
|
||||
expected = pa.array([None, [3.0, 4.0]])
|
||||
|
||||
|
||||
+8
-60
@@ -7,14 +7,12 @@
|
||||
//! build type-safe filter / projection expressions that map directly to
|
||||
//! DataFusion [`Expr`] nodes, bypassing SQL string parsing.
|
||||
|
||||
use std::ops::{Add, Div, Mul, Not, Sub};
|
||||
|
||||
use arrow::{datatypes::DataType, pyarrow::PyArrowType};
|
||||
use datafusion_common::ScalarValue;
|
||||
use lancedb::expr::{
|
||||
DfExpr, col as ldb_col, contains, expr_cast, is_in, lit as df_lit, lower, upper,
|
||||
};
|
||||
use pyo3::types::{PyBytes, PyDate, PyDateTime};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, PyAny, PyResult, exceptions::PyValueError, prelude::*, pyfunction};
|
||||
|
||||
/// A type-safe DataFusion expression.
|
||||
@@ -65,30 +63,30 @@ impl PyExpr {
|
||||
Self(self.0.clone().or(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Logical NOT.
|
||||
fn not_(&self) -> Self {
|
||||
use std::ops::Not;
|
||||
Self(self.0.clone().not())
|
||||
}
|
||||
|
||||
// ── arithmetic ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Add expressions.
|
||||
fn add(&self, other: &Self) -> Self {
|
||||
use std::ops::Add;
|
||||
Self(self.0.clone().add(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Subtract expressions.
|
||||
fn sub(&self, other: &Self) -> Self {
|
||||
use std::ops::Sub;
|
||||
Self(self.0.clone().sub(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Multiply expressions.
|
||||
fn mul(&self, other: &Self) -> Self {
|
||||
use std::ops::Mul;
|
||||
Self(self.0.clone().mul(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Divide expressions.
|
||||
fn div(&self, other: &Self) -> Self {
|
||||
use std::ops::Div;
|
||||
Self(self.0.clone().div(other.0.clone()))
|
||||
}
|
||||
|
||||
@@ -155,8 +153,7 @@ pub fn expr_col(name: &str) -> PyExpr {
|
||||
|
||||
/// Create a literal value expression.
|
||||
///
|
||||
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`, `date`,
|
||||
/// `datetime`, `Decimal`.
|
||||
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`.
|
||||
#[pyfunction]
|
||||
pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
// bool must be checked before int because bool is a subclass of int in Python
|
||||
@@ -166,19 +163,6 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
if let Ok(i) = value.extract::<i64>() {
|
||||
return Ok(PyExpr(df_lit(i)));
|
||||
}
|
||||
// Decimal must be checked before f64: Python's Decimal implements __float__,
|
||||
// so value.extract::<f64>() would succeed and silently truncate the value to
|
||||
// f64, losing precision. Build a Decimal128 scalar to preserve it instead.
|
||||
if value.get_type().name()? == "Decimal" {
|
||||
let s = value.call_method0("__str__")?.extract::<String>()?;
|
||||
// Parse the decimal string into an i128 value, precision, and scale.
|
||||
let (val, precision, scale) = parse_decimal(&s)?;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Decimal128(
|
||||
Some(val),
|
||||
precision,
|
||||
scale,
|
||||
))));
|
||||
}
|
||||
if let Ok(f) = value.extract::<f64>() {
|
||||
return Ok(PyExpr(df_lit(f)));
|
||||
}
|
||||
@@ -189,48 +173,12 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
let bytes = value.extract::<Vec<u8>>()?;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Binary(Some(bytes)))));
|
||||
}
|
||||
|
||||
// datetime.datetime is a subclass of datetime.date, so it must be checked first.
|
||||
if let Ok(dt) = value.cast::<PyDateTime>() {
|
||||
let ts: f64 = dt.call_method0("timestamp")?.extract()?;
|
||||
let micros = (ts * 1_000_000.0).round() as i64;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::TimestampMicrosecond(
|
||||
Some(micros),
|
||||
None,
|
||||
))));
|
||||
}
|
||||
if let Ok(d) = value.cast::<PyDate>() {
|
||||
let ordinal: i32 = d.call_method0("toordinal")?.extract()?;
|
||||
let days = ordinal - 719163; // Unix epoch is 1970-01-01
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Date32(Some(days)))));
|
||||
}
|
||||
|
||||
Err(PyValueError::new_err(format!(
|
||||
"unsupported literal type: {}. Supported: bool, int, float, str, bytes, date, datetime, Decimal",
|
||||
"unsupported literal type: {}. Supported: bool, int, float, str, bytes",
|
||||
value.get_type().name()?
|
||||
)))
|
||||
}
|
||||
|
||||
fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> {
|
||||
let s = s.trim();
|
||||
let dot_pos = s.find('.');
|
||||
let scale = if let Some(pos) = dot_pos {
|
||||
(s.len() - pos - 1) as i8
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let digits = s.replace('.', "");
|
||||
let val = digits
|
||||
.parse::<i128>()
|
||||
.map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?;
|
||||
|
||||
// Precision is total number of digits
|
||||
let precision = digits.trim_start_matches('-').len() as u8;
|
||||
|
||||
Ok((val, precision, scale))
|
||||
}
|
||||
|
||||
/// Call an arbitrary registered SQL function by name.
|
||||
///
|
||||
/// See `lancedb::expr::func` for the list of supported function names.
|
||||
|
||||
+2
-15
@@ -15,8 +15,8 @@ use pyo3::{
|
||||
use query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
||||
use session::Session;
|
||||
use table::{
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
||||
LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult,
|
||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, LsmWriteSpec,
|
||||
MergeResult, Table, UpdateFieldMetadataResult, UpdateResult,
|
||||
};
|
||||
|
||||
pub mod arrow;
|
||||
@@ -27,7 +27,6 @@ pub mod header;
|
||||
pub mod index;
|
||||
pub mod namespace;
|
||||
pub mod oauth;
|
||||
pub mod otel;
|
||||
pub mod permutation;
|
||||
pub mod query;
|
||||
pub mod runtime;
|
||||
@@ -44,7 +43,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<PyBlobFile>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
m.add_class::<Query>()?;
|
||||
m.add_class::<FTSQuery>()?;
|
||||
@@ -60,23 +58,12 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<DeleteResult>()?;
|
||||
m.add_class::<DropColumnsResult>()?;
|
||||
m.add_class::<UpdateResult>()?;
|
||||
m.add_class::<FtsToken>()?;
|
||||
m.add_class::<PyAsyncPermutationBuilder>()?;
|
||||
m.add_class::<PyPermutationReader>()?;
|
||||
m.add_class::<PyExpr>()?;
|
||||
// OpenTelemetry metrics bridge
|
||||
m.add_class::<otel::PyMetricPoint>()?;
|
||||
m.add_class::<otel::PyMetricDescription>()?;
|
||||
m.add_function(wrap_pyfunction!(
|
||||
otel::register_lancedb_metrics_recorder,
|
||||
m
|
||||
)?)?;
|
||||
m.add_function(wrap_pyfunction!(otel::lancedb_metrics_catalog, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(otel::snapshot_lancedb_metrics, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace_client, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(table::tokenize, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?;
|
||||
|
||||
@@ -1,96 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Python-facing wrappers 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 PyO3
|
||||
//! classes and exposes the three entry points to Python, where
|
||||
//! `lancedb/otel.py` bridges them into the user's OpenTelemetry `MeterProvider`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use lancedb::metrics_otel::{MetricPoint, MetricValue};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
/// One metric data point exposed to Python. For counters and gauges only
|
||||
/// `value` is set; for histograms `buckets` (cumulative `le` counts), `count`,
|
||||
/// and `sum` are set.
|
||||
#[pyclass(name = "MetricPoint", get_all)]
|
||||
pub struct PyMetricPoint {
|
||||
name: String,
|
||||
kind: String,
|
||||
attributes: HashMap<String, String>,
|
||||
value: Option<f64>,
|
||||
buckets: Option<Vec<(String, u64)>>,
|
||||
count: Option<u64>,
|
||||
sum: Option<f64>,
|
||||
}
|
||||
|
||||
impl From<MetricPoint> for PyMetricPoint {
|
||||
fn from(point: MetricPoint) -> 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), Some(count), Some(sum)),
|
||||
};
|
||||
Self {
|
||||
name: point.name,
|
||||
kind,
|
||||
attributes: point.attributes,
|
||||
value,
|
||||
buckets,
|
||||
count,
|
||||
sum,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A described metric, used by the Python layer to create instruments up front.
|
||||
#[pyclass(name = "MetricDescription", get_all)]
|
||||
pub struct PyMetricDescription {
|
||||
name: String,
|
||||
kind: String,
|
||||
unit: Option<String>,
|
||||
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.
|
||||
#[pyfunction]
|
||||
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.
|
||||
#[pyfunction]
|
||||
pub fn lancedb_metrics_catalog() -> Vec<PyMetricDescription> {
|
||||
lancedb::metrics_otel::metrics_catalog()
|
||||
.into_iter()
|
||||
.map(|desc| PyMetricDescription {
|
||||
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.
|
||||
///
|
||||
/// The read is lock-free but not O(1): it walks every registered series and
|
||||
/// allocates owned copies of their names and labels. The GIL is released across
|
||||
/// that work so a periodic collection doesn't stall other Python threads.
|
||||
#[pyfunction]
|
||||
pub fn snapshot_lancedb_metrics(py: Python<'_>) -> Vec<PyMetricPoint> {
|
||||
let points = py.detach(lancedb::metrics_otel::snapshot_metrics);
|
||||
points.into_iter().map(PyMetricPoint::from).collect()
|
||||
}
|
||||
@@ -79,14 +79,13 @@ impl PyAsyncPermutationBuilder {
|
||||
|
||||
#[pymethods]
|
||||
impl PyAsyncPermutationBuilder {
|
||||
#[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, clump_size=None, split_names=None))]
|
||||
#[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, split_names=None))]
|
||||
pub fn split_random(
|
||||
slf: PyRefMut<'_, Self>,
|
||||
ratios: Option<Vec<f64>>,
|
||||
counts: Option<Vec<u64>>,
|
||||
fixed: Option<u64>,
|
||||
seed: Option<u64>,
|
||||
clump_size: Option<u64>,
|
||||
split_names: Option<Vec<String>>,
|
||||
) -> PyResult<Self> {
|
||||
// Check that exactly one split type is provided
|
||||
@@ -112,14 +111,7 @@ impl PyAsyncPermutationBuilder {
|
||||
};
|
||||
|
||||
slf.modify(|builder| {
|
||||
builder.with_split_strategy(
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
},
|
||||
split_names,
|
||||
)
|
||||
builder.with_split_strategy(SplitStrategy::Random { seed, sizes }, split_names)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+5
-239
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use crate::runtime::{block_on, future_into_py};
|
||||
use crate::runtime::future_into_py;
|
||||
use crate::{
|
||||
connection::Connection,
|
||||
error::PythonErrorExt,
|
||||
@@ -12,23 +12,19 @@ use crate::{
|
||||
table::scannable::PyScannable,
|
||||
};
|
||||
use arrow::{
|
||||
array::{Array, LargeBinaryArray},
|
||||
datatypes::{DataType, Schema},
|
||||
ffi_stream::ArrowArrayStreamReader,
|
||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||
};
|
||||
use lancedb::blob::BlobFile;
|
||||
use lancedb::index::scalar::FtsIndexBuilder;
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use lancedb::tokenize as lancedb_tokenize;
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
exceptions::{PyRuntimeError, PyValueError},
|
||||
pyclass, pyfunction, pymethods,
|
||||
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
|
||||
pyclass, pymethods,
|
||||
types::{IntoPyDict, PyAnyMethods, PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
mod scannable;
|
||||
@@ -326,12 +322,6 @@ impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AddColumnsResult {
|
||||
@@ -416,150 +406,6 @@ impl From<lancedb::table::DropColumnsResult> for DropColumnsResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Lazy blob handle from ``Table.fetch_blob_files``.
|
||||
#[pyclass(name = "BlobFile")]
|
||||
pub struct PyBlobFile {
|
||||
inner: Arc<BlobFile>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyBlobFile {
|
||||
fn read_bytes(self_: PyRef<'_, Self>) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
pub fn read(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let bytes = inner
|
||||
.read()
|
||||
.await
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Python::attach(|py| Ok(PyBytes::new(py, bytes.as_ref()).unbind()))
|
||||
})
|
||||
}
|
||||
|
||||
fn close(self_: PyRef<'_, Self>) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.close().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob close failed: {e}")))
|
||||
}
|
||||
|
||||
fn is_closed(self_: PyRef<'_, Self>) -> bool {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.is_closed().await })
|
||||
}
|
||||
|
||||
fn seek(self_: PyRef<'_, Self>, position: u64) -> PyResult<()> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.seek(position).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob seek failed: {e}")))
|
||||
}
|
||||
|
||||
fn tell(self_: PyRef<'_, Self>) -> PyResult<u64> {
|
||||
let inner = self_.inner.clone();
|
||||
block_on(async move { inner.tell().await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob tell failed: {e}")))
|
||||
}
|
||||
|
||||
fn size(self_: PyRef<'_, Self>) -> u64 {
|
||||
self_.inner.size()
|
||||
}
|
||||
|
||||
/// Read a blob-local byte range without moving the cursor.
|
||||
fn read_range(self_: PyRef<'_, Self>, offset: u64, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let end = offset
|
||||
.checked_add(length as u64)
|
||||
.ok_or_else(|| PyValueError::new_err("offset + length overflowed"))?;
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_range(offset..end).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read_range failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
|
||||
fn read_up_to(self_: PyRef<'_, Self>, length: usize) -> PyResult<Py<PyBytes>> {
|
||||
let inner = self_.inner.clone();
|
||||
let bytes = block_on(async move { inner.read_up_to(length).await })
|
||||
.map_err(|e| PyRuntimeError::new_err(format!("blob read failed: {e}")))?;
|
||||
Ok(PyBytes::new(self_.py(), bytes.as_ref()).unbind())
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FtsToken {
|
||||
pub text: String,
|
||||
pub position: u32,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl FtsToken {
|
||||
pub fn __repr__(&self) -> String {
|
||||
format!("FtsToken(text={:?}, position={})", self.text, self.position)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<LanceDbFtsToken> for FtsToken {
|
||||
fn from(token: LanceDbFtsToken) -> Self {
|
||||
Self {
|
||||
text: token.text,
|
||||
position: token.position,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction(signature = (
|
||||
query,
|
||||
*,
|
||||
base_tokenizer = "simple".to_string(),
|
||||
language = "English".to_string(),
|
||||
max_token_length = Some(40),
|
||||
lower_case = true,
|
||||
stem = true,
|
||||
remove_stop_words = true,
|
||||
ascii_folding = true,
|
||||
ngram_min_length = 3,
|
||||
ngram_max_length = 3,
|
||||
prefix_only = false
|
||||
))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn tokenize(
|
||||
query: String,
|
||||
base_tokenizer: String,
|
||||
language: String,
|
||||
max_token_length: Option<u32>,
|
||||
lower_case: bool,
|
||||
stem: bool,
|
||||
remove_stop_words: bool,
|
||||
ascii_folding: bool,
|
||||
ngram_min_length: u32,
|
||||
ngram_max_length: u32,
|
||||
prefix_only: bool,
|
||||
) -> PyResult<Vec<FtsToken>> {
|
||||
let params = FtsIndexBuilder::default()
|
||||
.base_tokenizer(base_tokenizer)
|
||||
.language(&language)
|
||||
.map_err(|_| {
|
||||
PyValueError::new_err(format!(
|
||||
"LanceDB does not support the requested language: '{}'",
|
||||
language
|
||||
))
|
||||
})?
|
||||
.max_token_length(max_token_length.map(|value| value as usize))
|
||||
.lower_case(lower_case)
|
||||
.stem(stem)
|
||||
.remove_stop_words(remove_stop_words)
|
||||
.ascii_folding(ascii_folding)
|
||||
.ngram_min_length(ngram_min_length)
|
||||
.ngram_max_length(ngram_max_length)
|
||||
.ngram_prefix_only(prefix_only);
|
||||
let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect())
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Table {
|
||||
// We keep a copy of the name to use if the inner table is dropped
|
||||
@@ -858,29 +704,6 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (query, *, column=None, index_name=None))]
|
||||
pub fn tokenize(
|
||||
self_: PyRef<'_, Self>,
|
||||
query: String,
|
||||
column: Option<String>,
|
||||
index_name: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let tokens = match (column.as_deref(), index_name.as_deref()) {
|
||||
(Some(_), Some(_)) | (None, None) => {
|
||||
return Err(PyValueError::new_err(
|
||||
"Specify exactly one of 'column' or 'index_name'",
|
||||
));
|
||||
}
|
||||
(Some(column), None) => inner.tokenize_with_column(&query, column).await,
|
||||
(None, Some(index_name)) => inner.tokenize(&query, index_name).await,
|
||||
}
|
||||
.infer_error()?;
|
||||
Ok(tokens.into_iter().map(FtsToken::from).collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn index_stats(self_: PyRef<'_, Self>, index_name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -1072,55 +895,6 @@ impl Table {
|
||||
))
|
||||
}
|
||||
|
||||
/// Names of the blob v2 columns declared on this table, in declaration order.
|
||||
pub fn blob_columns(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.blob_columns().await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Read blob bytes for `row_ids` from blob v2 column `column`.
|
||||
#[pyo3(signature = (column, row_ids))]
|
||||
pub fn fetch_blobs(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
row_ids: Vec<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let blobs: LargeBinaryArray = inner
|
||||
.fetch_blobs(column.as_str(), &row_ids)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Python::attach(|py| blobs.to_data().to_pyarrow(py).map(|obj| obj.unbind()))
|
||||
})
|
||||
}
|
||||
|
||||
/// Open lazy blob handles for `row_ids` from blob v2 column `column`.
|
||||
#[pyo3(signature = (column, row_ids))]
|
||||
pub fn fetch_blob_files(
|
||||
self_: PyRef<'_, Self>,
|
||||
column: String,
|
||||
row_ids: Vec<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let handles = inner
|
||||
.fetch_blob_files(column.as_str(), &row_ids)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(handles
|
||||
.into_iter()
|
||||
.map(|handle| {
|
||||
handle.map(|file| PyBlobFile {
|
||||
inner: Arc::new(file),
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
|
||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
||||
pub fn optimize(
|
||||
@@ -1255,14 +1029,6 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let spec = inner.get_lsm_write_spec().await.infer_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
|
||||
@@ -3,14 +3,10 @@
|
||||
|
||||
"""Tests for the type-safe expression builder API."""
|
||||
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from decimal import Decimal
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
from lancedb.expr import Expr, col, func, lit
|
||||
from lancedb.expr import Expr, col, lit, func
|
||||
|
||||
|
||||
# ── unit tests for Expr construction ─────────────────────────────────────────
|
||||
@@ -58,28 +54,6 @@ 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):
|
||||
@@ -168,20 +142,6 @@ 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):
|
||||
@@ -322,40 +282,6 @@ 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 ─────────────────
|
||||
|
||||
|
||||
@@ -506,72 +432,6 @@ 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 ───────────────────────────────────
|
||||
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
# 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.
|
||||
|
||||
|
||||
def _metrics_by_name(reader):
|
||||
data = reader.get_metrics_data()
|
||||
result = {}
|
||||
for resource_metrics in data.resource_metrics:
|
||||
for scope_metrics in resource_metrics.scope_metrics:
|
||||
for metric in scope_metrics.metrics:
|
||||
result[metric.name] = metric
|
||||
return result
|
||||
|
||||
|
||||
def test_instrument_lancedb_metrics_exports_object_store_metrics(tmp_path):
|
||||
pytest.importorskip("opentelemetry.sdk.metrics")
|
||||
from lancedb.otel import instrument_lancedb_metrics
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
|
||||
|
||||
reader = InMemoryMetricReader()
|
||||
provider = MeterProvider(metric_readers=[reader])
|
||||
assert instrument_lancedb_metrics(provider)
|
||||
|
||||
# The catalog is populated once the recorder is installed.
|
||||
from lancedb._lancedb import lancedb_metrics_catalog
|
||||
|
||||
catalog = {desc.name: desc for desc in lancedb_metrics_catalog()}
|
||||
# Every metric kind emitted by the object store must be described so it is
|
||||
# surfaced by the bridge (counter, histogram, and gauge).
|
||||
assert catalog["lance_object_store_requests_total"].kind == "counter"
|
||||
assert catalog["lance_object_store_request_duration_seconds"].kind == "histogram"
|
||||
assert catalog["lance_object_store_in_flight_requests"].kind == "gauge"
|
||||
assert catalog["lance_object_store_retryable_responses_total"].kind == "counter"
|
||||
|
||||
# Generate object store activity on the local filesystem (scheme "file").
|
||||
db = lancedb.connect(str(tmp_path))
|
||||
table = db.create_table("t", pa.table({"id": pa.array(range(256))}))
|
||||
assert table.count_rows() == 256
|
||||
assert table.to_arrow().num_rows == 256
|
||||
|
||||
metrics = _metrics_by_name(reader)
|
||||
|
||||
requests = metrics["lance_object_store_requests_total"]
|
||||
points = list(requests.data.data_points)
|
||||
assert points, "expected at least one request data point"
|
||||
# Object store metrics are labelled by `operation` and `base` (the store
|
||||
# scheme, e.g. "file", by default).
|
||||
assert all("base" in p.attributes and "operation" in p.attributes for p in points)
|
||||
assert sum(p.value for p in points) > 0
|
||||
|
||||
# Histograms are decomposed into bucket / count / sum observable counters.
|
||||
bucket = metrics["lance_object_store_request_duration_seconds_bucket"]
|
||||
bucket_points = list(bucket.data.data_points)
|
||||
assert bucket_points
|
||||
assert all("le" in p.attributes for p in bucket_points)
|
||||
# The implicit +Inf bucket must be present and is the cumulative maximum.
|
||||
assert any(p.attributes["le"] == "+Inf" for p in bucket_points)
|
||||
|
||||
count = metrics["lance_object_store_request_duration_seconds_count"]
|
||||
assert sum(p.value for p in count.data.data_points) > 0
|
||||
|
||||
# The `_sum` instrument must also be wired and report positive latency.
|
||||
duration_sum = metrics["lance_object_store_request_duration_seconds_sum"]
|
||||
assert sum(p.value for p in duration_sum.data.data_points) > 0
|
||||
|
||||
# Unit handling: only `_sum` keeps the histogram's unit (seconds); `_bucket`
|
||||
# and `_count` observe cumulative counts and are unitless.
|
||||
assert duration_sum.unit == "s"
|
||||
assert bucket.unit == ""
|
||||
assert count.unit == ""
|
||||
|
||||
|
||||
def test_snapshot_empty_before_install_is_safe():
|
||||
# snapshot is callable regardless of installation state and never raises.
|
||||
from lancedb._lancedb import snapshot_lancedb_metrics
|
||||
|
||||
assert isinstance(snapshot_lancedb_metrics(), list)
|
||||
|
||||
|
||||
def test_instrument_warns_when_recorder_unavailable(monkeypatch):
|
||||
# A foreign `metrics` recorder already installed -> register returns False;
|
||||
# instrument_lancedb_metrics must warn and return False without instrumenting.
|
||||
pytest.importorskip("opentelemetry.sdk.metrics")
|
||||
import lancedb.otel as otel
|
||||
from opentelemetry.sdk.metrics import MeterProvider
|
||||
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
|
||||
|
||||
monkeypatch.setattr(otel, "register_lancedb_metrics_recorder", lambda: False)
|
||||
|
||||
reader = InMemoryMetricReader()
|
||||
provider = MeterProvider(metric_readers=[reader])
|
||||
with pytest.warns(UserWarning, match="recorder"):
|
||||
assert otel.instrument_lancedb_metrics(provider) is False
|
||||
Generated
+18
-63
@@ -780,7 +780,7 @@ name = "cuda-bindings"
|
||||
version = "13.3.1"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "cuda-pathfinder", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "cuda-pathfinder" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
|
||||
@@ -815,37 +815,37 @@ wheels = [
|
||||
|
||||
[package.optional-dependencies]
|
||||
cublas = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cublas", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
cudart = [
|
||||
{ name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
cufft = [
|
||||
{ name = "nvidia-cufft", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
cufile = [
|
||||
{ name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
|
||||
]
|
||||
cupti = [
|
||||
{ name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
curand = [
|
||||
{ name = "nvidia-curand", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
cusolver = [
|
||||
{ name = "nvidia-cusolver", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
cusparse = [
|
||||
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
nvjitlink = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
nvrtc = [
|
||||
{ name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
nvtx = [
|
||||
{ name = "nvidia-nvtx", marker = "(python_full_version < '3.14' and sys_platform == 'win32') or sys_platform == 'linux'" },
|
||||
{ name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1909,9 +1909,6 @@ embeddings = [
|
||||
{ name = "sentencepiece" },
|
||||
{ name = "torch" },
|
||||
]
|
||||
otel = [
|
||||
{ name = "opentelemetry-api" },
|
||||
]
|
||||
pylance = [
|
||||
{ name = "pylance" },
|
||||
]
|
||||
@@ -1926,7 +1923,6 @@ tests = [
|
||||
{ name = "boto3" },
|
||||
{ name = "datafusion" },
|
||||
{ name = "duckdb" },
|
||||
{ name = "opentelemetry-sdk" },
|
||||
{ name = "pandas", version = "2.2.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
|
||||
{ name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||
{ name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" },
|
||||
@@ -1967,8 +1963,6 @@ requires-dist = [
|
||||
{ name = "open-clip-torch", marker = "extra == 'clip'" },
|
||||
{ name = "open-clip-torch", marker = "extra == 'embeddings'", specifier = ">=2.20.0" },
|
||||
{ name = "openai", marker = "extra == 'embeddings'", specifier = ">=1.6.1" },
|
||||
{ name = "opentelemetry-api", marker = "extra == 'otel'" },
|
||||
{ name = "opentelemetry-sdk", marker = "extra == 'tests'", specifier = ">=1.30.0" },
|
||||
{ name = "overrides", marker = "python_full_version < '3.12'", specifier = ">=0.7" },
|
||||
{ name = "packaging", specifier = ">=23.0" },
|
||||
{ name = "pandas", marker = "extra == 'tests'", specifier = ">=1.4" },
|
||||
@@ -2000,7 +1994,7 @@ requires-dist = [
|
||||
{ name = "transformers", marker = "extra == 'siglip'", specifier = ">=4.41.0" },
|
||||
{ name = "typing-extensions", marker = "python_full_version < '3.11' and extra == 'dev'", specifier = ">=4.0.0" },
|
||||
]
|
||||
provides-extras = ["azure", "clip", "dev", "docs", "embeddings", "otel", "pylance", "siglip", "tests"]
|
||||
provides-extras = ["azure", "clip", "dev", "docs", "embeddings", "pylance", "siglip", "tests"]
|
||||
|
||||
[[package]]
|
||||
name = "lomond"
|
||||
@@ -2781,7 +2775,7 @@ name = "nvidia-cudnn-cu13"
|
||||
version = "9.19.0.56"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cublas" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" },
|
||||
@@ -2793,7 +2787,7 @@ name = "nvidia-cufft"
|
||||
version = "12.0.0.61"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
|
||||
@@ -2823,9 +2817,9 @@ name = "nvidia-cusolver"
|
||||
version = "12.0.4.66"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-cublas", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-cublas" },
|
||||
{ name = "nvidia-cusparse" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
|
||||
@@ -2837,7 +2831,7 @@ name = "nvidia-cusparse"
|
||||
version = "12.6.3.3"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.14' and sys_platform == 'emscripten') or (python_full_version < '3.14' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
|
||||
{ name = "nvidia-nvjitlink" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
|
||||
@@ -2940,45 +2934,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f6/46/180e14be801a75bc13f234cb1b594b232adeb9c84e60a9ab1832e8333591/openai-2.40.0-py3-none-any.whl", hash = "sha256:2b205637ff214477f9ce9ab035e9f494db0e3fa8f1e599008953735fbf6ff1ff", size = 1350935, upload-time = "2026-06-01T21:48:21.462Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-api"
|
||||
version = "1.43.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/ae/cc/e4c9584181f86494df0f6bdec1a4f3280c50db44704dc2a407e994fc87bb/opentelemetry_api-1.43.0.tar.gz", hash = "sha256:107d0d03857ea8fc7c5fcbbbd83f800c281f0d560553d61c1d675fccfd1761c1", size = 73476, upload-time = "2026-06-24T15:19:55.323Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/17/83/6dba32b85f31868400440dc7ad2ca1eab94cbbf3a7b0459ed39f8311a9e2/opentelemetry_api-1.43.0-py3-none-any.whl", hash = "sha256:20acf45e9b21851926835292e4045d290acade1edd2ff3de86d2f069687ba1fd", size = 61912, upload-time = "2026-06-24T15:19:35.434Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-sdk"
|
||||
version = "1.43.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "opentelemetry-api" },
|
||||
{ name = "opentelemetry-semantic-conventions" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/3e/eb/5041074274ac0956b03637cc039d434569112468e875eddfcc9a0674ce06/opentelemetry_sdk-1.43.0.tar.gz", hash = "sha256:d8187c81c162df9913e4003dd6485f7390d9a24fc17026ec7387b8b8218b08e9", size = 254744, upload-time = "2026-06-24T15:20:08.467Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/49/e3/b17be23af124201c9f52eececd4cc8ddfed1597d37b4ee771895d325805c/opentelemetry_sdk-1.43.0-py3-none-any.whl", hash = "sha256:d1323a547c1ce69d6a069a17a44b7da82bb8b332051ecb074041f87642c86823", size = 178852, upload-time = "2026-06-24T15:19:52.169Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opentelemetry-semantic-conventions"
|
||||
version = "0.64b0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "opentelemetry-api" },
|
||||
{ name = "typing-extensions" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/5a/30/5f26df29509eccd86b99b481ac9ffa39da49ba9577cc69071c552ae30447/opentelemetry_semantic_conventions-0.64b0.tar.gz", hash = "sha256:72f76fb2d1582d9d033dd1fcd84532e961e6ff3d90d24ba6fabc72975a83864c", size = 148340, upload-time = "2026-06-24T15:20:09.267Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/f2/ca/23ba87a221b574a7c5a99d48849d80bfe8b047624681357e2b002e566187/opentelemetry_semantic_conventions-0.64b0-py3-none-any.whl", hash = "sha256:ea77e85e354b8f604ddbe5f3d9135216f982fa4d77e5859ac30f6d8a50505aa6", size = 203713, upload-time = "2026-06-24T15:19:53.339Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "overrides"
|
||||
version = "7.7.0"
|
||||
|
||||
+2
-24
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.32.0-beta.2"
|
||||
version = "0.31.0-beta.6"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
@@ -44,12 +44,11 @@ lance-io = { workspace = true }
|
||||
lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lindera"] }
|
||||
lance-table = { workspace = true }
|
||||
lance-linalg = { workspace = true }
|
||||
lance-testing = { workspace = true }
|
||||
lance-encoding = { workspace = true }
|
||||
lance-arrow = { workspace = true }
|
||||
lance-namespace = { workspace = true }
|
||||
lance-namespace-impls = { workspace = true }
|
||||
metrics = { workspace = true, optional = true }
|
||||
metrics-util = { workspace = true, optional = true }
|
||||
moka = { workspace = true }
|
||||
pin-project = { workspace = true }
|
||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||
@@ -94,7 +93,6 @@ semver = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
anyhow = "1"
|
||||
lance-testing = { workspace = true }
|
||||
tempfile = "3.5.0"
|
||||
random_word = { version = "0.4.3", features = ["en"] }
|
||||
tokio = { version = "1.23", features = ["io-util", "macros", "net", "rt-multi-thread", "sync"] }
|
||||
@@ -110,8 +108,6 @@ http-body = "1" # Matching reqwest
|
||||
rstest = "0.23.0"
|
||||
test-log = "0.2"
|
||||
serial_test = "3"
|
||||
[target.'cfg(unix)'.dev-dependencies]
|
||||
pprof = { version = "0.14", features = ["flamegraph"] }
|
||||
|
||||
|
||||
[features]
|
||||
@@ -130,12 +126,6 @@ azure = [
|
||||
"lance-namespace-impls/dir-azure",
|
||||
"lance-namespace-impls/credential-vendor-azure",
|
||||
]
|
||||
cos = ["lance/tencent", "lance-io/tencent"]
|
||||
goosefs = [
|
||||
"lance/goosefs",
|
||||
"lance-io/goosefs",
|
||||
"lance-namespace-impls/dir-goosefs",
|
||||
]
|
||||
huggingface = [
|
||||
"lance/huggingface",
|
||||
"lance-io/huggingface",
|
||||
@@ -149,15 +139,6 @@ remote = [
|
||||
"lance-namespace-impls/rest",
|
||||
"lance-namespace-impls/rest-adapter",
|
||||
]
|
||||
# Publish LanceDB's internal metrics (currently object store request counts,
|
||||
# bytes, latency, errors, and throttles) through the `metrics` crate facade,
|
||||
# and re-export the `metrics` crate as `lancedb::metrics`. Install any
|
||||
# `metrics`-compatible recorder to collect them.
|
||||
metrics = ["dep:metrics", "lance/metrics", "lance-io/metrics"]
|
||||
# Additional adapter on top of `metrics` that installs a process-global recorder
|
||||
# and exposes a pull-based snapshot/catalog API (see `lancedb::metrics_otel`)
|
||||
# for bridging metrics into OpenTelemetry or other pull-based exporters.
|
||||
metrics-otel = ["metrics", "dep:metrics-util"]
|
||||
fp16kernels = ["lance-linalg/fp16kernels"]
|
||||
s3-test = []
|
||||
bedrock = ["dep:aws-sdk-bedrockruntime"]
|
||||
@@ -183,9 +164,6 @@ required-features = ["sentence-transformers"]
|
||||
name = "bedrock"
|
||||
required-features = ["bedrock"]
|
||||
|
||||
[[example]]
|
||||
name = "bench_streaming_dataloader"
|
||||
|
||||
[[example]]
|
||||
name = "simple"
|
||||
|
||||
|
||||
@@ -1,272 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! Benchmark + CPU profiler for the PermutationReader used by the elastic
|
||||
//! streaming dataloader.
|
||||
//!
|
||||
//! Normal sweep:
|
||||
//! cargo run --release --example bench_streaming_dataloader
|
||||
//!
|
||||
//! Flamegraph (self-contained, no perf/dtrace needed):
|
||||
//! BENCH_PROFILE=1 BENCH_CHUNK=64 cargo run --release \
|
||||
//! --example bench_streaming_dataloader
|
||||
//! # writes flamegraph.svg in the current directory
|
||||
//!
|
||||
//! Environment variables:
|
||||
//! BENCH_NUM_ROWS – total rows (default 49152 = 24 × 2048)
|
||||
//! BENCH_NUM_SPLITS – number of splits (default 24)
|
||||
//! BENCH_STEPS – round-robin cycles per chunk-size trial (default 200)
|
||||
//! BENCH_ROW_BYTES – bytes of payload per row (default 4096)
|
||||
//! BENCH_CHUNK – restrict sweep to this single chunk size
|
||||
//! BENCH_PROFILE – if set to "1", capture a pprof flamegraph SVG
|
||||
|
||||
use std::{sync::Arc, time::Instant};
|
||||
|
||||
use arrow_array::{Int32Array, LargeBinaryArray, RecordBatch};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use lancedb::{
|
||||
Result, Table,
|
||||
arrow::{SendableRecordBatchStream, SimpleRecordBatchStream},
|
||||
connect,
|
||||
dataloader::permutation::{
|
||||
builder::{PermutationBuilder, ShuffleStrategy},
|
||||
reader::PermutationReader,
|
||||
split::{SplitSizes, SplitStrategy},
|
||||
},
|
||||
query::Select,
|
||||
};
|
||||
|
||||
fn env_usize(key: &str, default: usize) -> usize {
|
||||
std::env::var(key)
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Table creation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn make_base_table(num_rows: usize, row_bytes: usize) -> Result<Table> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int32, false),
|
||||
Field::new("payload", DataType::LargeBinary, false),
|
||||
]));
|
||||
let payload = vec![0u8; row_bytes];
|
||||
let ids: Int32Array = (0..num_rows as i32).collect();
|
||||
let payloads: LargeBinaryArray = (0..num_rows).map(|_| Some(payload.as_slice())).collect();
|
||||
let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(ids), Arc::new(payloads)])?;
|
||||
let stream: SendableRecordBatchStream = Box::pin(SimpleRecordBatchStream::new(
|
||||
futures::stream::once(std::future::ready(Ok(batch))),
|
||||
schema,
|
||||
));
|
||||
let db = connect("memory:///").execute().await?;
|
||||
db.create_table("base", stream).execute().await
|
||||
}
|
||||
|
||||
async fn make_permutation_table(base: &Table, num_splits: usize) -> Result<Table> {
|
||||
PermutationBuilder::new(base.clone())
|
||||
.with_split_strategy(
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Fixed(num_splits as u64),
|
||||
clump_size: None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.with_shuffle_strategy(ShuffleStrategy::Random {
|
||||
seed: Some(42),
|
||||
clump_size: None,
|
||||
})
|
||||
.build()
|
||||
.await
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Round-robin hot loop (mirrors StreamingDataset.__iter__)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async fn run_hot_loop(
|
||||
readers: &[PermutationReader],
|
||||
chunk_size: usize,
|
||||
steps: usize,
|
||||
) -> Result<(usize, f64)> {
|
||||
let n = readers.len();
|
||||
let split_sizes: Vec<usize> = readers.iter().map(|r| r.count_rows() as usize).collect();
|
||||
|
||||
struct SplitBuf {
|
||||
batch: Option<RecordBatch>,
|
||||
row_in_batch: usize,
|
||||
consumed: usize,
|
||||
}
|
||||
let mut bufs: Vec<SplitBuf> = (0..n)
|
||||
.map(|_| SplitBuf {
|
||||
batch: None,
|
||||
row_in_batch: 0,
|
||||
consumed: 0,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Pre-fill
|
||||
for i in 0..n {
|
||||
let fetch = chunk_size.min(split_sizes[i]);
|
||||
if fetch > 0 {
|
||||
let offsets: Vec<u64> = (0..fetch as u64).collect();
|
||||
bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?);
|
||||
}
|
||||
}
|
||||
|
||||
let mut total_rows = 0usize;
|
||||
let t0 = Instant::now();
|
||||
|
||||
'outer: for _step in 0..steps {
|
||||
for i in 0..n {
|
||||
if bufs[i].consumed >= split_sizes[i] {
|
||||
break 'outer;
|
||||
}
|
||||
let need_refill = bufs[i]
|
||||
.batch
|
||||
.as_ref()
|
||||
.map(|b| bufs[i].row_in_batch >= b.num_rows())
|
||||
.unwrap_or(true);
|
||||
if need_refill {
|
||||
let start = bufs[i].consumed as u64;
|
||||
let remaining = (split_sizes[i] - bufs[i].consumed) as u64;
|
||||
let fetch = chunk_size.min(remaining as usize);
|
||||
let offsets: Vec<u64> = (start..start + fetch as u64).collect();
|
||||
bufs[i].batch = Some(readers[i].take_offsets(&offsets, Select::All).await?);
|
||||
bufs[i].row_in_batch = 0;
|
||||
}
|
||||
bufs[i].row_in_batch += 1;
|
||||
bufs[i].consumed += 1;
|
||||
total_rows += 1;
|
||||
}
|
||||
}
|
||||
|
||||
Ok((total_rows, t0.elapsed().as_secs_f64()))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Main
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let num_splits = env_usize("BENCH_NUM_SPLITS", 24);
|
||||
let num_rows = env_usize("BENCH_NUM_ROWS", num_splits * 2048);
|
||||
let steps = env_usize("BENCH_STEPS", 200);
|
||||
let row_bytes = env_usize("BENCH_ROW_BYTES", 4096);
|
||||
let single_chunk: Option<usize> = std::env::var("BENCH_CHUNK")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok());
|
||||
let do_profile = std::env::var("BENCH_PROFILE")
|
||||
.map(|v| v == "1")
|
||||
.unwrap_or(false);
|
||||
|
||||
assert_eq!(
|
||||
num_rows % num_splits,
|
||||
0,
|
||||
"NUM_ROWS must be divisible by NUM_SPLITS"
|
||||
);
|
||||
|
||||
println!("Benchmark config:");
|
||||
println!(
|
||||
" num_rows={} num_splits={} rows/split={} steps={} row_bytes={}",
|
||||
num_rows,
|
||||
num_splits,
|
||||
num_rows / num_splits,
|
||||
steps,
|
||||
row_bytes,
|
||||
);
|
||||
println!(
|
||||
" ~{:.1} MB total",
|
||||
(num_rows * row_bytes) as f64 / (1024.0 * 1024.0)
|
||||
);
|
||||
println!();
|
||||
|
||||
print!("Building base table... ");
|
||||
let _ = std::io::Write::flush(&mut std::io::stdout());
|
||||
let base = make_base_table(num_rows, row_bytes).await?;
|
||||
println!("done");
|
||||
|
||||
print!("Building permutation table... ");
|
||||
let _ = std::io::Write::flush(&mut std::io::stdout());
|
||||
let perm = make_permutation_table(&base, num_splits).await?;
|
||||
println!("done");
|
||||
|
||||
print!("Building {} PermutationReaders... ", num_splits);
|
||||
let _ = std::io::Write::flush(&mut std::io::stdout());
|
||||
let base_inner = base.base_table().clone();
|
||||
let perm_inner = perm.base_table().clone();
|
||||
let mut readers = Vec::with_capacity(num_splits);
|
||||
for split in 0..num_splits {
|
||||
readers.push(
|
||||
PermutationReader::try_from_tables(
|
||||
base_inner.clone(),
|
||||
perm_inner.clone(),
|
||||
split as u64,
|
||||
)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
println!("done ({} rows/split)", readers[0].count_rows());
|
||||
println!();
|
||||
|
||||
let chunk_sizes: Vec<usize> = if let Some(c) = single_chunk {
|
||||
vec![c]
|
||||
} else {
|
||||
vec![1, 4, 16, 64, 256, 1024, 4096, 16384]
|
||||
};
|
||||
|
||||
if do_profile {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
let chunk = chunk_sizes[0];
|
||||
println!("Profiling chunk={chunk} for {steps} steps...");
|
||||
// Warm-up outside the profiler window
|
||||
let _ = run_hot_loop(&readers, chunk, 1).await?;
|
||||
|
||||
let guard = pprof::ProfilerGuardBuilder::default()
|
||||
.frequency(1000)
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?;
|
||||
|
||||
if let Ok(report) = guard.report().build() {
|
||||
let svg_path = "flamegraph.svg";
|
||||
let file = std::fs::File::create(svg_path).unwrap();
|
||||
report.flamegraph(file).unwrap();
|
||||
println!("Flamegraph written to {svg_path}");
|
||||
}
|
||||
|
||||
let rows_per_sec = rows as f64 / elapsed;
|
||||
println!("chunk={chunk} {rows} rows {elapsed:.3}s {rows_per_sec:.0} rows/s");
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
println!("Flamegraph profiling (BENCH_PROFILE=1) is not supported on this platform.");
|
||||
println!("Run without BENCH_PROFILE to get throughput numbers.");
|
||||
}
|
||||
} else {
|
||||
println!(
|
||||
"{:>6} {:>7} {:>8} {:>11} {:>10}",
|
||||
"chunk", "rows", "elapsed", "rows/s", "ms/step"
|
||||
);
|
||||
println!("{}", "-".repeat(52));
|
||||
|
||||
for &chunk in &chunk_sizes {
|
||||
let _ = run_hot_loop(&readers, chunk, 1).await?;
|
||||
let (rows, elapsed) = run_hot_loop(&readers, chunk, steps).await?;
|
||||
let rows_per_sec = rows as f64 / elapsed;
|
||||
let ms_per_step = elapsed / steps as f64 * 1000.0;
|
||||
println!(
|
||||
"{:>6} {:>7} {:>7.3}s {:>11.0} {:>9.1}ms",
|
||||
chunk, rows, elapsed, rows_per_sec, ms_per_step,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
println!("\nDone.");
|
||||
Ok(())
|
||||
}
|
||||
@@ -342,19 +342,6 @@ impl ListingDatabase {
|
||||
))
|
||||
}
|
||||
|
||||
fn storage_base_uri(uri: &str) -> String {
|
||||
let Ok(mut url) = url::Url::parse(uri) else {
|
||||
return uri.to_string();
|
||||
};
|
||||
url.set_query(None);
|
||||
let Some((storage_scheme, _commit_scheme)) = url.scheme().split_once('+') else {
|
||||
return url.to_string();
|
||||
};
|
||||
let storage_scheme = storage_scheme.to_string();
|
||||
let _ = url.set_scheme(&storage_scheme);
|
||||
url.to_string()
|
||||
}
|
||||
|
||||
async fn prepare_namespace_root(
|
||||
uri: &str,
|
||||
storage_options: &HashMap<String, String>,
|
||||
@@ -533,8 +520,6 @@ impl ListingDatabase {
|
||||
// will add a trailing '?' to the url
|
||||
url.set_query(None);
|
||||
|
||||
let storage_base_uri = Self::storage_base_uri(url.as_str());
|
||||
|
||||
let table_base_uri = if let Some(store) = engine {
|
||||
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
|
||||
WARN_ONCE.call_once(|| {
|
||||
@@ -547,6 +532,8 @@ impl ListingDatabase {
|
||||
url.to_string()
|
||||
};
|
||||
|
||||
let plain_uri = url.to_string();
|
||||
|
||||
let session = request
|
||||
.session
|
||||
.clone()
|
||||
@@ -563,13 +550,13 @@ impl ListingDatabase {
|
||||
};
|
||||
let (object_store, base_path) = ObjectStore::from_uri_and_params(
|
||||
session.store_registry(),
|
||||
&storage_base_uri,
|
||||
&plain_uri,
|
||||
&os_params,
|
||||
)
|
||||
.await?;
|
||||
if object_store.is_local() {
|
||||
Self::try_create_dir(&storage_base_uri).context(CreateDirSnafu {
|
||||
path: storage_base_uri.clone(),
|
||||
Self::try_create_dir(&plain_uri).context(CreateDirSnafu {
|
||||
path: plain_uri.clone(),
|
||||
})?;
|
||||
}
|
||||
|
||||
@@ -583,7 +570,7 @@ impl ListingDatabase {
|
||||
};
|
||||
|
||||
let namespace_database = Self::connect_namespace_database(
|
||||
&storage_base_uri,
|
||||
&table_base_uri,
|
||||
options.storage_options.clone(),
|
||||
request.namespace_client_properties.clone(),
|
||||
request.read_consistency_interval,
|
||||
@@ -1322,60 +1309,6 @@ mod tests {
|
||||
(tempdir, db)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_listing_database_root_ops_do_not_create_manifest() {
|
||||
let tempdir = tempdir().unwrap();
|
||||
let uri = tempdir.path().to_str().unwrap();
|
||||
|
||||
let request = ConnectRequest {
|
||||
uri: uri.to_string(),
|
||||
#[cfg(feature = "remote")]
|
||||
client_config: Default::default(),
|
||||
options: Default::default(),
|
||||
namespace_client_properties: Default::default(),
|
||||
manifest_enabled: false,
|
||||
read_consistency_interval: None,
|
||||
session: None,
|
||||
};
|
||||
|
||||
let db = ListingDatabase::connect_with_options(&request)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!tempdir.path().join("__manifest").exists());
|
||||
|
||||
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
|
||||
db.create_table(CreateTableRequest {
|
||||
name: "root_table".to_string(),
|
||||
namespace_path: vec![],
|
||||
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
|
||||
mode: CreateTableMode::Create,
|
||||
write_options: Default::default(),
|
||||
location: None,
|
||||
namespace_client: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
db.open_table(OpenTableRequest {
|
||||
name: "root_table".to_string(),
|
||||
namespace_path: vec![],
|
||||
index_cache_size: None,
|
||||
lance_read_params: None,
|
||||
location: None,
|
||||
namespace_client: None,
|
||||
managed_versioning: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
#[allow(deprecated)]
|
||||
let table_names = db.table_names(TableNamesRequest::default()).await.unwrap();
|
||||
|
||||
assert_eq!(table_names, vec!["root_table".to_string()]);
|
||||
assert!(!tempdir.path().join("__manifest").exists());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clone_table_basic() {
|
||||
let (_tempdir, db) = setup_database().await;
|
||||
@@ -2350,24 +2283,6 @@ mod tests {
|
||||
assert_eq!(captured.as_deref(), Some("foo=bar"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_storage_base_uri_strips_commit_engine_scheme() {
|
||||
assert_eq!(
|
||||
ListingDatabase::storage_base_uri("s3+ddb://bucket/prefix?ddbTableName=commit_table"),
|
||||
"s3://bucket/prefix"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
ListingDatabase::storage_base_uri("s3://bucket/prefix?foo=bar"),
|
||||
"s3://bucket/prefix"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
ListingDatabase::storage_base_uri("/tmp/lancedb"),
|
||||
"/tmp/lancedb"
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: connecting via a URL-style URI (which goes through
|
||||
/// `url::Url::parse` and the `query_pairs_mut()` path) must not
|
||||
/// append a trailing `?` to per-table URIs when the input URI has
|
||||
|
||||
@@ -391,7 +391,6 @@ mod tests {
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Percentages(vec![0.05, 0.30]),
|
||||
clump_size: None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user