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
|
name: lancedb-branch-ops
|
||||||
description: >-
|
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.
|
||||||
Manage LanceDB table branches through the REST API: list, create, and delete
|
|
||||||
branches; target schema reads, field-metadata updates, and index creation to a
|
|
||||||
named branch; and verify that branch changes remain isolated from main. Use
|
|
||||||
when a task involves branch lifecycle, an experimental or isolated table
|
|
||||||
version, directing an operation to a non-main branch, or confirming that a
|
|
||||||
mutation did not affect main. This skill also explains that LanceDB has no
|
|
||||||
checkout operation; each request selects its target branch in the request
|
|
||||||
body.
|
|
||||||
---
|
---
|
||||||
|
|
||||||
## Goal
|
## 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]
|
[tool.bumpversion]
|
||||||
current_version = "0.32.0-beta.1"
|
current_version = "0.31.0-beta.6"
|
||||||
parse = """(?x)
|
parse = """(?x)
|
||||||
(?P<major>0|[1-9]\\d*)\\.
|
(?P<major>0|[1-9]\\d*)\\.
|
||||||
(?P<minor>0|[1-9]\\d*)\\.
|
(?P<minor>0|[1-9]\\d*)\\.
|
||||||
|
|||||||
@@ -34,16 +34,15 @@ runs:
|
|||||||
maturin-version: "1.12.4"
|
maturin-version: "1.12.4"
|
||||||
command: build
|
command: build
|
||||||
working-directory: python
|
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
|
target: x86_64-unknown-linux-gnu
|
||||||
manylinux: ${{ inputs.manylinux }}
|
manylinux: ${{ inputs.manylinux }}
|
||||||
args: ${{ inputs.args }}
|
args: ${{ inputs.args }}
|
||||||
before-script-linux: |
|
before-script-linux: |
|
||||||
set -e
|
set -e
|
||||||
curl -fsSL https://github.com/protocolbuffers/protobuf/releases/download/v24.4/protoc-24.4-linux-x86_64.zip -o /tmp/protoc.zip
|
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
|
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||||
rm /tmp/protoc.zip
|
&& rm /tmp/protoc.zip
|
||||||
/usr/local/bin/protoc --version
|
|
||||||
- name: Build Arm Manylinux Wheel
|
- name: Build Arm Manylinux Wheel
|
||||||
if: ${{ inputs.arm-build == 'true' }}
|
if: ${{ inputs.arm-build == 'true' }}
|
||||||
uses: PyO3/maturin-action@v1
|
uses: PyO3/maturin-action@v1
|
||||||
@@ -51,14 +50,13 @@ runs:
|
|||||||
maturin-version: "1.12.4"
|
maturin-version: "1.12.4"
|
||||||
command: build
|
command: build
|
||||||
working-directory: python
|
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
|
target: aarch64-unknown-linux-gnu
|
||||||
manylinux: ${{ inputs.manylinux }}
|
manylinux: ${{ inputs.manylinux }}
|
||||||
args: ${{ inputs.args }}
|
args: ${{ inputs.args }}
|
||||||
before-script-linux: |
|
before-script-linux: |
|
||||||
set -e
|
set -e
|
||||||
yum install -y clang
|
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
|
&& 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
|
&& unzip /tmp/protoc.zip -d /usr/local \
|
||||||
rm /tmp/protoc.zip
|
&& rm /tmp/protoc.zip
|
||||||
/usr/local/bin/protoc --version
|
|
||||||
|
|||||||
@@ -103,7 +103,7 @@ jobs:
|
|||||||
features: fp16kernels
|
features: fp16kernels
|
||||||
pre_build: brew install protobuf
|
pre_build: brew install protobuf
|
||||||
- target: x86_64-pc-windows-msvc
|
- target: x86_64-pc-windows-msvc
|
||||||
host: windows-2025-8x-x64
|
host: windows-latest
|
||||||
features: ","
|
features: ","
|
||||||
pre_build: |-
|
pre_build: |-
|
||||||
choco install --no-progress protoc ninja nasm
|
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
|
# There is an issue where choco doesn't add nasm to the path
|
||||||
export PATH="$PATH:/c/Program Files/NASM"
|
export PATH="$PATH:/c/Program Files/NASM"
|
||||||
nasm -v
|
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
|
- target: aarch64-pc-windows-msvc
|
||||||
host: windows-2025-8x-x64
|
host: windows-latest
|
||||||
features: ","
|
features: ","
|
||||||
pre_build: |-
|
pre_build: |-
|
||||||
choco install --no-progress protoc
|
choco install --no-progress protoc
|
||||||
rustup target add aarch64-pc-windows-msvc
|
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
|
- target: x86_64-unknown-linux-gnu
|
||||||
host: ubuntu-latest
|
host: ubuntu-latest
|
||||||
features: fp16kernels
|
features: fp16kernels
|
||||||
|
|||||||
@@ -125,26 +125,10 @@ jobs:
|
|||||||
- uses: rui314/setup-mold@v1
|
- uses: rui314/setup-mold@v1
|
||||||
- name: Make Swap
|
- name: Make Swap
|
||||||
run: |
|
run: |
|
||||||
swapfile=/swapfile
|
sudo fallocate -l 16G /swapfile
|
||||||
min_swap_bytes=$((15 * 1024 * 1024 * 1024))
|
sudo chmod 600 /swapfile
|
||||||
active_swap_bytes="$(sudo swapon --show=NAME,SIZE --bytes --noheadings | awk '$1 == "/swapfile" { print $2 }')"
|
sudo mkswap /swapfile
|
||||||
if [ -n "$active_swap_bytes" ]; then
|
sudo swapon /swapfile
|
||||||
if [ "$active_swap_bytes" -ge "$min_swap_bytes" ]; then
|
|
||||||
echo "/swapfile is already active with enough space; skipping swap creation"
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
echo "/swapfile is already active but smaller than 16G; using /mnt/lancedb-swapfile"
|
|
||||||
swapfile=/mnt/lancedb-swapfile
|
|
||||||
fi
|
|
||||||
if sudo swapon --show=NAME --noheadings | grep -Fxq "$swapfile"; then
|
|
||||||
echo "$swapfile is already active; skipping swap creation"
|
|
||||||
exit 0
|
|
||||||
fi
|
|
||||||
sudo rm -f "$swapfile"
|
|
||||||
sudo fallocate -l 16G "$swapfile"
|
|
||||||
sudo chmod 600 "$swapfile"
|
|
||||||
sudo mkswap "$swapfile"
|
|
||||||
sudo swapon "$swapfile"
|
|
||||||
- name: Build
|
- name: Build
|
||||||
run: cargo build --profile ci --all-features --tests --locked --examples
|
run: cargo build --profile ci --all-features --tests --locked --examples
|
||||||
- name: Run feature tests
|
- name: Run feature tests
|
||||||
|
|||||||
Generated
+163
-393
File diff suppressed because it is too large
Load Diff
+14
-16
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
|||||||
rust-version = "1.91.0"
|
rust-version = "1.91.0"
|
||||||
|
|
||||||
[workspace.dependencies]
|
[workspace.dependencies]
|
||||||
lance = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance = { "version" = "=8.0.0", default-features = false }
|
||||||
lance-core = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-core = "=8.0.0"
|
||||||
lance-datagen = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-datagen = "=8.0.0"
|
||||||
lance-file = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-file = "=8.0.0"
|
||||||
lance-io = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-io = { "version" = "=8.0.0", default-features = false }
|
||||||
lance-index = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-index = "=8.0.0"
|
||||||
lance-linalg = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-linalg = "=8.0.0"
|
||||||
lance-namespace = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-namespace = "=8.0.0"
|
||||||
lance-namespace-impls = { "version" = "=9.0.0-beta.23", default-features = false, "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-namespace-impls = { "version" = "=8.0.0", default-features = false }
|
||||||
lance-table = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-table = "=8.0.0"
|
||||||
lance-testing = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-testing = "=8.0.0"
|
||||||
lance-datafusion = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-datafusion = "=8.0.0"
|
||||||
lance-encoding = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-encoding = "=8.0.0"
|
||||||
lance-arrow = { "version" = "=9.0.0-beta.23", "tag" = "v9.0.0-beta.23", "git" = "https://github.com/lance-format/lance.git" }
|
lance-arrow = "=8.0.0"
|
||||||
ahash = "0.8"
|
ahash = "0.8"
|
||||||
# Note that this one does not include pyarrow
|
# Note that this one does not include pyarrow
|
||||||
arrow = { version = "58.0.0", optional = false }
|
arrow = { version = "58.0.0", optional = false }
|
||||||
@@ -54,8 +54,6 @@ half = { "version" = "2.7.1", default-features = false, features = [
|
|||||||
] }
|
] }
|
||||||
futures = "0"
|
futures = "0"
|
||||||
log = "0.4"
|
log = "0.4"
|
||||||
metrics = "0.24"
|
|
||||||
metrics-util = "0.19"
|
|
||||||
moka = { version = "0.12", features = ["future"] }
|
moka = { version = "0.12", features = ["future"] }
|
||||||
object_store = "0.13.2"
|
object_store = "0.13.2"
|
||||||
pin-project = "1.0.7"
|
pin-project = "1.0.7"
|
||||||
|
|||||||
@@ -51,6 +51,18 @@ ignore = [
|
|||||||
# https://rustsec.org/advisories/RUSTSEC-2024-0436
|
# https://rustsec.org/advisories/RUSTSEC-2024-0436
|
||||||
{ id = "RUSTSEC-2024-0436", reason = "transitive via datafusion; awaiting ecosystem migration" },
|
{ 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.
|
# tantivy: segfault on malformed input due to missing bounds check.
|
||||||
# Pulled in via lance for full-text search. We only feed tantivy
|
# Pulled in via lance for full-text search. We only feed tantivy
|
||||||
# documents we construct ourselves, not attacker-controlled bytes.
|
# documents we construct ourselves, not attacker-controlled bytes.
|
||||||
@@ -68,6 +80,18 @@ ignore = [
|
|||||||
# https://rustsec.org/advisories/RUSTSEC-2025-0119
|
# https://rustsec.org/advisories/RUSTSEC-2025-0119
|
||||||
{ id = "RUSTSEC-2025-0119", reason = "transitive via hf-hub/indicatif; cosmetic formatting crate" },
|
{ 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
|
# 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
|
# 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.
|
# 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
|
# 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" },
|
{ 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.
|
# 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-0176
|
||||||
# https://rustsec.org/advisories/RUSTSEC-2026-0177
|
# https://rustsec.org/advisories/RUSTSEC-2026-0177
|
||||||
{ id = "RUSTSEC-2026-0176", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
{ 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" },
|
{ 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>
|
<dependency>
|
||||||
<groupId>com.lancedb</groupId>
|
<groupId>com.lancedb</groupId>
|
||||||
<artifactId>lancedb-core</artifactId>
|
<artifactId>lancedb-core</artifactId>
|
||||||
<version>0.32.0-beta.1</version>
|
<version>0.31.0-beta.6</version>
|
||||||
</dependency>
|
</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()
|
### indexStats()
|
||||||
|
|
||||||
```ts
|
```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()
|
### unsetLsmWriteSpec()
|
||||||
|
|
||||||
```ts
|
```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)
|
- [FragmentStatistics](interfaces/FragmentStatistics.md)
|
||||||
- [FragmentSummaryStats](interfaces/FragmentSummaryStats.md)
|
- [FragmentSummaryStats](interfaces/FragmentSummaryStats.md)
|
||||||
- [FtsOptions](interfaces/FtsOptions.md)
|
- [FtsOptions](interfaces/FtsOptions.md)
|
||||||
- [FtsToken](interfaces/FtsToken.md)
|
|
||||||
- [FullTextQuery](interfaces/FullTextQuery.md)
|
- [FullTextQuery](interfaces/FullTextQuery.md)
|
||||||
- [FullTextSearchOptions](interfaces/FullTextSearchOptions.md)
|
- [FullTextSearchOptions](interfaces/FullTextSearchOptions.md)
|
||||||
- [HnswPqOptions](interfaces/HnswPqOptions.md)
|
- [HnswPqOptions](interfaces/HnswPqOptions.md)
|
||||||
@@ -108,7 +107,6 @@
|
|||||||
- [TimeoutConfig](interfaces/TimeoutConfig.md)
|
- [TimeoutConfig](interfaces/TimeoutConfig.md)
|
||||||
- [TlsConfig](interfaces/TlsConfig.md)
|
- [TlsConfig](interfaces/TlsConfig.md)
|
||||||
- [TokenResponse](interfaces/TokenResponse.md)
|
- [TokenResponse](interfaces/TokenResponse.md)
|
||||||
- [TokenizeOptions](interfaces/TokenizeOptions.md)
|
|
||||||
- [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md)
|
- [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md)
|
||||||
- [UpdateOptions](interfaces/UpdateOptions.md)
|
- [UpdateOptions](interfaces/UpdateOptions.md)
|
||||||
- [UpdateResult](interfaces/UpdateResult.md)
|
- [UpdateResult](interfaces/UpdateResult.md)
|
||||||
@@ -118,7 +116,6 @@
|
|||||||
|
|
||||||
## Type Aliases
|
## Type Aliases
|
||||||
|
|
||||||
- [BaseTokenizer](type-aliases/BaseTokenizer.md)
|
|
||||||
- [Data](type-aliases/Data.md)
|
- [Data](type-aliases/Data.md)
|
||||||
- [DataLike](type-aliases/DataLike.md)
|
- [DataLike](type-aliases/DataLike.md)
|
||||||
- [FieldLike](type-aliases/FieldLike.md)
|
- [FieldLike](type-aliases/FieldLike.md)
|
||||||
@@ -128,15 +125,12 @@
|
|||||||
- [RecordBatchLike](type-aliases/RecordBatchLike.md)
|
- [RecordBatchLike](type-aliases/RecordBatchLike.md)
|
||||||
- [SchemaLike](type-aliases/SchemaLike.md)
|
- [SchemaLike](type-aliases/SchemaLike.md)
|
||||||
- [TableLike](type-aliases/TableLike.md)
|
- [TableLike](type-aliases/TableLike.md)
|
||||||
- [TokenizeTableOptions](type-aliases/TokenizeTableOptions.md)
|
|
||||||
|
|
||||||
## Functions
|
## Functions
|
||||||
|
|
||||||
- [RecordBatchIterator](functions/RecordBatchIterator.md)
|
- [RecordBatchIterator](functions/RecordBatchIterator.md)
|
||||||
- [connect](functions/connect.md)
|
- [connect](functions/connect.md)
|
||||||
- [connectNamespace](functions/connectNamespace.md)
|
- [connectNamespace](functions/connectNamespace.md)
|
||||||
- [instrumentLanceDbMetrics](functions/instrumentLanceDbMetrics.md)
|
|
||||||
- [makeArrowTable](functions/makeArrowTable.md)
|
- [makeArrowTable](functions/makeArrowTable.md)
|
||||||
- [packBits](functions/packBits.md)
|
- [packBits](functions/packBits.md)
|
||||||
- [permutationBuilder](functions/permutationBuilder.md)
|
- [permutationBuilder](functions/permutationBuilder.md)
|
||||||
- [tokenize](functions/tokenize.md)
|
|
||||||
|
|||||||
@@ -23,7 +23,7 @@ whether to remove punctuation
|
|||||||
### baseTokenizer?
|
### baseTokenizer?
|
||||||
|
|
||||||
```ts
|
```ts
|
||||||
optional baseTokenizer: BaseTokenizer;
|
optional baseTokenizer: "raw" | "simple" | "whitespace" | "ngram";
|
||||||
```
|
```
|
||||||
|
|
||||||
The tokenizer to use when building the index.
|
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.
|
"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?
|
### 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
|
## Properties
|
||||||
|
|
||||||
### clumpSize?
|
|
||||||
|
|
||||||
```ts
|
|
||||||
optional clumpSize: number;
|
|
||||||
```
|
|
||||||
|
|
||||||
***
|
|
||||||
|
|
||||||
### counts?
|
### counts?
|
||||||
|
|
||||||
```ts
|
```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>
|
<parent>
|
||||||
<groupId>com.lancedb</groupId>
|
<groupId>com.lancedb</groupId>
|
||||||
<artifactId>lancedb-parent</artifactId>
|
<artifactId>lancedb-parent</artifactId>
|
||||||
<version>0.32.0-beta.1</version>
|
<version>0.31.0-beta.6</version>
|
||||||
<relativePath>../pom.xml</relativePath>
|
<relativePath>../pom.xml</relativePath>
|
||||||
</parent>
|
</parent>
|
||||||
|
|
||||||
|
|||||||
+2
-2
@@ -6,7 +6,7 @@
|
|||||||
|
|
||||||
<groupId>com.lancedb</groupId>
|
<groupId>com.lancedb</groupId>
|
||||||
<artifactId>lancedb-parent</artifactId>
|
<artifactId>lancedb-parent</artifactId>
|
||||||
<version>0.32.0-beta.1</version>
|
<version>0.31.0-beta.6</version>
|
||||||
<packaging>pom</packaging>
|
<packaging>pom</packaging>
|
||||||
<name>${project.artifactId}</name>
|
<name>${project.artifactId}</name>
|
||||||
<description>LanceDB Java SDK Parent POM</description>
|
<description>LanceDB Java SDK Parent POM</description>
|
||||||
@@ -28,7 +28,7 @@
|
|||||||
<properties>
|
<properties>
|
||||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||||
<arrow.version>15.0.0</arrow.version>
|
<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.skip>false</spotless.skip>
|
||||||
<spotless.version>2.30.0</spotless.version>
|
<spotless.version>2.30.0</spotless.version>
|
||||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||||
|
|||||||
+2
-2
@@ -1,7 +1,7 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "lancedb-nodejs"
|
name = "lancedb-nodejs"
|
||||||
edition.workspace = true
|
edition.workspace = true
|
||||||
version = "0.32.0-beta.1"
|
version = "0.31.0-beta.6"
|
||||||
publish = false
|
publish = false
|
||||||
license.workspace = true
|
license.workspace = true
|
||||||
description.workspace = true
|
description.workspace = true
|
||||||
@@ -44,6 +44,6 @@ aws-lc-rs = "=1.16.3"
|
|||||||
napi-build = "2.3.1"
|
napi-build = "2.3.1"
|
||||||
|
|
||||||
[features]
|
[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"]
|
fp16kernels = ["lancedb/fp16kernels"]
|
||||||
remote = ["lancedb/remote"]
|
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,
|
PhraseQuery,
|
||||||
Table,
|
Table,
|
||||||
connect,
|
connect,
|
||||||
tokenize,
|
|
||||||
} from "../lancedb";
|
} from "../lancedb";
|
||||||
import {
|
import {
|
||||||
Table as ArrowTable,
|
Table as ArrowTable,
|
||||||
@@ -2308,75 +2307,6 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
|||||||
expect(results2[0].text).toBe(data[1].text);
|
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 () => {
|
test("full text search fast search", async () => {
|
||||||
const db = await connect(tmpDir.name);
|
const db = await connect(tmpDir.name);
|
||||||
const data = [{ text: "hello world", vector: [0.1, 0.2, 0.3], id: 1 }];
|
const data = [{ text: "hello world", vector: [0.1, 0.2, 0.3], id: 1 }];
|
||||||
@@ -3062,56 +2992,6 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => {
|
|||||||
}),
|
}),
|
||||||
).rejects.toThrow();
|
).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", () => {
|
describe("LSM merge insert", () => {
|
||||||
|
|||||||
@@ -13,21 +13,13 @@ import {
|
|||||||
Connection as LanceDbConnection,
|
Connection as LanceDbConnection,
|
||||||
JsHeaderProvider as NativeJsHeaderProvider,
|
JsHeaderProvider as NativeJsHeaderProvider,
|
||||||
Session,
|
Session,
|
||||||
tokenize as nativeTokenize,
|
|
||||||
} from "./native.js";
|
} from "./native.js";
|
||||||
|
|
||||||
import { HeaderProvider } from "./header";
|
import { HeaderProvider } from "./header";
|
||||||
import type { BaseTokenizer } from "./indices";
|
|
||||||
import type { FtsToken } from "./table";
|
|
||||||
|
|
||||||
// Re-export native header provider for use with connectWithHeaderProvider
|
// Re-export native header provider for use with connectWithHeaderProvider
|
||||||
export { JsHeaderProvider as NativeJsHeaderProvider } from "./native.js";
|
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 {
|
export {
|
||||||
AddColumnsSql,
|
AddColumnsSql,
|
||||||
ConnectionOptions,
|
ConnectionOptions,
|
||||||
@@ -117,7 +109,6 @@ export {
|
|||||||
HnswPqOptions,
|
HnswPqOptions,
|
||||||
HnswSqOptions,
|
HnswSqOptions,
|
||||||
FtsOptions,
|
FtsOptions,
|
||||||
BaseTokenizer,
|
|
||||||
} from "./indices";
|
} from "./indices";
|
||||||
|
|
||||||
export {
|
export {
|
||||||
@@ -128,8 +119,6 @@ export {
|
|||||||
OptimizeOptions,
|
OptimizeOptions,
|
||||||
Version,
|
Version,
|
||||||
WriteProgress,
|
WriteProgress,
|
||||||
FtsToken,
|
|
||||||
TokenizeTableOptions,
|
|
||||||
LsmWriteSpec,
|
LsmWriteSpec,
|
||||||
ColumnAlteration,
|
ColumnAlteration,
|
||||||
FieldMetadataUpdate,
|
FieldMetadataUpdate,
|
||||||
@@ -161,68 +150,6 @@ export {
|
|||||||
} from "./arrow";
|
} from "./arrow";
|
||||||
export { IntoSql, packBits } from "./util";
|
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.
|
* Connect to a LanceDB instance at the given URI.
|
||||||
*
|
*
|
||||||
|
|||||||
@@ -486,16 +486,6 @@ export interface IvfFlatOptions {
|
|||||||
sampleRate?: number;
|
sampleRate?: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
export type BaseTokenizer =
|
|
||||||
| "simple"
|
|
||||||
| "whitespace"
|
|
||||||
| "raw"
|
|
||||||
| "ngram"
|
|
||||||
| "icu"
|
|
||||||
| "icu/split"
|
|
||||||
| `jieba/${string}`
|
|
||||||
| `lindera/${string}`;
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Options to create a full text search index
|
* 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.
|
* "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.
|
* "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
|
* 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>;
|
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
|
* Specification selecting Lance's MemWAL LSM-style write path for
|
||||||
* `mergeInsert`.
|
* `mergeInsert`.
|
||||||
@@ -605,17 +585,6 @@ export abstract class Table {
|
|||||||
* @returns {Promise<void>}
|
* @returns {Promise<void>}
|
||||||
*/
|
*/
|
||||||
abstract unsetLsmWriteSpec(): 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.
|
* 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>;
|
abstract optimize(options?: Partial<OptimizeOptions>): Promise<OptimizeStats>;
|
||||||
/** List all indices that have been created with {@link Table.createIndex} */
|
/** List all indices that have been created with {@link Table.createIndex} */
|
||||||
abstract listIndices(): Promise<IndexConfig[]>;
|
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 */
|
/** Return the table as an arrow table */
|
||||||
abstract toArrow(): Promise<ArrowTable>;
|
abstract toArrow(): Promise<ArrowTable>;
|
||||||
|
|
||||||
@@ -1135,15 +1091,6 @@ export class LocalTable extends Table {
|
|||||||
return await this.inner.unsetLsmWriteSpec();
|
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> {
|
async closeLsmWriters(): Promise<void> {
|
||||||
return await this.inner.closeLsmWriters();
|
return await this.inner.closeLsmWriters();
|
||||||
}
|
}
|
||||||
@@ -1206,17 +1153,6 @@ export class LocalTable extends Table {
|
|||||||
return await this.inner.listIndices();
|
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> {
|
async toArrow(): Promise<ArrowTable> {
|
||||||
return await this.query().toArrow();
|
return await this.query().toArrow();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-darwin-arm64",
|
"name": "@lancedb/lancedb-darwin-arm64",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["darwin"],
|
"os": ["darwin"],
|
||||||
"cpu": ["arm64"],
|
"cpu": ["arm64"],
|
||||||
"main": "lancedb.darwin-arm64.node",
|
"main": "lancedb.darwin-arm64.node",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["linux"],
|
"os": ["linux"],
|
||||||
"cpu": ["arm64"],
|
"cpu": ["arm64"],
|
||||||
"main": "lancedb.linux-arm64-gnu.node",
|
"main": "lancedb.linux-arm64-gnu.node",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["linux"],
|
"os": ["linux"],
|
||||||
"cpu": ["arm64"],
|
"cpu": ["arm64"],
|
||||||
"main": "lancedb.linux-arm64-musl.node",
|
"main": "lancedb.linux-arm64-musl.node",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["linux"],
|
"os": ["linux"],
|
||||||
"cpu": ["x64"],
|
"cpu": ["x64"],
|
||||||
"main": "lancedb.linux-x64-gnu.node",
|
"main": "lancedb.linux-x64-gnu.node",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["linux"],
|
"os": ["linux"],
|
||||||
"cpu": ["x64"],
|
"cpu": ["x64"],
|
||||||
"main": "lancedb.linux-x64-musl.node",
|
"main": "lancedb.linux-x64-musl.node",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": [
|
"os": [
|
||||||
"win32"
|
"win32"
|
||||||
],
|
],
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"os": ["win32"],
|
"os": ["win32"],
|
||||||
"cpu": ["x64"],
|
"cpu": ["x64"],
|
||||||
"main": "lancedb.win32-x64-msvc.node",
|
"main": "lancedb.win32-x64-msvc.node",
|
||||||
|
|||||||
Generated
+2
-73
@@ -1,12 +1,12 @@
|
|||||||
{
|
{
|
||||||
"name": "@lancedb/lancedb",
|
"name": "@lancedb/lancedb",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"lockfileVersion": 3,
|
"lockfileVersion": 3,
|
||||||
"requires": true,
|
"requires": true,
|
||||||
"packages": {
|
"packages": {
|
||||||
"": {
|
"": {
|
||||||
"name": "@lancedb/lancedb",
|
"name": "@lancedb/lancedb",
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"cpu": [
|
"cpu": [
|
||||||
"x64",
|
"x64",
|
||||||
"arm64"
|
"arm64"
|
||||||
@@ -18,7 +18,6 @@
|
|||||||
"win32"
|
"win32"
|
||||||
],
|
],
|
||||||
"dependencies": {
|
"dependencies": {
|
||||||
"@opentelemetry/api": "^1.9.0",
|
|
||||||
"reflect-metadata": "^0.2.2"
|
"reflect-metadata": "^0.2.2"
|
||||||
},
|
},
|
||||||
"devDependencies": {
|
"devDependencies": {
|
||||||
@@ -28,7 +27,6 @@
|
|||||||
"@biomejs/biome": "^1.7.3",
|
"@biomejs/biome": "^1.7.3",
|
||||||
"@jest/globals": "^29.7.0",
|
"@jest/globals": "^29.7.0",
|
||||||
"@napi-rs/cli": "3.7.0",
|
"@napi-rs/cli": "3.7.0",
|
||||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
|
||||||
"@types/axios": "^0.14.0",
|
"@types/axios": "^0.14.0",
|
||||||
"@types/jest": "^29.1.2",
|
"@types/jest": "^29.1.2",
|
||||||
"@types/node": "22.7.4",
|
"@types/node": "22.7.4",
|
||||||
@@ -4150,75 +4148,6 @@
|
|||||||
"@octokit/openapi-types": "^27.0.0"
|
"@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": {
|
"node_modules/@protobufjs/aspromise": {
|
||||||
"version": "1.1.2",
|
"version": "1.1.2",
|
||||||
"resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz",
|
"resolved": "https://registry.npmjs.org/@protobufjs/aspromise/-/aspromise-1.1.2.tgz",
|
||||||
|
|||||||
+1
-3
@@ -11,7 +11,7 @@
|
|||||||
"ann"
|
"ann"
|
||||||
],
|
],
|
||||||
"private": false,
|
"private": false,
|
||||||
"version": "0.32.0-beta.1",
|
"version": "0.31.0-beta.6",
|
||||||
"main": "dist/index.js",
|
"main": "dist/index.js",
|
||||||
"exports": {
|
"exports": {
|
||||||
".": "./dist/index.js",
|
".": "./dist/index.js",
|
||||||
@@ -44,7 +44,6 @@
|
|||||||
"@biomejs/biome": "^1.7.3",
|
"@biomejs/biome": "^1.7.3",
|
||||||
"@jest/globals": "^29.7.0",
|
"@jest/globals": "^29.7.0",
|
||||||
"@napi-rs/cli": "3.7.0",
|
"@napi-rs/cli": "3.7.0",
|
||||||
"@opentelemetry/sdk-metrics": "^1.30.0",
|
|
||||||
"@types/axios": "^0.14.0",
|
"@types/axios": "^0.14.0",
|
||||||
"@types/jest": "^29.1.2",
|
"@types/jest": "^29.1.2",
|
||||||
"@types/node": "22.7.4",
|
"@types/node": "22.7.4",
|
||||||
@@ -93,7 +92,6 @@
|
|||||||
"version": "napi version"
|
"version": "napi version"
|
||||||
},
|
},
|
||||||
"dependencies": {
|
"dependencies": {
|
||||||
"@opentelemetry/api": "^1.9.0",
|
|
||||||
"reflect-metadata": "^0.2.2"
|
"reflect-metadata": "^0.2.2"
|
||||||
},
|
},
|
||||||
"optionalDependencies": {
|
"optionalDependencies": {
|
||||||
|
|||||||
Generated
-53
@@ -8,9 +8,6 @@ importers:
|
|||||||
|
|
||||||
.:
|
.:
|
||||||
dependencies:
|
dependencies:
|
||||||
'@opentelemetry/api':
|
|
||||||
specifier: ^1.9.0
|
|
||||||
version: 1.9.1
|
|
||||||
apache-arrow:
|
apache-arrow:
|
||||||
specifier: '>=15.0.0 <=18.1.0'
|
specifier: '>=15.0.0 <=18.1.0'
|
||||||
version: 18.1.0
|
version: 18.1.0
|
||||||
@@ -36,9 +33,6 @@ importers:
|
|||||||
'@napi-rs/cli':
|
'@napi-rs/cli':
|
||||||
specifier: 3.7.0
|
specifier: 3.7.0
|
||||||
version: 3.7.0(@emnapi/core@1.10.0)(@emnapi/runtime@1.10.0)(@types/node@22.7.4)
|
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':
|
'@types/axios':
|
||||||
specifier: ^0.14.0
|
specifier: ^0.14.0
|
||||||
version: 0.14.4
|
version: 0.14.4
|
||||||
@@ -1313,32 +1307,6 @@ packages:
|
|||||||
'@octokit/types@16.0.0':
|
'@octokit/types@16.0.0':
|
||||||
resolution: {integrity: sha512-sKq+9r1Mm4efXW1FCk7hFSeJo4QKreL/tTbR0rz/qx/r1Oa2VV83LTA/H/MuCOX7uCIJmQVRKBcbmWoySjAnSg==}
|
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':
|
'@protobufjs/aspromise@1.1.2':
|
||||||
resolution: {integrity: sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==}
|
resolution: {integrity: sha512-j+gKExEuLmKwvz3OgROXtrJ2UG2x8Ch2YZUxahh+s1F2HZ+wAceUNLkvy6zKCPVRkU++ZWQrdxsUeQXmcg4uoQ==}
|
||||||
|
|
||||||
@@ -4957,27 +4925,6 @@ snapshots:
|
|||||||
dependencies:
|
dependencies:
|
||||||
'@octokit/openapi-types': 27.0.0
|
'@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':
|
'@protobufjs/aspromise@1.1.2':
|
||||||
optional: true
|
optional: true
|
||||||
|
|
||||||
|
|||||||
@@ -9,11 +9,8 @@ use lancedb::index::vector::{
|
|||||||
IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
|
IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder,
|
||||||
IvfRqIndexBuilder,
|
IvfRqIndexBuilder,
|
||||||
};
|
};
|
||||||
use lancedb::tokenize as lancedb_tokenize;
|
|
||||||
use napi_derive::napi;
|
use napi_derive::napi;
|
||||||
|
|
||||||
use crate::error::NapiErrorExt;
|
|
||||||
use crate::table::FtsToken;
|
|
||||||
use crate::util::parse_distance_type;
|
use crate::util::parse_distance_type;
|
||||||
|
|
||||||
#[napi]
|
#[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]
|
#[napi]
|
||||||
impl Index {
|
impl Index {
|
||||||
#[napi(factory)]
|
#[napi(factory)]
|
||||||
|
|||||||
@@ -12,7 +12,6 @@ mod header;
|
|||||||
mod index;
|
mod index;
|
||||||
mod iterator;
|
mod iterator;
|
||||||
pub mod merge;
|
pub mod merge;
|
||||||
pub mod otel;
|
|
||||||
pub mod permutation;
|
pub mod permutation;
|
||||||
mod query;
|
mod query;
|
||||||
pub mod remote;
|
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 counts: Option<Vec<i64>>,
|
||||||
pub fixed: Option<i64>,
|
pub fixed: Option<i64>,
|
||||||
pub seed: Option<i64>,
|
pub seed: Option<i64>,
|
||||||
pub clump_size: Option<i64>,
|
|
||||||
pub split_names: Option<Vec<String>>,
|
pub split_names: Option<Vec<String>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -126,15 +125,10 @@ impl PermutationBuilder {
|
|||||||
};
|
};
|
||||||
|
|
||||||
let seed = options.seed.map(|s| s as u64);
|
let seed = options.seed.map(|s| s as u64);
|
||||||
let clump_size = options.clump_size.map(|c| c as u64);
|
|
||||||
|
|
||||||
self.modify(|builder| {
|
self.modify(|builder| {
|
||||||
builder.with_split_strategy(
|
builder.with_split_strategy(
|
||||||
SplitStrategy::Random {
|
SplitStrategy::Random { seed, sizes },
|
||||||
seed,
|
|
||||||
sizes,
|
|
||||||
clump_size,
|
|
||||||
},
|
|
||||||
options.split_names.clone(),
|
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::ipc::{ipc_file_to_batches, ipc_file_to_schema};
|
||||||
use lancedb::table::{
|
use lancedb::table::{
|
||||||
AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration,
|
AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration,
|
||||||
FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
FieldMetadataUpdate as LanceFieldMetadataUpdate, NewColumnTransform, OptimizeAction,
|
||||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
OptimizeOptions, Ref, Table as LanceDbTable,
|
||||||
};
|
};
|
||||||
use napi::bindgen_prelude::*;
|
use napi::bindgen_prelude::*;
|
||||||
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
|
use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode};
|
||||||
@@ -411,16 +411,6 @@ impl Table {
|
|||||||
.default_error()
|
.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)]
|
#[napi(catch_unwind)]
|
||||||
pub async fn close_lsm_writers(&self) -> napi::Result<()> {
|
pub async fn close_lsm_writers(&self) -> napi::Result<()> {
|
||||||
self.inner_ref()?.close_lsm_writers().await.default_error()
|
self.inner_ref()?.close_lsm_writers().await.default_error()
|
||||||
@@ -574,27 +564,6 @@ impl Table {
|
|||||||
.collect::<Vec<_>>())
|
.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)]
|
#[napi(catch_unwind)]
|
||||||
pub async fn index_stats(&self, index_name: String) -> napi::Result<Option<IndexStatistics>> {
|
pub async fn index_stats(&self, index_name: String) -> napi::Result<Option<IndexStatistics>> {
|
||||||
let tbl = self.inner_ref()?;
|
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
|
/// Specification selecting Lance's MemWAL LSM-style write path for
|
||||||
/// `mergeInsert`.
|
/// `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.
|
/// Statistics about a compaction operation.
|
||||||
#[napi(object)]
|
#[napi(object)]
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone, Debug)]
|
||||||
|
|||||||
@@ -1,5 +1,5 @@
|
|||||||
[tool.bumpversion]
|
[tool.bumpversion]
|
||||||
current_version = "0.35.0-beta.2"
|
current_version = "0.34.0"
|
||||||
parse = """(?x)
|
parse = """(?x)
|
||||||
(?P<major>0|[1-9]\\d*)\\.
|
(?P<major>0|[1-9]\\d*)\\.
|
||||||
(?P<minor>0|[1-9]\\d*)\\.
|
(?P<minor>0|[1-9]\\d*)\\.
|
||||||
|
|||||||
+2
-2
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "lancedb-python"
|
name = "lancedb-python"
|
||||||
version = "0.35.0-beta.2"
|
version = "0.34.0"
|
||||||
publish = false
|
publish = false
|
||||||
edition.workspace = true
|
edition.workspace = true
|
||||||
description = "Python bindings for LanceDB"
|
description = "Python bindings for LanceDB"
|
||||||
@@ -47,6 +47,6 @@ pyo3-build-config = { version = "0.28", features = [
|
|||||||
] }
|
] }
|
||||||
|
|
||||||
[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"]
|
fp16kernels = ["lancedb/fp16kernels"]
|
||||||
remote = ["lancedb/remote"]
|
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 = [
|
||||||
"pylance>=5.0.0b5",
|
"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 = [
|
tests = [
|
||||||
"aiohttp>=3.9.0",
|
"aiohttp>=3.9.0",
|
||||||
"boto3>=1.28.57",
|
"boto3>=1.28.57",
|
||||||
@@ -65,7 +61,6 @@ tests = [
|
|||||||
"pylance>=5.0.0b5",
|
"pylance>=5.0.0b5",
|
||||||
"requests>=2.31.0",
|
"requests>=2.31.0",
|
||||||
"datafusion>=52,<53",
|
"datafusion>=52,<53",
|
||||||
"opentelemetry-sdk>=1.30.0",
|
|
||||||
]
|
]
|
||||||
dev = [
|
dev = [
|
||||||
"ruff>=0.3.0",
|
"ruff>=0.3.0",
|
||||||
|
|||||||
@@ -6,22 +6,19 @@ import importlib.metadata
|
|||||||
import os
|
import os
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from datetime import timedelta
|
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")
|
__version__ = importlib.metadata.version("lancedb")
|
||||||
|
|
||||||
from ._lancedb import connect as lancedb_connect
|
from ._lancedb import connect as lancedb_connect
|
||||||
from ._lancedb import FtsToken
|
|
||||||
from ._lancedb import tokenize as _tokenize
|
|
||||||
from .common import URI, sanitize_uri
|
from .common import URI, sanitize_uri
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
from .db import AsyncConnection, DBConnection, LanceDBConnection
|
from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||||
from .remote import ClientConfig
|
from .remote import ClientConfig
|
||||||
from .remote.db import RemoteDBConnection
|
from .remote.db import RemoteDBConnection
|
||||||
from .expr import Expr, col, lit, func
|
from .expr import Expr, col, lit, func
|
||||||
from .schema import blob, vector, BlobType
|
from .schema import vector
|
||||||
from .table import AsyncTable, Table
|
from .table import AsyncTable, Table
|
||||||
from .types import BaseTokenizerType
|
|
||||||
from ._lancedb import Session
|
from ._lancedb import Session
|
||||||
from .namespace import (
|
from .namespace import (
|
||||||
connect_namespace,
|
connect_namespace,
|
||||||
@@ -152,14 +149,8 @@ def connect(
|
|||||||
|
|
||||||
For object storage, use a URI prefix:
|
For object storage, use a URI prefix:
|
||||||
|
|
||||||
>>> db = lancedb.connect( # doctest: +SKIP
|
>>> db = lancedb.connect("s3://my-bucket/lancedb",
|
||||||
... "s3://my-bucket/lancedb",
|
... storage_options={"aws_access_key_id": "***"})
|
||||||
... storage_options={
|
|
||||||
... "aws_access_key_id": "***",
|
|
||||||
... "aws_secret_access_key": "***",
|
|
||||||
... "aws_region": "us-east-1",
|
|
||||||
... },
|
|
||||||
... )
|
|
||||||
|
|
||||||
For tests and temporary data, use an in-memory database:
|
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_"
|
WORKER_PROPERTY_PREFIX = "_lancedb_worker_"
|
||||||
|
|
||||||
|
|
||||||
@@ -493,21 +450,17 @@ async def connect_async(
|
|||||||
__all__ = [
|
__all__ = [
|
||||||
"connect",
|
"connect",
|
||||||
"connect_async",
|
"connect_async",
|
||||||
"tokenize",
|
|
||||||
"connect_namespace",
|
"connect_namespace",
|
||||||
"connect_namespace_async",
|
"connect_namespace_async",
|
||||||
"AsyncConnection",
|
"AsyncConnection",
|
||||||
"AsyncLanceNamespaceDBConnection",
|
"AsyncLanceNamespaceDBConnection",
|
||||||
"AsyncTable",
|
"AsyncTable",
|
||||||
"FtsToken",
|
|
||||||
"col",
|
"col",
|
||||||
"Expr",
|
"Expr",
|
||||||
"func",
|
"func",
|
||||||
"lit",
|
"lit",
|
||||||
"URI",
|
"URI",
|
||||||
"sanitize_uri",
|
"sanitize_uri",
|
||||||
"blob",
|
|
||||||
"BlobType",
|
|
||||||
"vector",
|
"vector",
|
||||||
"DBConnection",
|
"DBConnection",
|
||||||
"LanceDBConnection",
|
"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 datetime import datetime, timedelta
|
||||||
from decimal import Decimal
|
|
||||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||||
|
|
||||||
import pyarrow as pa
|
import pyarrow as pa
|
||||||
@@ -25,45 +24,11 @@ from lance_namespace import (
|
|||||||
ListTablesResponse,
|
ListTablesResponse,
|
||||||
)
|
)
|
||||||
from .remote import ClientConfig
|
from .remote import ClientConfig
|
||||||
from .types import BaseTokenizerType
|
|
||||||
|
|
||||||
IvfHnswPq: type[HnswPq] = HnswPq
|
IvfHnswPq: type[HnswPq] = HnswPq
|
||||||
IvfHnswSq: type[HnswSq] = HnswSq
|
IvfHnswSq: type[HnswSq] = HnswSq
|
||||||
IvfHnswFlat: type[HnswFlat] = HnswFlat
|
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:
|
class PyExpr:
|
||||||
"""A type-safe DataFusion expression node (Rust-side handle)."""
|
"""A type-safe DataFusion expression node (Rust-side handle)."""
|
||||||
|
|
||||||
@@ -88,9 +53,7 @@ class PyExpr:
|
|||||||
def to_sql(self) -> str: ...
|
def to_sql(self) -> str: ...
|
||||||
|
|
||||||
def expr_col(name: str) -> PyExpr: ...
|
def expr_col(name: str) -> PyExpr: ...
|
||||||
def expr_lit(
|
def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ...
|
||||||
value: Union[bool, int, float, str, bytes, date, datetime, Decimal],
|
|
||||||
) -> PyExpr: ...
|
|
||||||
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
||||||
|
|
||||||
class Session:
|
class Session:
|
||||||
@@ -196,17 +159,6 @@ class Connection(object):
|
|||||||
self,
|
self,
|
||||||
) -> Dict[str, Any]: ...
|
) -> 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:
|
class Table:
|
||||||
def name(self) -> str: ...
|
def name(self) -> str: ...
|
||||||
def __repr__(self) -> str: ...
|
def __repr__(self) -> str: ...
|
||||||
@@ -253,13 +205,6 @@ class Table:
|
|||||||
async def prewarm_index(self, index_name: str) -> None: ...
|
async def prewarm_index(self, index_name: str) -> None: ...
|
||||||
async def prewarm_data(self, columns: Optional[List[str]] = None) -> None: ...
|
async def prewarm_data(self, columns: Optional[List[str]] = None) -> None: ...
|
||||||
async def list_indices(self) -> list[IndexConfig]: ...
|
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 delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ...
|
||||||
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
|
async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ...
|
||||||
async def add_columns_with_schema(self, schema: pa.Schema) -> 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_unenforced_primary_key(self, columns: List[str]) -> None: ...
|
||||||
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
||||||
async def unset_lsm_write_spec(self) -> 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: ...
|
async def close_lsm_writers(self) -> None: ...
|
||||||
@property
|
@property
|
||||||
def tags(self) -> Tags: ...
|
def tags(self) -> Tags: ...
|
||||||
@@ -291,13 +235,6 @@ class Table:
|
|||||||
def query(self) -> Query: ...
|
def query(self) -> Query: ...
|
||||||
def take_offsets(self, offsets: list[int]) -> TakeQuery: ...
|
def take_offsets(self, offsets: list[int]) -> TakeQuery: ...
|
||||||
def take_row_ids(self, row_ids: 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: ...
|
def vector_search(self) -> VectorQuery: ...
|
||||||
|
|
||||||
class Tags:
|
class Tags:
|
||||||
@@ -533,10 +470,6 @@ class MergeResult:
|
|||||||
num_attempts: int
|
num_attempts: int
|
||||||
num_rows: int
|
num_rows: int
|
||||||
|
|
||||||
class FtsToken:
|
|
||||||
text: str
|
|
||||||
position: int
|
|
||||||
|
|
||||||
class LsmWriteSpec:
|
class LsmWriteSpec:
|
||||||
"""Specification selecting Lance's MemWAL LSM-style write path for
|
"""Specification selecting Lance's MemWAL LSM-style write path for
|
||||||
`merge_insert`."""
|
`merge_insert`."""
|
||||||
|
|||||||
@@ -4,7 +4,7 @@
|
|||||||
|
|
||||||
import os
|
import os
|
||||||
from functools import cached_property
|
from functools import cached_property
|
||||||
from typing import List, Optional, Union
|
from typing import List, Union
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
@@ -15,8 +15,6 @@ from .base import TextEmbeddingFunction
|
|||||||
from .registry import register
|
from .registry import register
|
||||||
from .utils import TEXT, api_key_not_found_help
|
from .utils import TEXT, api_key_not_found_help
|
||||||
|
|
||||||
EMBEDDING_BATCH_SIZE = 100
|
|
||||||
|
|
||||||
|
|
||||||
@register("gemini-text")
|
@register("gemini-text")
|
||||||
class GeminiText(TextEmbeddingFunction):
|
class GeminiText(TextEmbeddingFunction):
|
||||||
@@ -83,7 +81,6 @@ class GeminiText(TextEmbeddingFunction):
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
name: str = "gemini-embedding-001"
|
name: str = "gemini-embedding-001"
|
||||||
dim: Optional[int] = None
|
|
||||||
query_task_type: str = "retrieval_query"
|
query_task_type: str = "retrieval_query"
|
||||||
source_task_type: str = "retrieval_document"
|
source_task_type: str = "retrieval_document"
|
||||||
|
|
||||||
@@ -96,8 +93,6 @@ class GeminiText(TextEmbeddingFunction):
|
|||||||
model_config["ignored_types"] = (cached_property,)
|
model_config["ignored_types"] = (cached_property,)
|
||||||
|
|
||||||
def ndims(self):
|
def ndims(self):
|
||||||
if self.dim:
|
|
||||||
return self.dim
|
|
||||||
# TODO: fix hardcoding
|
# TODO: fix hardcoding
|
||||||
return 768
|
return 768
|
||||||
|
|
||||||
@@ -138,22 +133,22 @@ class GeminiText(TextEmbeddingFunction):
|
|||||||
contents.append({"parts": [{"text": text}]})
|
contents.append({"parts": [{"text": text}]})
|
||||||
|
|
||||||
# Build config
|
# Build config
|
||||||
config_kwargs = {"output_dimensionality": self.ndims()}
|
config_kwargs = {}
|
||||||
if task_type:
|
if task_type:
|
||||||
config_kwargs["task_type"] = task_type.upper() # API expects uppercase
|
config_kwargs["task_type"] = task_type.upper() # API expects uppercase
|
||||||
|
|
||||||
config = types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
# Call embed_content for each content
|
||||||
|
|
||||||
# Call embed_content in groups of at most EMBEDDING_BATCH_SIZE docs at a time
|
|
||||||
embeddings = []
|
embeddings = []
|
||||||
for i in range(0, len(contents), EMBEDDING_BATCH_SIZE):
|
for content in contents:
|
||||||
chunk = contents[i : i + EMBEDDING_BATCH_SIZE]
|
config = (
|
||||||
|
types.EmbedContentConfig(**config_kwargs) if config_kwargs else None
|
||||||
|
)
|
||||||
response = self.client.models.embed_content(
|
response = self.client.models.embed_content(
|
||||||
model=self.name,
|
model=self.name,
|
||||||
contents=chunk,
|
contents=content,
|
||||||
config=config,
|
config=config,
|
||||||
)
|
)
|
||||||
embeddings.extend([np.array(e.values) for e in response.embeddings])
|
embeddings.append(response.embeddings[0].values)
|
||||||
|
|
||||||
return embeddings
|
return embeddings
|
||||||
|
|
||||||
@@ -165,13 +160,5 @@ class GeminiText(TextEmbeddingFunction):
|
|||||||
api_key_not_found_help("google")
|
api_key_not_found_help("google")
|
||||||
|
|
||||||
from google import genai as genai_module
|
from google import genai as genai_module
|
||||||
from lancedb import __version__
|
|
||||||
|
|
||||||
return genai_module.Client(
|
return genai_module.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
|
||||||
api_key=os.environ.get("GOOGLE_API_KEY"),
|
|
||||||
http_options={
|
|
||||||
"headers": {
|
|
||||||
"x-goog-api-client": f"lancedb/{__version__}",
|
|
||||||
}
|
|
||||||
},
|
|
||||||
)
|
|
||||||
|
|||||||
@@ -2,7 +2,7 @@
|
|||||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
from functools import cached_property
|
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
|
import numpy as np
|
||||||
|
|
||||||
@@ -56,16 +56,6 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
|||||||
embeddings = self._compute_embedding(texts)
|
embeddings = self._compute_embedding(texts)
|
||||||
return list(embeddings)
|
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
|
@cached_property
|
||||||
def _ollama_client(self) -> "ollama.Client":
|
def _ollama_client(self) -> "ollama.Client":
|
||||||
ollama = attempt_import_or_raise("ollama")
|
ollama = attempt_import_or_raise("ollama")
|
||||||
|
|||||||
@@ -19,8 +19,6 @@ operators::
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from datetime import date, datetime
|
|
||||||
from decimal import Decimal
|
|
||||||
from typing import Iterable, Union
|
from typing import Iterable, Union
|
||||||
|
|
||||||
import pyarrow as pa
|
import pyarrow as pa
|
||||||
@@ -65,7 +63,7 @@ def _coerce(value: "ExprLike") -> "Expr":
|
|||||||
|
|
||||||
|
|
||||||
# Type alias used in annotations.
|
# 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:
|
class Expr:
|
||||||
@@ -120,18 +118,10 @@ class Expr:
|
|||||||
"""Logical AND (``expr_a & expr_b``)."""
|
"""Logical AND (``expr_a & expr_b``)."""
|
||||||
return Expr(self._inner.and_(_coerce(other)._inner))
|
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":
|
def __or__(self, other: "Expr") -> "Expr":
|
||||||
"""Logical OR (``expr_a | expr_b``)."""
|
"""Logical OR (``expr_a | expr_b``)."""
|
||||||
return Expr(self._inner.or_(_coerce(other)._inner))
|
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":
|
def __invert__(self) -> "Expr":
|
||||||
"""Logical NOT (``~expr``)."""
|
"""Logical NOT (``~expr``)."""
|
||||||
return Expr(self._inner.not_())
|
return Expr(self._inner.not_())
|
||||||
@@ -276,14 +266,13 @@ def col(name: str) -> Expr:
|
|||||||
return Expr(expr_col(name))
|
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.
|
"""Create a literal (constant) value expression.
|
||||||
|
|
||||||
Parameters
|
Parameters
|
||||||
----------
|
----------
|
||||||
value:
|
value:
|
||||||
A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``,
|
A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``.
|
||||||
``datetime``, or ``Decimal``.
|
|
||||||
|
|
||||||
Examples
|
Examples
|
||||||
--------
|
--------
|
||||||
@@ -291,9 +280,6 @@ def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) ->
|
|||||||
>>> col("price") * lit(1.1)
|
>>> col("price") * lit(1.1)
|
||||||
Expr((price * 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))
|
return Expr(expr_lit(value))
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -127,8 +127,6 @@ class FTS:
|
|||||||
- "whitespace": Split text by whitespace, but not punctuation.
|
- "whitespace": Split text by whitespace, but not punctuation.
|
||||||
- "raw": No tokenization. The entire text is treated as a single token.
|
- "raw": No tokenization. The entire text is treated as a single token.
|
||||||
- "ngram": N-gram tokenizer for substring-style matching.
|
- "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.
|
- "jieba/*": Jieba tokenizer loaded from Lance's language model home.
|
||||||
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
||||||
language : str, default "English"
|
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 ._lancedb import async_permutation_builder, PermutationReader
|
||||||
from .table import LanceTable, Table
|
from .table import LanceTable, Table
|
||||||
from .background_loop import LOOP
|
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
|
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -65,7 +65,6 @@ class PermutationBuilder:
|
|||||||
counts: Optional[list[int]] = None,
|
counts: Optional[list[int]] = None,
|
||||||
fixed: Optional[int] = None,
|
fixed: Optional[int] = None,
|
||||||
seed: Optional[int] = None,
|
seed: Optional[int] = None,
|
||||||
clump_size: Optional[int] = None,
|
|
||||||
split_names: Optional[list[str]] = None,
|
split_names: Optional[list[str]] = None,
|
||||||
) -> "PermutationBuilder":
|
) -> "PermutationBuilder":
|
||||||
"""
|
"""
|
||||||
@@ -88,9 +87,6 @@ class PermutationBuilder:
|
|||||||
Rows will be randomly assigned to splits. The optional seed can be provided to
|
Rows will be randomly assigned to splits. The optional seed can be provided to
|
||||||
make the assignment deterministic.
|
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 optional split_names can be provided to name the splits. If not provided,
|
||||||
the splits can only be referenced by their index.
|
the splits can only be referenced by their index.
|
||||||
"""
|
"""
|
||||||
@@ -99,7 +95,6 @@ class PermutationBuilder:
|
|||||||
counts=counts,
|
counts=counts,
|
||||||
fixed=fixed,
|
fixed=fixed,
|
||||||
seed=seed,
|
seed=seed,
|
||||||
clump_size=clump_size,
|
|
||||||
split_names=split_names,
|
split_names=split_names,
|
||||||
)
|
)
|
||||||
return self
|
return self
|
||||||
@@ -946,7 +941,6 @@ class Permutation:
|
|||||||
"pandas",
|
"pandas",
|
||||||
"arrow",
|
"arrow",
|
||||||
"torch",
|
"torch",
|
||||||
"torch_row",
|
|
||||||
"torch_col",
|
"torch_col",
|
||||||
"polars",
|
"polars",
|
||||||
],
|
],
|
||||||
@@ -962,19 +956,15 @@ class Permutation:
|
|||||||
- "python_col" - the batch will be a dict of lists (one entry per column)
|
- "python_col" - the batch will be a dict of lists (one entry per column)
|
||||||
- "pandas" - the batch will be a pandas DataFrame
|
- "pandas" - the batch will be a pandas DataFrame
|
||||||
- "arrow" - the batch will be a pyarrow RecordBatch
|
- "arrow" - the batch will be a pyarrow RecordBatch
|
||||||
- "torch" - the batch will be a list of per-row dicts mapping column
|
- "torch" - the batch will be a list of tensors, one per row
|
||||||
name to a 0-D torch tensor. Works with the default
|
|
||||||
``torch.utils.data.DataLoader`` collate, which stacks the per-row
|
|
||||||
dicts back into a dict of batched tensors.
|
|
||||||
- "torch_row" - the batch will be a list of tensors, one per row
|
|
||||||
- "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns)
|
- "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns)
|
||||||
- "polars" - the batch will be a polars DataFrame
|
- "polars" - the batch will be a polars DataFrame
|
||||||
|
|
||||||
Conversion may or may not involve a data copy. Lance uses Arrow internally
|
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.
|
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
|
Conversion to torch_col will be zero-copy but will only support a subset of data
|
||||||
subset of data types (numeric types).
|
types (numeric types).
|
||||||
|
|
||||||
Conversion to numpy and/or pandas will typically be zero-copy for numeric
|
Conversion to numpy and/or pandas will typically be zero-copy for numeric
|
||||||
types. Conversion of strings, lists, and structs will require creating python
|
types. Conversion of strings, lists, and structs will require creating python
|
||||||
@@ -995,8 +985,6 @@ class Permutation:
|
|||||||
elif format == "arrow":
|
elif format == "arrow":
|
||||||
return self.with_transform(Transforms.arrow2arrow)
|
return self.with_transform(Transforms.arrow2arrow)
|
||||||
elif format == "torch":
|
elif format == "torch":
|
||||||
return self.with_transform(batch_to_tensor_dict)
|
|
||||||
elif format == "torch_row":
|
|
||||||
return self.with_transform(batch_to_tensor_rows)
|
return self.with_transform(batch_to_tensor_rows)
|
||||||
elif format == "torch_col":
|
elif format == "torch_col":
|
||||||
return self.with_transform(batch_to_tensor)
|
return self.with_transform(batch_to_tensor)
|
||||||
|
|||||||
+78
-299
@@ -15,12 +15,10 @@ from typing import (
|
|||||||
List,
|
List,
|
||||||
Literal,
|
Literal,
|
||||||
Optional,
|
Optional,
|
||||||
Protocol,
|
|
||||||
Tuple,
|
Tuple,
|
||||||
Type,
|
Type,
|
||||||
TypeVar,
|
TypeVar,
|
||||||
Union,
|
Union,
|
||||||
runtime_checkable,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
import deprecation
|
import deprecation
|
||||||
@@ -41,21 +39,15 @@ from .expr import Expr
|
|||||||
from .rerankers.base import Reranker
|
from .rerankers.base import Reranker
|
||||||
from .rerankers.rrf import RRFReranker
|
from .rerankers.rrf import RRFReranker
|
||||||
from .rerankers.util import check_reranker_result
|
from .rerankers.util import check_reranker_result
|
||||||
from .schema import is_blob_like_field, schema_has_blob_field
|
|
||||||
from .util import flatten_columns
|
from .util import flatten_columns
|
||||||
from ._blob import (
|
|
||||||
BLOB_MODE_TO_HANDLING,
|
BlobMode = Literal["lazy", "bytes", "descriptions"]
|
||||||
FetchBlobsAsync,
|
|
||||||
FetchBlobsSync,
|
_BLOB_MODE_TO_HANDLING = {
|
||||||
blob_auto_row_id_for_scan,
|
"lazy": "blobs_descriptions",
|
||||||
blob_v2_projection_sources,
|
"bytes": "all_binary",
|
||||||
finalize_blob_query_table,
|
"descriptions": "blobs_descriptions",
|
||||||
replace_v2_blob_columns_with_bytes,
|
}
|
||||||
replace_v2_blob_columns_with_bytes_sync,
|
|
||||||
supports_blob_auto_row_id,
|
|
||||||
validate_blob_mode,
|
|
||||||
)
|
|
||||||
from .types import BlobMode, QueryProjection
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
import sys
|
import sys
|
||||||
@@ -81,22 +73,25 @@ if TYPE_CHECKING:
|
|||||||
T = TypeVar("T", bound="LanceModel")
|
T = TypeVar("T", bound="LanceModel")
|
||||||
|
|
||||||
|
|
||||||
@runtime_checkable
|
def _validate_blob_mode(blob_mode: BlobMode) -> None:
|
||||||
class _LanceScanner(Protocol):
|
if blob_mode not in _BLOB_MODE_TO_HANDLING:
|
||||||
projected_schema: pa.Schema | None
|
modes = ", ".join(repr(mode) for mode in _BLOB_MODE_TO_HANDLING)
|
||||||
schema: pa.Schema | None
|
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:
|
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:
|
def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
|
||||||
@@ -145,7 +140,13 @@ def _combine_where(
|
|||||||
return f"({existing_sql}) AND ({new_sql})"
|
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:
|
if columns is None:
|
||||||
return {}
|
return {}
|
||||||
if isinstance(columns, list):
|
if isinstance(columns, list):
|
||||||
@@ -170,11 +171,7 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
|||||||
|
|
||||||
|
|
||||||
def _scanner_kwargs_for_query(
|
def _scanner_kwargs_for_query(
|
||||||
query: Query,
|
query: Query, blob_mode: BlobMode, dataset: Optional[Any] = None
|
||||||
blob_mode: BlobMode,
|
|
||||||
dataset: Optional[Any] = None,
|
|
||||||
*,
|
|
||||||
with_row_id: Optional[bool] = None,
|
|
||||||
) -> Dict[str, Any]:
|
) -> Dict[str, Any]:
|
||||||
fragments = _scanner_fragments_for_query(query, dataset)
|
fragments = _scanner_fragments_for_query(query, dataset)
|
||||||
kwargs = {
|
kwargs = {
|
||||||
@@ -182,10 +179,10 @@ def _scanner_kwargs_for_query(
|
|||||||
"filter": _filter_to_sql(query.filter),
|
"filter": _filter_to_sql(query.filter),
|
||||||
"limit": query.limit,
|
"limit": query.limit,
|
||||||
"offset": query.offset,
|
"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,
|
"with_row_address": query.with_row_address,
|
||||||
"fast_search": query.fast_search,
|
"fast_search": query.fast_search,
|
||||||
"blob_handling": BLOB_MODE_TO_HANDLING[blob_mode],
|
"blob_handling": _BLOB_MODE_TO_HANDLING[blob_mode],
|
||||||
"fragments": fragments,
|
"fragments": fragments,
|
||||||
}
|
}
|
||||||
return {key: value for key, value in kwargs.items() if value is not None}
|
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(
|
def _ensure_lazy_blob_frame(
|
||||||
df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode
|
df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode
|
||||||
) -> "pd.DataFrame":
|
) -> "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
|
return df
|
||||||
|
|
||||||
for field in schema:
|
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
|
continue
|
||||||
value = df[field.name].iloc[0]
|
value = df[field.name].iloc[0]
|
||||||
if value is not None and not hasattr(value, "readall"):
|
if value is not None and not hasattr(value, "readall"):
|
||||||
@@ -232,7 +229,7 @@ def _ensure_lazy_blob_frame(
|
|||||||
return df
|
return df
|
||||||
|
|
||||||
|
|
||||||
def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
def _scanner_to_table(scanner: Any) -> pa.Table:
|
||||||
if hasattr(scanner, "to_pyarrow"):
|
if hasattr(scanner, "to_pyarrow"):
|
||||||
reader = scanner.to_pyarrow()
|
reader = scanner.to_pyarrow()
|
||||||
return reader.read_all()
|
return reader.read_all()
|
||||||
@@ -242,9 +239,7 @@ def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
|||||||
return reader.read_all()
|
return reader.read_all()
|
||||||
|
|
||||||
|
|
||||||
def _scanner_to_pandas(
|
def _scanner_to_pandas(scanner: Any, blob_mode: BlobMode, **kwargs) -> "pd.DataFrame":
|
||||||
scanner: _LanceScanner, blob_mode: BlobMode, **kwargs
|
|
||||||
) -> pd.DataFrame:
|
|
||||||
schema = getattr(scanner, "projected_schema", None)
|
schema = getattr(scanner, "projected_schema", None)
|
||||||
if schema is None:
|
if schema is None:
|
||||||
schema = getattr(scanner, "schema", None)
|
schema = getattr(scanner, "schema", None)
|
||||||
@@ -265,71 +260,13 @@ def _scanner_to_pandas(
|
|||||||
return df
|
return df
|
||||||
|
|
||||||
tbl = _scanner_to_table(scanner)
|
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(
|
raise _unsupported_blob_pandas_error(
|
||||||
"the Lance scanner does not expose to_pandas"
|
"the Lance scanner does not expose to_pandas"
|
||||||
)
|
)
|
||||||
return tbl.to_pandas(**kwargs)
|
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
|
# Pydantic validation function for vector queries
|
||||||
def ensure_vector_query(
|
def ensure_vector_query(
|
||||||
val: Any,
|
val: Any,
|
||||||
@@ -737,7 +674,7 @@ class Query(pydantic.BaseModel):
|
|||||||
distance_type: Optional[str] = None
|
distance_type: Optional[str] = None
|
||||||
|
|
||||||
# which columns to return in the results (dict values may be str or Expr)
|
# 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
|
# minimum number of IVF partitions to search
|
||||||
#
|
#
|
||||||
@@ -1021,7 +958,7 @@ class LanceQueryBuilder(ABC):
|
|||||||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||||
optional flattening.
|
optional flattening.
|
||||||
"""
|
"""
|
||||||
validate_blob_mode(blob_mode)
|
_validate_blob_mode(blob_mode)
|
||||||
output_schema = getattr(self, "output_schema", None)
|
output_schema = getattr(self, "output_schema", None)
|
||||||
if output_schema is not None:
|
if output_schema is not None:
|
||||||
schema = output_schema()
|
schema = output_schema()
|
||||||
@@ -1080,11 +1017,6 @@ class LanceQueryBuilder(ABC):
|
|||||||
Execute the query and return the results as a pyarrow
|
Execute the query and return the results as a pyarrow
|
||||||
[RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html)
|
[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
|
Parameters
|
||||||
----------
|
----------
|
||||||
batch_size: int
|
batch_size: int
|
||||||
@@ -1263,42 +1195,6 @@ class LanceQueryBuilder(ABC):
|
|||||||
self._with_row_id = with_row_id
|
self._with_row_id = with_row_id
|
||||||
return self
|
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:
|
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||||||
"""Set whether to return row addresses.
|
"""Set whether to return row addresses.
|
||||||
|
|
||||||
@@ -1475,29 +1371,13 @@ class LanceQueryBuilder(ABC):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
dataset = self._table.to_lance()
|
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 = dataset.scanner(
|
||||||
**_scanner_kwargs_for_query(
|
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||||
query,
|
|
||||||
"descriptions" if blob_sources else blob_mode,
|
|
||||||
dataset,
|
|
||||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return _finish_plain_scan_pandas(
|
|
||||||
scanner,
|
|
||||||
blob_mode=blob_mode,
|
|
||||||
blob_sources=blob_sources,
|
|
||||||
fetch_blobs=self._table.fetch_blobs,
|
|
||||||
strip_auto_row_id=blob_auto_row_id,
|
|
||||||
flatten=flatten,
|
|
||||||
**kwargs,
|
|
||||||
)
|
)
|
||||||
|
if flatten is not None:
|
||||||
|
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||||
|
return tbl.to_pandas(**kwargs)
|
||||||
|
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def to_query_object(self) -> Query:
|
def to_query_object(self) -> Query:
|
||||||
@@ -1745,9 +1625,7 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
|||||||
The maximum time to wait for the query to complete.
|
The maximum time to wait for the query to complete.
|
||||||
If None, wait indefinitely.
|
If None, wait indefinitely.
|
||||||
"""
|
"""
|
||||||
return self._finalize_blob_query_table(
|
return self.to_batches(timeout=timeout).read_all()
|
||||||
self.to_batches(timeout=timeout).read_all()
|
|
||||||
)
|
|
||||||
|
|
||||||
def to_query_object(self) -> Query:
|
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()
|
vector = self._query if isinstance(self._query, list) else self._query.tolist()
|
||||||
if isinstance(vector[0], np.ndarray):
|
if isinstance(vector[0], np.ndarray):
|
||||||
vector = [v.tolist() for v in vector]
|
vector = [v.tolist() for v in vector]
|
||||||
query = self._query_for_scan()
|
query = self.to_query_object()
|
||||||
result_set = self._table._execute_query(
|
result_set = self._table._execute_query(
|
||||||
query, batch_size=batch_size, timeout=timeout
|
query, batch_size=batch_size, timeout=timeout
|
||||||
)
|
)
|
||||||
@@ -1951,7 +1829,8 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
|||||||
Parameters
|
Parameters
|
||||||
----------
|
----------
|
||||||
phrase_query: bool, default True
|
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
|
Returns
|
||||||
-------
|
-------
|
||||||
@@ -1961,21 +1840,6 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
|||||||
self._phrase_query = phrase_query
|
self._phrase_query = phrase_query
|
||||||
return self
|
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:
|
def fast_search(self) -> LanceFtsQueryBuilder:
|
||||||
"""
|
"""
|
||||||
Skip a flat search of unindexed data. This will improve
|
Skip a flat search of unindexed data. This will improve
|
||||||
@@ -2000,7 +1864,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
|||||||
fragments=self._fragments,
|
fragments=self._fragments,
|
||||||
fragment_ids=self._fragment_ids,
|
fragment_ids=self._fragment_ids,
|
||||||
full_text_query=FullTextSearchQuery(
|
full_text_query=FullTextSearchQuery(
|
||||||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
query=self._query, columns=self._fts_columns
|
||||||
),
|
),
|
||||||
offset=self._offset,
|
offset=self._offset,
|
||||||
fast_search=self._fast_search,
|
fast_search=self._fast_search,
|
||||||
@@ -2018,13 +1882,22 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
|||||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||||
self._table._ensure_no_legacy_fts_index()
|
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 = self._table._execute_query(query, timeout=timeout)
|
||||||
results = results.read_all()
|
results = results.read_all()
|
||||||
if self._reranker is not None:
|
if self._reranker is not None:
|
||||||
results = self._reranker.rerank_fts(self._query, results)
|
results = self._reranker.rerank_fts(self._query, results)
|
||||||
check_reranker_result(results)
|
check_reranker_result(results)
|
||||||
return self._finalize_blob_query_table(results)
|
return results
|
||||||
|
|
||||||
def to_batches(
|
def to_batches(
|
||||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||||
@@ -2052,9 +1925,7 @@ class LanceFtsQueryBuilder(LanceQueryBuilder):
|
|||||||
|
|
||||||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||||||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||||||
return self._finalize_blob_query_table(
|
return self.to_batches(timeout=timeout).read_all()
|
||||||
self.to_batches(timeout=timeout).read_all()
|
|
||||||
)
|
|
||||||
|
|
||||||
def to_query_object(self) -> Query:
|
def to_query_object(self) -> Query:
|
||||||
return Query(
|
return Query(
|
||||||
@@ -2076,7 +1947,7 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
|||||||
def to_batches(
|
def to_batches(
|
||||||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||||||
) -> pa.RecordBatchReader:
|
) -> pa.RecordBatchReader:
|
||||||
query = self._query_for_scan()
|
query = self.to_query_object()
|
||||||
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
||||||
|
|
||||||
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
||||||
@@ -2148,13 +2019,14 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
|||||||
|
|
||||||
return vector_query, text_query
|
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.
|
"""Set whether to use phrase query.
|
||||||
|
|
||||||
Parameters
|
Parameters
|
||||||
----------
|
----------
|
||||||
phrase_query: bool, default True
|
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
|
Returns
|
||||||
-------
|
-------
|
||||||
@@ -2179,25 +2051,15 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
|||||||
fts_results = fts_future.result()
|
fts_results = fts_future.result()
|
||||||
vector_results = vector_future.result()
|
vector_results = vector_future.result()
|
||||||
|
|
||||||
results = self._combine_hybrid_results(
|
return self._combine_hybrid_results(
|
||||||
fts_results=fts_results,
|
fts_results=fts_results,
|
||||||
vector_results=vector_results,
|
vector_results=vector_results,
|
||||||
norm=self._norm,
|
norm=self._norm,
|
||||||
fts_query=self._fts_query._query,
|
fts_query=self._fts_query._query,
|
||||||
reranker=self._reranker,
|
reranker=self._reranker,
|
||||||
limit=self._limit,
|
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
|
@staticmethod
|
||||||
def _combine_hybrid_results(
|
def _combine_hybrid_results(
|
||||||
@@ -2638,7 +2500,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
|
|||||||
self._vector_query.ef(self._ef)
|
self._vector_query.ef(self._ef)
|
||||||
if self._bypass_vector_index:
|
if self._bypass_vector_index:
|
||||||
self._vector_query.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(
|
self._vector_query.distance_range(
|
||||||
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
||||||
)
|
)
|
||||||
@@ -2668,9 +2530,6 @@ class AsyncQueryBase(object):
|
|||||||
self._with_row_address = None
|
self._with_row_address = None
|
||||||
self._fragments = None
|
self._fragments = None
|
||||||
self._fragment_ids = 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:
|
def to_query_object(self) -> Query:
|
||||||
"""
|
"""
|
||||||
@@ -2680,46 +2539,11 @@ class AsyncQueryBase(object):
|
|||||||
python and more easily serializable.
|
python and more easily serializable.
|
||||||
"""
|
"""
|
||||||
query = Query.from_inner(self._inner.to_query_request())
|
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.with_row_address = self._with_row_address
|
||||||
query.fragments = self._fragments
|
query.fragments = self._fragments
|
||||||
query.fragment_ids = self._fragment_ids
|
query.fragment_ids = self._fragment_ids
|
||||||
return query
|
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:
|
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||||||
"""
|
"""
|
||||||
Return only the specified columns.
|
Return only the specified columns.
|
||||||
@@ -2772,7 +2596,6 @@ class AsyncQueryBase(object):
|
|||||||
"""
|
"""
|
||||||
Include the _rowid column in the results.
|
Include the _rowid column in the results.
|
||||||
"""
|
"""
|
||||||
self._with_row_id = True
|
|
||||||
self._inner.with_row_id()
|
self._inner.with_row_id()
|
||||||
return self
|
return self
|
||||||
|
|
||||||
@@ -2819,7 +2642,6 @@ class AsyncQueryBase(object):
|
|||||||
If not specified, no timeout is applied. If the query does not
|
If not specified, no timeout is applied. If the query does not
|
||||||
complete within the specified time, an error will be raised.
|
complete within the specified time, an error will be raised.
|
||||||
"""
|
"""
|
||||||
await self._maybe_add_blob_row_id()
|
|
||||||
return AsyncRecordBatchReader(
|
return AsyncRecordBatchReader(
|
||||||
await self._inner.execute(
|
await self._inner.execute(
|
||||||
max_batch_length=max_batch_length, timeout=timeout
|
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.
|
complete within the specified time, an error will be raised.
|
||||||
"""
|
"""
|
||||||
batch_iter = await self.to_batches(timeout=timeout)
|
batch_iter = await self.to_batches(timeout=timeout)
|
||||||
return self._finalize_blob_query_table(
|
return pa.Table.from_batches(
|
||||||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
await batch_iter.read_all(), schema=batch_iter.schema
|
||||||
)
|
)
|
||||||
|
|
||||||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
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
|
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||||||
optional flattening.
|
optional flattening.
|
||||||
"""
|
"""
|
||||||
validate_blob_mode(blob_mode)
|
_validate_blob_mode(blob_mode)
|
||||||
if hasattr(self._inner, "output_schema"):
|
if hasattr(self._inner, "output_schema"):
|
||||||
schema = await self.output_schema()
|
schema = await self.output_schema()
|
||||||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||||||
@@ -2959,36 +2781,14 @@ class AsyncQueryBase(object):
|
|||||||
if not _query_is_plain_scan(query):
|
if not _query_is_plain_scan(query):
|
||||||
return None
|
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()
|
dataset = await self._table._to_lance()
|
||||||
scanner = dataset.scanner(
|
scanner = dataset.scanner(
|
||||||
**_scanner_kwargs_for_query(
|
**_scanner_kwargs_for_query(query, blob_mode, dataset)
|
||||||
query,
|
|
||||||
"descriptions" if blob_sources else blob_mode,
|
|
||||||
dataset,
|
|
||||||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return await _finish_plain_scan_pandas_async(
|
|
||||||
scanner,
|
|
||||||
blob_mode=blob_mode,
|
|
||||||
blob_sources=blob_sources,
|
|
||||||
fetch_blobs=self._table.fetch_blobs,
|
|
||||||
strip_auto_row_id=blob_auto_row_id,
|
|
||||||
flatten=flatten,
|
|
||||||
**kwargs,
|
|
||||||
)
|
)
|
||||||
|
if flatten is not None:
|
||||||
|
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||||||
|
return tbl.to_pandas(**kwargs)
|
||||||
|
return _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||||||
|
|
||||||
async def to_polars(
|
async def to_polars(
|
||||||
self,
|
self,
|
||||||
@@ -3773,24 +3573,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
|||||||
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
||||||
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||||||
|
|
||||||
req = fts_query._inner.to_query_request()
|
# save the row ID choice that was made on the query builder and force it
|
||||||
blob_auto_row_id = False
|
# to actually fetch the row ids because we need this for reranking
|
||||||
blob_paths: tuple[str, ...] = ()
|
with_row_ids = self._inner.get_with_row_id()
|
||||||
if self._table is not None and supports_blob_auto_row_id(self._table):
|
|
||||||
schema = await self._table.schema()
|
|
||||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
|
||||||
self._table,
|
|
||||||
schema,
|
|
||||||
req.select,
|
|
||||||
with_row_id=self._with_row_id,
|
|
||||||
)
|
|
||||||
if blob_auto_row_id:
|
|
||||||
blob_paths = tuple(
|
|
||||||
blob_v2_projection_sources(schema, req.select).keys()
|
|
||||||
)
|
|
||||||
self._blob_auto_row_id = blob_auto_row_id
|
|
||||||
self._blob_paths = blob_paths
|
|
||||||
|
|
||||||
fts_query.with_row_id()
|
fts_query.with_row_id()
|
||||||
vec_query.with_row_id()
|
vec_query.with_row_id()
|
||||||
|
|
||||||
@@ -3806,14 +3591,8 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
|||||||
fts_query=fts_query.get_query(),
|
fts_query=fts_query.get_query(),
|
||||||
reranker=self._reranker,
|
reranker=self._reranker,
|
||||||
limit=self._inner.get_limit(),
|
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)
|
return AsyncRecordBatchReader(result, max_batch_length=max_batch_length)
|
||||||
|
|
||||||
|
|||||||
@@ -28,7 +28,6 @@ from lancedb._lancedb import (
|
|||||||
UpdateFieldMetadataResult,
|
UpdateFieldMetadataResult,
|
||||||
DeleteResult,
|
DeleteResult,
|
||||||
DropColumnsResult,
|
DropColumnsResult,
|
||||||
FtsToken,
|
|
||||||
IndexConfig,
|
IndexConfig,
|
||||||
LsmWriteSpec,
|
LsmWriteSpec,
|
||||||
MergeResult,
|
MergeResult,
|
||||||
@@ -245,23 +244,6 @@ class RemoteTable(Table):
|
|||||||
"""List all the indices on the table"""
|
"""List all the indices on the table"""
|
||||||
return LOOP.run(self._table.list_indices())
|
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]:
|
def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]:
|
||||||
"""List all the stats of a specified index"""
|
"""List all the stats of a specified index"""
|
||||||
return LOOP.run(self._table.index_stats(index_uuid))
|
return LOOP.run(self._table.index_stats(index_uuid))
|
||||||
@@ -930,10 +912,6 @@ class RemoteTable(Table):
|
|||||||
"""Not supported on LanceDB Cloud."""
|
"""Not supported on LanceDB Cloud."""
|
||||||
return LOOP.run(self._table.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``."""
|
|
||||||
return LOOP.run(self._table.get_lsm_write_spec())
|
|
||||||
|
|
||||||
def close_lsm_writers(self) -> None:
|
def close_lsm_writers(self) -> None:
|
||||||
"""No-op on LanceDB Cloud (no local shard writers)."""
|
"""No-op on LanceDB Cloud (no local shard writers)."""
|
||||||
return LOOP.run(self._table.close_lsm_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"
|
"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:
|
def head(self, n=5) -> pa.Table:
|
||||||
"""
|
"""
|
||||||
Return the first `n` rows of the table.
|
Return the first `n` rows of the table.
|
||||||
|
|||||||
@@ -12,7 +12,6 @@ from .rrf import RRFReranker
|
|||||||
from .mrr import MRRReranker
|
from .mrr import MRRReranker
|
||||||
from .answerdotai import AnswerdotaiRerankers
|
from .answerdotai import AnswerdotaiRerankers
|
||||||
from .voyageai import VoyageAIReranker
|
from .voyageai import VoyageAIReranker
|
||||||
from .watsonx import WatsonxReranker
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"Reranker",
|
"Reranker",
|
||||||
@@ -26,5 +25,4 @@ __all__ = [
|
|||||||
"AnswerdotaiRerankers",
|
"AnswerdotaiRerankers",
|
||||||
"VoyageAIReranker",
|
"VoyageAIReranker",
|
||||||
"MRRReranker",
|
"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
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
|
|
||||||
"""Schema helpers for Lance blob columns."""
|
"""Schema related utilities."""
|
||||||
|
|
||||||
import pyarrow as pa
|
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:
|
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||||
"""A help function to create a vector type.
|
"""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 lancedb.scannable import _register_optional_converters, to_scannable
|
||||||
|
|
||||||
from . import __version__
|
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.arrow import peek_reader
|
||||||
from lancedb.background_loop import LOOP, embedding_executor
|
from lancedb.background_loop import LOOP, embedding_executor
|
||||||
from .dependencies import (
|
from .dependencies import (
|
||||||
@@ -96,7 +88,10 @@ from .util import (
|
|||||||
value_to_sql,
|
value_to_sql,
|
||||||
)
|
)
|
||||||
from .index import lang_mapping
|
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(
|
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
|
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_PREFIXES = ("jieba", "lindera")
|
||||||
_MODEL_BACKED_TOKENIZER_ERRORS = (
|
_MODEL_BACKED_TOKENIZER_ERRORS = (
|
||||||
"unknown base tokenizer",
|
"unknown base tokenizer",
|
||||||
@@ -173,7 +185,6 @@ if TYPE_CHECKING:
|
|||||||
UpdateFieldMetadataResult,
|
UpdateFieldMetadataResult,
|
||||||
DeleteResult,
|
DeleteResult,
|
||||||
DropColumnsResult,
|
DropColumnsResult,
|
||||||
FtsToken,
|
|
||||||
LsmWriteSpec,
|
LsmWriteSpec,
|
||||||
MergeResult,
|
MergeResult,
|
||||||
UpdateResult,
|
UpdateResult,
|
||||||
@@ -640,16 +651,6 @@ def _append_vector_columns(
|
|||||||
col_data = func.compute_source_embeddings_with_retry(
|
col_data = func.compute_source_embeddings_with_retry(
|
||||||
batch[conf.source_column]
|
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:
|
if no_vector_column:
|
||||||
batch = batch.append_column(
|
batch = batch.append_column(
|
||||||
schema.field(vector_column),
|
schema.field(vector_column),
|
||||||
@@ -1148,8 +1149,6 @@ class Table(ABC):
|
|||||||
- "whitespace": Split text by whitespace, but not punctuation.
|
- "whitespace": Split text by whitespace, but not punctuation.
|
||||||
- "raw": No tokenization. The entire text is treated as a single token.
|
- "raw": No tokenization. The entire text is treated as a single token.
|
||||||
- "ngram": N-Gram tokenizer.
|
- "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.
|
- "jieba/*": Jieba tokenizer loaded from Lance's language model home.
|
||||||
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
- "lindera/*": Lindera tokenizer loaded from Lance's language model home.
|
||||||
language : str, default "English"
|
language : str, default "English"
|
||||||
@@ -1514,31 +1513,6 @@ class Table(ABC):
|
|||||||
A query object that can be executed to get the rows.
|
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
|
@abstractmethod
|
||||||
def _execute_query(
|
def _execute_query(
|
||||||
self,
|
self,
|
||||||
@@ -1802,24 +1776,6 @@ class Table(ABC):
|
|||||||
[Table.create_index][lancedb.table.Table.create_index]
|
[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
|
@abstractmethod
|
||||||
def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
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:
|
def take_row_ids(self, row_ids: list[int]) -> LanceTakeQueryBuilder:
|
||||||
return LanceTakeQueryBuilder(self._table.take_row_ids(row_ids))
|
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
|
@property
|
||||||
def tags(self) -> Tags:
|
def tags(self) -> Tags:
|
||||||
"""Tag management for the table.
|
"""Tag management for the table.
|
||||||
@@ -2446,14 +2389,9 @@ class LanceTable(Table):
|
|||||||
-------
|
-------
|
||||||
pd.DataFrame
|
pd.DataFrame
|
||||||
"""
|
"""
|
||||||
validate_blob_mode(blob_mode)
|
_validate_blob_mode(blob_mode)
|
||||||
if blob_mode == "descriptions" or not schema_has_blob_field(self.schema):
|
if blob_mode == "descriptions" or not _schema_has_blob_field(self.schema):
|
||||||
arrow_tbl = self.to_arrow()
|
return self.to_arrow().to_pandas(**kwargs)
|
||||||
if blob_mode == "descriptions":
|
|
||||||
arrow_tbl = strip_auto_row_ids(
|
|
||||||
arrow_tbl, blob_v2_column_paths(self.schema)
|
|
||||||
)
|
|
||||||
return arrow_tbl.to_pandas(**kwargs)
|
|
||||||
|
|
||||||
if (
|
if (
|
||||||
blob_mode == "lazy"
|
blob_mode == "lazy"
|
||||||
@@ -2462,9 +2400,6 @@ class LanceTable(Table):
|
|||||||
):
|
):
|
||||||
return self.to_arrow().to_pandas(**kwargs)
|
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)
|
return self.to_lance().to_pandas(blob_mode=blob_mode, **kwargs)
|
||||||
|
|
||||||
def to_arrow(self) -> pa.Table:
|
def to_arrow(self) -> pa.Table:
|
||||||
@@ -3765,26 +3700,6 @@ class LanceTable(Table):
|
|||||||
"""
|
"""
|
||||||
return LOOP.run(self._table.list_indices())
|
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]:
|
def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
||||||
"""
|
"""
|
||||||
Retrieve statistics about an index
|
Retrieve statistics about an index
|
||||||
@@ -3834,11 +3749,6 @@ class LanceTable(Table):
|
|||||||
[`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
|
[`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
|
||||||
return LOOP.run(self._table.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:
|
def close_lsm_writers(self) -> None:
|
||||||
"""Close cached MemWAL shard writers. See
|
"""Close cached MemWAL shard writers. See
|
||||||
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
|
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
|
||||||
@@ -4110,16 +4020,7 @@ def _handle_bad_vector_column(
|
|||||||
dim = _infer_vector_dim(vec_arr)
|
dim = _infer_vector_dim(vec_arr)
|
||||||
if dim is None:
|
if dim is None:
|
||||||
return data
|
return data
|
||||||
|
has_wrong_dim = pc.not_equal(pc.list_value_length(vec_arr), dim)
|
||||||
is_null = pc.is_null(vec_arr)
|
|
||||||
# pc.list_value_length returns null for null list entries, so
|
|
||||||
# pc.not_equal(null, dim) also returns null. Use or_kleene so that
|
|
||||||
# True OR null = True (Kleene three-valued logic), ensuring null vectors
|
|
||||||
# are counted as wrong-dim.
|
|
||||||
has_wrong_dim = pc.or_kleene(
|
|
||||||
is_null,
|
|
||||||
pc.not_equal(pc.list_value_length(vec_arr), dim),
|
|
||||||
)
|
|
||||||
|
|
||||||
has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
|
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(
|
raise ValueError(
|
||||||
"`fill_value` must not be None if `on_bad_vectors` is 'fill'"
|
"`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:
|
else:
|
||||||
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}")
|
||||||
|
|
||||||
return data.set_column(position, vector_column_name, vec_arr)
|
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:
|
def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray:
|
||||||
if isinstance(arr, pa.ChunkedArray):
|
if isinstance(arr, pa.ChunkedArray):
|
||||||
values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
|
values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks])
|
||||||
@@ -4538,17 +4398,6 @@ class AsyncTable:
|
|||||||
"""
|
"""
|
||||||
await self._inner.unset_lsm_write_spec()
|
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:
|
async def close_lsm_writers(self) -> None:
|
||||||
"""Drain and close any cached MemWAL shard writers for this table.
|
"""Drain and close any cached MemWAL shard writers for this table.
|
||||||
|
|
||||||
@@ -4655,18 +4504,14 @@ class AsyncTable:
|
|||||||
-------
|
-------
|
||||||
pd.DataFrame
|
pd.DataFrame
|
||||||
"""
|
"""
|
||||||
validate_blob_mode(blob_mode)
|
_validate_blob_mode(blob_mode)
|
||||||
schema = await self.schema()
|
if blob_mode == "descriptions" or not _schema_has_blob_field(
|
||||||
if blob_mode == "descriptions" or not schema_has_blob_field(schema):
|
await self.schema()
|
||||||
arrow_tbl = await self.to_arrow()
|
):
|
||||||
if blob_mode == "descriptions":
|
return (await self.to_arrow()).to_pandas(**kwargs)
|
||||||
arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema))
|
|
||||||
return arrow_tbl.to_pandas(**kwargs)
|
|
||||||
|
|
||||||
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
|
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
|
||||||
return (await self.to_arrow()).to_pandas(**kwargs)
|
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)
|
return (await self._to_lance()).to_pandas(blob_mode=blob_mode, **kwargs)
|
||||||
|
|
||||||
async def to_arrow(self) -> pa.Table:
|
async def to_arrow(self) -> pa.Table:
|
||||||
@@ -5767,24 +5612,6 @@ class AsyncTable:
|
|||||||
"""
|
"""
|
||||||
return AsyncTakeQuery(self._inner.take_row_ids(row_ids), self)
|
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
|
@property
|
||||||
def tags(self) -> AsyncTags:
|
def tags(self) -> AsyncTags:
|
||||||
"""Tag management for the dataset.
|
"""Tag management for the dataset.
|
||||||
@@ -5886,24 +5713,6 @@ class AsyncTable:
|
|||||||
"""
|
"""
|
||||||
return await self._inner.list_indices()
|
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]:
|
async def index_stats(self, index_name: str) -> Optional[IndexStatistics]:
|
||||||
"""
|
"""
|
||||||
Retrieve statistics about an index
|
Retrieve statistics about an index
|
||||||
|
|||||||
@@ -1,24 +1,11 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
# SPDX-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
from __future__ import annotations
|
from typing import Literal
|
||||||
|
|
||||||
from typing import Dict, List, Literal, Optional, Tuple, Union
|
|
||||||
|
|
||||||
from .expr import Expr
|
|
||||||
|
|
||||||
# Query type literals
|
# Query type literals
|
||||||
QueryType = Literal["vector", "fts", "hybrid", "auto"]
|
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
|
# Distance type literals
|
||||||
DistanceType = Literal["l2", "cosine", "dot"]
|
DistanceType = Literal["l2", "cosine", "dot"]
|
||||||
DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"]
|
DistanceTypeWithHamming = Literal["l2", "cosine", "dot", "hamming"]
|
||||||
@@ -55,7 +42,5 @@ IndexType = Literal[
|
|||||||
]
|
]
|
||||||
|
|
||||||
# Tokenizer literals
|
# Tokenizer literals
|
||||||
BuiltinTokenizerType = Literal[
|
BuiltinTokenizerType = Literal["simple", "raw", "whitespace", "ngram"]
|
||||||
"simple", "raw", "whitespace", "ngram", "icu", "icu/split"
|
|
||||||
]
|
|
||||||
BaseTokenizerType = BuiltinTokenizerType | str
|
BaseTokenizerType = BuiltinTokenizerType | str
|
||||||
|
|||||||
@@ -177,10 +177,7 @@ def flatten_columns(tbl: pa.Table, flatten: Optional[Union[int, bool]] = None):
|
|||||||
continue
|
continue
|
||||||
else:
|
else:
|
||||||
break
|
break
|
||||||
# `bool` is a subclass of `int`, so guard against it explicitly: `flatten=False`
|
elif isinstance(flatten, int):
|
||||||
# (and `None`) must mean "do not flatten" rather than falling into the integer
|
|
||||||
# branch and raising on the `flatten <= 0` check.
|
|
||||||
elif isinstance(flatten, int) and not isinstance(flatten, bool):
|
|
||||||
if flatten <= 0:
|
if flatten <= 0:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"Please specify a positive integer for flatten or the boolean "
|
"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))
|
stacked = torch.tensor(numpy.column_stack(columns))
|
||||||
rows = list(stacked.unbind(dim=0))
|
rows = list(stacked.unbind(dim=0))
|
||||||
return rows
|
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
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
import os
|
import os
|
||||||
import pickle
|
|
||||||
from typing import List, Optional, Union
|
from typing import List, Optional, Union
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
@@ -243,49 +242,6 @@ def test_embedding_with_bad_results(tmp_path):
|
|||||||
assert tbl["vector"].null_count == 1
|
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):
|
def test_with_existing_vectors(tmp_path):
|
||||||
@register("mock-embedding")
|
@register("mock-embedding")
|
||||||
class MockEmbeddingFunction(TextEmbeddingFunction):
|
class MockEmbeddingFunction(TextEmbeddingFunction):
|
||||||
@@ -592,22 +548,6 @@ def test_openai_no_retry_on_401(mock_sleep):
|
|||||||
assert mock_sleep.call_count == 0
|
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():
|
def test_url_retrieve_downloads_image():
|
||||||
"""
|
"""
|
||||||
Embedding functions like open-clip, siglip, and jinaai use url_retrieve()
|
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
|
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):
|
def test_fts_on_list(mem_db: DBConnection):
|
||||||
data = pa.table(
|
data = pa.table(
|
||||||
{
|
{
|
||||||
@@ -1175,84 +1084,6 @@ def test_fts_query_to_json():
|
|||||||
assert json_str == expected
|
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):
|
def test_fts_fast_search(table):
|
||||||
table.create_fts_index("text")
|
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-License-Identifier: Apache-2.0
|
||||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
|
|
||||||
from unittest import mock
|
|
||||||
|
|
||||||
import lancedb
|
import lancedb
|
||||||
|
|
||||||
from lancedb.query import LanceHybridQueryBuilder
|
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
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_hybrid_query_distance_range_async(table: AsyncTable):
|
async def test_hybrid_query_distance_range_async(table: AsyncTable):
|
||||||
reranker = RRFReranker(return_score="all")
|
reranker = RRFReranker(return_score="all")
|
||||||
@@ -193,23 +177,6 @@ async def test_analyze_plan(table: AsyncTable):
|
|||||||
assert "metrics=" in res
|
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
|
@pytest.fixture
|
||||||
def table_with_id(tmpdir_factory) -> Table:
|
def table_with_id(tmpdir_factory) -> Table:
|
||||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
tmp_path = str(tmpdir_factory.mktemp("data"))
|
||||||
|
|||||||
@@ -11,7 +11,6 @@ import lancedb
|
|||||||
import pyarrow as pa
|
import pyarrow as pa
|
||||||
import pytest
|
import pytest
|
||||||
from lancedb._lancedb import LsmWriteSpec
|
from lancedb._lancedb import LsmWriteSpec
|
||||||
from lancedb.index import BTree
|
|
||||||
|
|
||||||
SCHEMA = pa.schema(
|
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"})
|
s = s.with_writer_config_defaults({"durable_write": "false"})
|
||||||
assert s.writer_config_defaults == {"durable_write": "false"}
|
assert s.writer_config_defaults == {"durable_write": "false"}
|
||||||
assert "durable_write" in repr(s)
|
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
|
|
||||||
|
|||||||
@@ -935,41 +935,14 @@ def test_transform_fn(mem_db):
|
|||||||
try:
|
try:
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
# "torch" returns a list of per-row dicts. Default DataLoader collate
|
torch_result = list(
|
||||||
# stacks the per-row dicts back into a dict of batched tensors.
|
permutation.with_format("torch").iter(10, skip_last_batch=False)
|
||||||
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)
|
|
||||||
)[0]
|
)[0]
|
||||||
assert isinstance(torch_rows, list)
|
assert isinstance(torch_result, list)
|
||||||
assert len(torch_rows) == 10
|
assert len(torch_result) == 10
|
||||||
assert isinstance(torch_rows[0], torch.Tensor)
|
assert isinstance(torch_result[0], torch.Tensor)
|
||||||
assert torch_rows[0].shape == (2,)
|
assert torch_result[0].shape == (2,)
|
||||||
assert torch_rows[0].dtype == torch.int64
|
assert torch_result[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
|
|
||||||
except ImportError:
|
except ImportError:
|
||||||
# Skip check if torch is not installed
|
# Skip check if torch is not installed
|
||||||
pass
|
pass
|
||||||
|
|||||||
@@ -11,7 +11,6 @@ import lancedb
|
|||||||
from lancedb.db import AsyncConnection
|
from lancedb.db import AsyncConnection
|
||||||
from lancedb.embeddings.base import TextEmbeddingFunction
|
from lancedb.embeddings.base import TextEmbeddingFunction
|
||||||
from lancedb.embeddings.registry import get_registry, register
|
from lancedb.embeddings.registry import get_registry, register
|
||||||
from lancedb.expr import col
|
|
||||||
from lancedb.index import FTS, IvfPq
|
from lancedb.index import FTS, IvfPq
|
||||||
import lancedb.pydantic
|
import lancedb.pydantic
|
||||||
import numpy as np
|
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):
|
def _assert_lazy_blob(value, expected: bytes):
|
||||||
assert hasattr(value, "readall")
|
assert hasattr(value, "readall")
|
||||||
assert value.readall() == expected
|
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")
|
@pytest.fixture(scope="module")
|
||||||
def table(tmpdir_factory) -> lancedb.table.Table:
|
def table(tmpdir_factory) -> lancedb.table.Table:
|
||||||
tmp_path = str(tmpdir_factory.mktemp("data"))
|
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")
|
assert not hasattr(first, "readall")
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
def test_plain_scan_query_to_pandas_blob_projection(tmp_db):
|
||||||
def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
|
||||||
pytest.importorskip("lance")
|
pytest.importorskip("lance")
|
||||||
table = _blob_query_table(
|
table = tmp_db.create_table(
|
||||||
tmp_db, f"test_query_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
"test_query_to_pandas_blob_projection", _blob_query_data()
|
||||||
)
|
)
|
||||||
|
|
||||||
df = (
|
df = (
|
||||||
@@ -312,8 +250,9 @@ def test_plain_scan_query_to_pandas_blob_bytes_projection(tmp_db, blob_schema):
|
|||||||
.to_pandas(blob_mode="bytes")
|
.to_pandas(blob_mode="bytes")
|
||||||
)
|
)
|
||||||
|
|
||||||
_assert_blob_bytes_projection(df)
|
assert df["id_alias"].tolist() == [3, 4]
|
||||||
assert "_rowid" not in df.columns
|
assert df["payload"].tolist() == [b"three", b"four"]
|
||||||
|
assert df["double_id"].tolist() == [6, 8]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
@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_df["id"].tolist() == [1]
|
||||||
_assert_lazy_blob(lazy_df["blob"].iloc[0], b"one")
|
_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 (
|
desc_df = await (
|
||||||
table.query()
|
table.query()
|
||||||
.where("id = 1")
|
.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")
|
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.asyncio
|
||||||
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
@pytest.mark.parametrize("blob_mode", ["bytes", "descriptions"])
|
||||||
async def test_async_plain_scan_query_to_pandas_blob_mode_does_not_collect_arrow(
|
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]
|
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):
|
def test_where_repeated_combines_with_and(table: lancedb.table.Table):
|
||||||
# Calling where() more than once should AND the filters together instead of
|
# Calling where() more than once should AND the filters together instead of
|
||||||
# silently replacing the previous one (regression test for #2649).
|
# 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)
|
# 2. Fast Search -> Should NOT include "LanceScan" (Uses Index)
|
||||||
plan = table.search(q).fast_search().explain_plan(True)
|
plan = table.search(q).fast_search().explain_plan(True)
|
||||||
assert "LanceScan" not in plan
|
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"))
|
content_len = int(request.headers.get("Content-Length"))
|
||||||
body = json.loads(request.rfile.read(content_len))
|
body = json.loads(request.rfile.read(content_len))
|
||||||
if "filter" in body:
|
if "filter" in body:
|
||||||
match = re.search(
|
match = re.search(r"_rowoffset in \((.*?)\)", body["filter"])
|
||||||
r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
|
offsets = [int(offset.strip()) for offset in match.group(1).split(",")]
|
||||||
)
|
|
||||||
offsets = [int(o.strip()) for o in match.group(1).split(",")]
|
|
||||||
else:
|
else:
|
||||||
offsets = list(range(len(rows)))
|
offsets = rows
|
||||||
table = pa.table({"a": [rows[offset] for offset in offsets]})
|
table = pa.table({"a": [rows[offset] for offset in offsets]})
|
||||||
|
|
||||||
request.send_response(200)
|
request.send_response(200)
|
||||||
|
|||||||
@@ -23,7 +23,6 @@ from lancedb.rerankers import (
|
|||||||
AnswerdotaiRerankers,
|
AnswerdotaiRerankers,
|
||||||
VoyageAIReranker,
|
VoyageAIReranker,
|
||||||
MRRReranker,
|
MRRReranker,
|
||||||
WatsonxReranker,
|
|
||||||
)
|
)
|
||||||
from lancedb.table import LanceTable
|
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 FTS score (rowid 0, {scores[0]:.4f}) should beat "
|
||||||
f"document with no FTS match (rowid 1, {scores[1]:.4f})"
|
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
|
Related issue: https://github.com/lancedb/lancedb/issues/1898
|
||||||
|
|
||||||
These tests validate the early error checking for S3 bucket names with dots.
|
These tests validate the early error checking for S3 bucket names with dots.
|
||||||
When validation succeeds, eager namespace initialization may still fail later
|
No actual S3 connection is made - validation happens before connection.
|
||||||
because these tests intentionally do not provide real S3 credentials.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import pytest
|
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"
|
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:
|
class TestS3BucketWithDotsSync:
|
||||||
"""Tests for connect()."""
|
"""Tests for connect()."""
|
||||||
|
|
||||||
@@ -35,22 +27,19 @@ class TestS3BucketWithDotsSync:
|
|||||||
with pytest.raises(ValueError, match="contains dots"):
|
with pytest.raises(ValueError, match="contains dots"):
|
||||||
lancedb.connect(BUCKET_WITH_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
|
uri, opts = BUCKET_WITH_DOTS_AND_REGION
|
||||||
assert_not_rejected_for_bucket_dots(
|
db = lancedb.connect(uri, storage_options=opts)
|
||||||
lambda: 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
|
uri, opts = BUCKET_WITH_DOTS_AND_AWS_REGION
|
||||||
assert_not_rejected_for_bucket_dots(
|
db = lancedb.connect(uri, storage_options=opts)
|
||||||
lambda: lancedb.connect(uri, storage_options=opts)
|
assert db is not None
|
||||||
)
|
|
||||||
|
|
||||||
def test_bucket_without_dots_is_not_rejected(self):
|
def test_bucket_without_dots_passes(self):
|
||||||
assert_not_rejected_for_bucket_dots(
|
db = lancedb.connect(BUCKET_WITHOUT_DOTS)
|
||||||
lambda: lancedb.connect(BUCKET_WITHOUT_DOTS)
|
assert db is not None
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class TestS3BucketWithDotsAsync:
|
class TestS3BucketWithDotsAsync:
|
||||||
@@ -62,24 +51,18 @@ class TestS3BucketWithDotsAsync:
|
|||||||
await lancedb.connect_async(BUCKET_WITH_DOTS)
|
await lancedb.connect_async(BUCKET_WITH_DOTS)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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
|
uri, opts = BUCKET_WITH_DOTS_AND_REGION
|
||||||
try:
|
db = await lancedb.connect_async(uri, storage_options=opts)
|
||||||
await lancedb.connect_async(uri, storage_options=opts)
|
assert db is not None
|
||||||
except ValueError as err:
|
|
||||||
assert "contains dots" not in str(err)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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
|
uri, opts = BUCKET_WITH_DOTS_AND_AWS_REGION
|
||||||
try:
|
db = await lancedb.connect_async(uri, storage_options=opts)
|
||||||
await lancedb.connect_async(uri, storage_options=opts)
|
assert db is not None
|
||||||
except ValueError as err:
|
|
||||||
assert "contains dots" not in str(err)
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_bucket_without_dots_is_not_rejected(self):
|
async def test_bucket_without_dots_passes(self):
|
||||||
try:
|
db = await lancedb.connect_async(BUCKET_WITHOUT_DOTS)
|
||||||
await lancedb.connect_async(BUCKET_WITHOUT_DOTS)
|
assert db is not None
|
||||||
except ValueError as err:
|
|
||||||
assert "contains dots" not in str(err)
|
|
||||||
|
|||||||
@@ -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):
|
def _assert_lazy_blob(value, expected: bytes):
|
||||||
assert hasattr(value, "readall")
|
assert hasattr(value, "readall")
|
||||||
assert value.readall() == expected
|
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")
|
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):
|
def test_table_to_pandas_kwargs(tmp_db: DBConnection):
|
||||||
pd = pytest.importorskip("pandas")
|
pd = pytest.importorskip("pandas")
|
||||||
data = pa.table({"id": pa.array([1, 2], pa.int64())})
|
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.asyncio
|
||||||
@pytest.mark.parametrize("blob_schema", ["v1", "v2"])
|
async def test_async_table_to_pandas_blob_bytes(tmp_db_async: AsyncConnection):
|
||||||
async def test_async_table_to_pandas_blob_bytes(
|
|
||||||
tmp_db_async: AsyncConnection, blob_schema
|
|
||||||
):
|
|
||||||
pytest.importorskip("lance")
|
pytest.importorskip("lance")
|
||||||
table = await _blob_table_async(
|
table = await tmp_db_async.create_table(
|
||||||
tmp_db_async, f"test_async_to_pandas_blob_{blob_schema}_bytes", blob_schema
|
"test_async_to_pandas_blob_bytes", data=_blob_test_data()
|
||||||
)
|
)
|
||||||
|
|
||||||
df = await table.to_pandas(blob_mode="bytes")
|
df = await table.to_pandas(blob_mode="bytes")
|
||||||
|
|
||||||
assert list(df.columns) == ["id", "blob"]
|
|
||||||
assert df["blob"].tolist() == [b"hello", b"world"]
|
assert df["blob"].tolist() == [b"hello", b"world"]
|
||||||
assert "_rowid" not in df.columns
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -1611,23 +1568,16 @@ def test_create_with_nans(mem_db: DBConnection):
|
|||||||
"fill_test",
|
"fill_test",
|
||||||
data=[
|
data=[
|
||||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
{"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], "item": "bar", "price": 20.0},
|
||||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||||
{"vector": [5], "item": "bar", "price": 22.0},
|
|
||||||
],
|
],
|
||||||
on_bad_vectors="fill",
|
on_bad_vectors="fill",
|
||||||
fill_value=0.0,
|
fill_value=0.0,
|
||||||
)
|
)
|
||||||
assert len(table) == 5
|
assert len(table) == 3
|
||||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||||
filled_vectors = {
|
v = arrow_tbl["vector"].to_pylist()[0]
|
||||||
row["price"]: row["vector"]
|
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
|
||||||
}
|
|
||||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
|
||||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
|
||||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
|
||||||
|
|
||||||
|
|
||||||
def test_add_with_nans(mem_db: DBConnection):
|
def test_add_with_nans(mem_db: DBConnection):
|
||||||
@@ -1670,21 +1620,15 @@ def test_add_with_nans(mem_db: DBConnection):
|
|||||||
data=[
|
data=[
|
||||||
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
|
||||||
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
{"vector": [np.nan], "item": "bar", "price": 20.0},
|
||||||
{"vector": [np.nan, 5.0], "item": "bar", "price": 21.0},
|
{"vector": [np.nan, np.nan], "item": "bar", "price": 20.0},
|
||||||
{"vector": [5], "item": "bar", "price": 22.0},
|
|
||||||
],
|
],
|
||||||
on_bad_vectors="fill",
|
on_bad_vectors="fill",
|
||||||
fill_value=0.0,
|
fill_value=0.0,
|
||||||
)
|
)
|
||||||
assert len(table) == 4
|
assert len(table) == 3
|
||||||
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
arrow_tbl = table.search().where("item == 'bar'").to_arrow()
|
||||||
filled_vectors = {
|
v = arrow_tbl["vector"].to_pylist()[0]
|
||||||
row["price"]: row["vector"]
|
assert np.allclose(v, np.array([0.0, 0.0]))
|
||||||
for row in arrow_tbl.select(["price", "vector"]).to_pylist()
|
|
||||||
}
|
|
||||||
assert np.allclose(filled_vectors[20.0], np.array([0.0, 0.0]))
|
|
||||||
assert np.allclose(filled_vectors[21.0], np.array([0.0, 5.0]))
|
|
||||||
assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0]))
|
|
||||||
|
|
||||||
|
|
||||||
def test_add_with_empty_fixed_size_list_drops_bad_rows(mem_db: DBConnection):
|
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,
|
fill_value=0.0,
|
||||||
)
|
)
|
||||||
|
|
||||||
vector = table.to_arrow()["vector"]
|
assert table.to_arrow()["vector"].to_pylist() == [[1.0, 2.0], [0.0, 0.0]]
|
||||||
assert vector.type == pa.list_(pa.float32())
|
|
||||||
assert vector.to_pylist() == [[1.0, 2.0], [0.0, 3.0]]
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
|
|||||||
@@ -148,47 +148,15 @@ def test_permutation_dataloader(mem_db):
|
|||||||
for batch in dataloader:
|
for batch in dataloader:
|
||||||
assert batch["a"].size(0) == 10
|
assert batch["a"].size(0) == 10
|
||||||
|
|
||||||
# "torch" produces a list of per-row dicts per batch. The default
|
permutation = permutation.with_format("torch")
|
||||||
# DataLoader collate stacks the per-row dicts back into a batched dict.
|
dataloader = torch.utils.data.DataLoader(permutation, batch_size=10, shuffle=True)
|
||||||
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)
|
|
||||||
for batch in dataloader:
|
for batch in dataloader:
|
||||||
assert batch.size(0) == 10
|
assert batch.size(0) == 10
|
||||||
assert batch.size(1) == 1
|
assert batch.size(1) == 1
|
||||||
|
|
||||||
col_perm = permutation.with_format("torch_col")
|
permutation = permutation.with_format("torch_col")
|
||||||
dataloader = torch.utils.data.DataLoader(
|
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:
|
for batch in dataloader:
|
||||||
assert batch.size(0) == 1
|
assert batch.size(0) == 1
|
||||||
|
|||||||
@@ -13,7 +13,6 @@ from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
|||||||
from lancedb.table import (
|
from lancedb.table import (
|
||||||
_append_vector_columns,
|
_append_vector_columns,
|
||||||
_cast_to_target_schema,
|
_cast_to_target_schema,
|
||||||
_fill_bad_vector_values,
|
|
||||||
_handle_bad_vectors,
|
_handle_bad_vectors,
|
||||||
_into_pyarrow_reader,
|
_into_pyarrow_reader,
|
||||||
_infer_target_schema,
|
_infer_target_schema,
|
||||||
@@ -26,42 +25,10 @@ import pandas as pd
|
|||||||
import polars as pl
|
import polars as pl
|
||||||
import pytest
|
import pytest
|
||||||
import lancedb
|
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
|
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():
|
def test_normalize_uri():
|
||||||
uris = [
|
uris = [
|
||||||
"relative/path",
|
"relative/path",
|
||||||
@@ -288,9 +255,7 @@ def test_append_vector_columns():
|
|||||||
|
|
||||||
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||||
def test_handle_bad_vectors_jagged(on_bad_vectors):
|
def test_handle_bad_vectors_jagged(on_bad_vectors):
|
||||||
vector = pa.array(
|
vector = pa.array([[1.0, 2.0], [3.0], [4.0, 5.0]])
|
||||||
[[1.0, 2.0], [3.0], [4.0, 5.0], [6.0, 7.0, 8.0], [None, 9.0], None]
|
|
||||||
)
|
|
||||||
schema = pa.schema({"vector": pa.list_(pa.float64())})
|
schema = pa.schema({"vector": pa.list_(pa.float64())})
|
||||||
data = pa.table({"vector": vector}, schema=schema)
|
data = pa.table({"vector": vector}, schema=schema)
|
||||||
|
|
||||||
@@ -316,54 +281,15 @@ def test_handle_bad_vectors_jagged(on_bad_vectors):
|
|||||||
).read_all()
|
).read_all()
|
||||||
|
|
||||||
if on_bad_vectors == "drop":
|
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":
|
elif on_bad_vectors == "fill":
|
||||||
expected = pa.array(
|
expected = pa.array([[1.0, 2.0], [42.0, 42.0], [4.0, 5.0]])
|
||||||
[
|
|
||||||
[1.0, 2.0],
|
|
||||||
[3.0, 42.0],
|
|
||||||
[4.0, 5.0],
|
|
||||||
[6.0, 7.0],
|
|
||||||
[None, 9.0],
|
|
||||||
[42.0, 42.0],
|
|
||||||
]
|
|
||||||
)
|
|
||||||
elif on_bad_vectors == "null":
|
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
|
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"])
|
@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"])
|
||||||
def test_handle_bad_vectors_nan(on_bad_vectors):
|
def test_handle_bad_vectors_nan(on_bad_vectors):
|
||||||
vector = pa.array([[1.0, float("nan")], [3.0, 4.0]])
|
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":
|
if on_bad_vectors == "drop":
|
||||||
expected = pa.array([[3.0, 4.0]])
|
expected = pa.array([[3.0, 4.0]])
|
||||||
elif on_bad_vectors == "fill":
|
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":
|
elif on_bad_vectors == "null":
|
||||||
expected = pa.array([None, [3.0, 4.0]])
|
expected = pa.array([None, [3.0, 4.0]])
|
||||||
|
|
||||||
|
|||||||
+8
-60
@@ -7,14 +7,12 @@
|
|||||||
//! build type-safe filter / projection expressions that map directly to
|
//! build type-safe filter / projection expressions that map directly to
|
||||||
//! DataFusion [`Expr`] nodes, bypassing SQL string parsing.
|
//! DataFusion [`Expr`] nodes, bypassing SQL string parsing.
|
||||||
|
|
||||||
use std::ops::{Add, Div, Mul, Not, Sub};
|
|
||||||
|
|
||||||
use arrow::{datatypes::DataType, pyarrow::PyArrowType};
|
use arrow::{datatypes::DataType, pyarrow::PyArrowType};
|
||||||
use datafusion_common::ScalarValue;
|
use datafusion_common::ScalarValue;
|
||||||
use lancedb::expr::{
|
use lancedb::expr::{
|
||||||
DfExpr, col as ldb_col, contains, expr_cast, is_in, lit as df_lit, lower, upper,
|
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};
|
use pyo3::{Bound, PyAny, PyResult, exceptions::PyValueError, prelude::*, pyfunction};
|
||||||
|
|
||||||
/// A type-safe DataFusion expression.
|
/// A type-safe DataFusion expression.
|
||||||
@@ -65,30 +63,30 @@ impl PyExpr {
|
|||||||
Self(self.0.clone().or(other.0.clone()))
|
Self(self.0.clone().or(other.0.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Logical NOT.
|
|
||||||
fn not_(&self) -> Self {
|
fn not_(&self) -> Self {
|
||||||
|
use std::ops::Not;
|
||||||
Self(self.0.clone().not())
|
Self(self.0.clone().not())
|
||||||
}
|
}
|
||||||
|
|
||||||
// ── arithmetic ───────────────────────────────────────────────────────────
|
// ── arithmetic ───────────────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Add expressions.
|
|
||||||
fn add(&self, other: &Self) -> Self {
|
fn add(&self, other: &Self) -> Self {
|
||||||
|
use std::ops::Add;
|
||||||
Self(self.0.clone().add(other.0.clone()))
|
Self(self.0.clone().add(other.0.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Subtract expressions.
|
|
||||||
fn sub(&self, other: &Self) -> Self {
|
fn sub(&self, other: &Self) -> Self {
|
||||||
|
use std::ops::Sub;
|
||||||
Self(self.0.clone().sub(other.0.clone()))
|
Self(self.0.clone().sub(other.0.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Multiply expressions.
|
|
||||||
fn mul(&self, other: &Self) -> Self {
|
fn mul(&self, other: &Self) -> Self {
|
||||||
|
use std::ops::Mul;
|
||||||
Self(self.0.clone().mul(other.0.clone()))
|
Self(self.0.clone().mul(other.0.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Divide expressions.
|
|
||||||
fn div(&self, other: &Self) -> Self {
|
fn div(&self, other: &Self) -> Self {
|
||||||
|
use std::ops::Div;
|
||||||
Self(self.0.clone().div(other.0.clone()))
|
Self(self.0.clone().div(other.0.clone()))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -155,8 +153,7 @@ pub fn expr_col(name: &str) -> PyExpr {
|
|||||||
|
|
||||||
/// Create a literal value expression.
|
/// Create a literal value expression.
|
||||||
///
|
///
|
||||||
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`, `date`,
|
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`.
|
||||||
/// `datetime`, `Decimal`.
|
|
||||||
#[pyfunction]
|
#[pyfunction]
|
||||||
pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||||
// bool must be checked before int because bool is a subclass of int in Python
|
// 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>() {
|
if let Ok(i) = value.extract::<i64>() {
|
||||||
return Ok(PyExpr(df_lit(i)));
|
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>() {
|
if let Ok(f) = value.extract::<f64>() {
|
||||||
return Ok(PyExpr(df_lit(f)));
|
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>>()?;
|
let bytes = value.extract::<Vec<u8>>()?;
|
||||||
return Ok(PyExpr(df_lit(ScalarValue::Binary(Some(bytes)))));
|
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!(
|
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()?
|
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.
|
/// Call an arbitrary registered SQL function by name.
|
||||||
///
|
///
|
||||||
/// See `lancedb::expr::func` for the list of supported function names.
|
/// 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 query::{FTSQuery, HybridQuery, Query, VectorQuery};
|
||||||
use session::Session;
|
use session::Session;
|
||||||
use table::{
|
use table::{
|
||||||
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken,
|
AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, LsmWriteSpec,
|
||||||
LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult,
|
MergeResult, Table, UpdateFieldMetadataResult, UpdateResult,
|
||||||
};
|
};
|
||||||
|
|
||||||
pub mod arrow;
|
pub mod arrow;
|
||||||
@@ -27,7 +27,6 @@ pub mod header;
|
|||||||
pub mod index;
|
pub mod index;
|
||||||
pub mod namespace;
|
pub mod namespace;
|
||||||
pub mod oauth;
|
pub mod oauth;
|
||||||
pub mod otel;
|
|
||||||
pub mod permutation;
|
pub mod permutation;
|
||||||
pub mod query;
|
pub mod query;
|
||||||
pub mod runtime;
|
pub mod runtime;
|
||||||
@@ -44,7 +43,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
m.add_class::<Connection>()?;
|
m.add_class::<Connection>()?;
|
||||||
m.add_class::<Session>()?;
|
m.add_class::<Session>()?;
|
||||||
m.add_class::<Table>()?;
|
m.add_class::<Table>()?;
|
||||||
m.add_class::<PyBlobFile>()?;
|
|
||||||
m.add_class::<IndexConfig>()?;
|
m.add_class::<IndexConfig>()?;
|
||||||
m.add_class::<Query>()?;
|
m.add_class::<Query>()?;
|
||||||
m.add_class::<FTSQuery>()?;
|
m.add_class::<FTSQuery>()?;
|
||||||
@@ -60,23 +58,12 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
|||||||
m.add_class::<DeleteResult>()?;
|
m.add_class::<DeleteResult>()?;
|
||||||
m.add_class::<DropColumnsResult>()?;
|
m.add_class::<DropColumnsResult>()?;
|
||||||
m.add_class::<UpdateResult>()?;
|
m.add_class::<UpdateResult>()?;
|
||||||
m.add_class::<FtsToken>()?;
|
|
||||||
m.add_class::<PyAsyncPermutationBuilder>()?;
|
m.add_class::<PyAsyncPermutationBuilder>()?;
|
||||||
m.add_class::<PyPermutationReader>()?;
|
m.add_class::<PyPermutationReader>()?;
|
||||||
m.add_class::<PyExpr>()?;
|
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, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(connect_namespace_client, 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!(permutation::async_permutation_builder, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
||||||
m.add_function(wrap_pyfunction!(query::fts_query_to_json, 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]
|
#[pymethods]
|
||||||
impl PyAsyncPermutationBuilder {
|
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(
|
pub fn split_random(
|
||||||
slf: PyRefMut<'_, Self>,
|
slf: PyRefMut<'_, Self>,
|
||||||
ratios: Option<Vec<f64>>,
|
ratios: Option<Vec<f64>>,
|
||||||
counts: Option<Vec<u64>>,
|
counts: Option<Vec<u64>>,
|
||||||
fixed: Option<u64>,
|
fixed: Option<u64>,
|
||||||
seed: Option<u64>,
|
seed: Option<u64>,
|
||||||
clump_size: Option<u64>,
|
|
||||||
split_names: Option<Vec<String>>,
|
split_names: Option<Vec<String>>,
|
||||||
) -> PyResult<Self> {
|
) -> PyResult<Self> {
|
||||||
// Check that exactly one split type is provided
|
// Check that exactly one split type is provided
|
||||||
@@ -112,14 +111,7 @@ impl PyAsyncPermutationBuilder {
|
|||||||
};
|
};
|
||||||
|
|
||||||
slf.modify(|builder| {
|
slf.modify(|builder| {
|
||||||
builder.with_split_strategy(
|
builder.with_split_strategy(SplitStrategy::Random { seed, sizes }, split_names)
|
||||||
SplitStrategy::Random {
|
|
||||||
seed,
|
|
||||||
sizes,
|
|
||||||
clump_size,
|
|
||||||
},
|
|
||||||
split_names,
|
|
||||||
)
|
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+5
-239
@@ -2,7 +2,7 @@
|
|||||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||||
use std::{collections::HashMap, sync::Arc};
|
use std::{collections::HashMap, sync::Arc};
|
||||||
|
|
||||||
use crate::runtime::{block_on, future_into_py};
|
use crate::runtime::future_into_py;
|
||||||
use crate::{
|
use crate::{
|
||||||
connection::Connection,
|
connection::Connection,
|
||||||
error::PythonErrorExt,
|
error::PythonErrorExt,
|
||||||
@@ -12,23 +12,19 @@ use crate::{
|
|||||||
table::scannable::PyScannable,
|
table::scannable::PyScannable,
|
||||||
};
|
};
|
||||||
use arrow::{
|
use arrow::{
|
||||||
array::{Array, LargeBinaryArray},
|
|
||||||
datatypes::{DataType, Schema},
|
datatypes::{DataType, Schema},
|
||||||
ffi_stream::ArrowArrayStreamReader,
|
ffi_stream::ArrowArrayStreamReader,
|
||||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||||
};
|
};
|
||||||
use lancedb::blob::BlobFile;
|
|
||||||
use lancedb::index::scalar::FtsIndexBuilder;
|
|
||||||
use lancedb::table::{
|
use lancedb::table::{
|
||||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken,
|
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform,
|
||||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||||
};
|
};
|
||||||
use lancedb::tokenize as lancedb_tokenize;
|
|
||||||
use pyo3::{
|
use pyo3::{
|
||||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||||
exceptions::{PyRuntimeError, PyValueError},
|
exceptions::{PyRuntimeError, PyValueError},
|
||||||
pyclass, pyfunction, pymethods,
|
pyclass, pymethods,
|
||||||
types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods},
|
types::{IntoPyDict, PyAnyMethods, PyDict, PyDictMethods},
|
||||||
};
|
};
|
||||||
|
|
||||||
mod scannable;
|
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)]
|
#[pyclass(get_all, from_py_object)]
|
||||||
#[derive(Clone, Debug)]
|
#[derive(Clone, Debug)]
|
||||||
pub struct AddColumnsResult {
|
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]
|
#[pyclass]
|
||||||
pub struct Table {
|
pub struct Table {
|
||||||
// We keep a copy of the name to use if the inner table is dropped
|
// 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>> {
|
pub fn index_stats(self_: PyRef<'_, Self>, index_name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||||
let inner = self_.inner_ref()?.clone();
|
let inner = self_.inner_ref()?.clone();
|
||||||
future_into_py(self_.py(), async move {
|
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.
|
/// Optimize the on-disk data by compacting and pruning old data, for better performance.
|
||||||
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
#[pyo3(signature = (cleanup_since_ms=None, delete_unverified=None))]
|
||||||
pub fn optimize(
|
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>> {
|
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||||
let inner = self_.inner_ref()?.clone();
|
let inner = self_.inner_ref()?.clone();
|
||||||
future_into_py(self_.py(), async move {
|
future_into_py(self_.py(), async move {
|
||||||
|
|||||||
@@ -3,14 +3,10 @@
|
|||||||
|
|
||||||
"""Tests for the type-safe expression builder API."""
|
"""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 pytest
|
||||||
|
import pyarrow as pa
|
||||||
import lancedb
|
import lancedb
|
||||||
from lancedb.expr import Expr, col, func, lit
|
from lancedb.expr import Expr, col, lit, func
|
||||||
|
|
||||||
|
|
||||||
# ── unit tests for Expr construction ─────────────────────────────────────────
|
# ── unit tests for Expr construction ─────────────────────────────────────────
|
||||||
@@ -58,28 +54,6 @@ class TestExprConstruction:
|
|||||||
with pytest.raises(Exception):
|
with pytest.raises(Exception):
|
||||||
func("not_a_real_function", col("x"))
|
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:
|
class TestExprOperators:
|
||||||
def test_eq_operator(self):
|
def test_eq_operator(self):
|
||||||
@@ -168,20 +142,6 @@ class TestExprOperators:
|
|||||||
assert isinstance(e, Expr)
|
assert isinstance(e, Expr)
|
||||||
assert e.to_sql() == "(name = 'alice')"
|
assert e.to_sql() == "(name = 'alice')"
|
||||||
|
|
||||||
def test_reflexive_comparisons(self):
|
|
||||||
# 10 < col("age") swaps to col("age") > 10
|
|
||||||
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:
|
class TestExprBytesLiteral:
|
||||||
def test_bytes_to_sql(self):
|
def test_bytes_to_sql(self):
|
||||||
@@ -322,40 +282,6 @@ class TestExprRepr:
|
|||||||
{e: 1}
|
{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 ─────────────────
|
# ── 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"]
|
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 ───────────────────────────────────
|
# ── 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"
|
version = "13.3.1"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
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 = [
|
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" },
|
{ 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]
|
[package.optional-dependencies]
|
||||||
cublas = [
|
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 = [
|
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 = [
|
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 = [
|
cufile = [
|
||||||
{ name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
|
{ name = "nvidia-cufile", marker = "sys_platform == 'linux'" },
|
||||||
]
|
]
|
||||||
cupti = [
|
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 = [
|
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 = [
|
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 = [
|
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 = [
|
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 = [
|
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 = [
|
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]]
|
[[package]]
|
||||||
@@ -1909,9 +1909,6 @@ embeddings = [
|
|||||||
{ name = "sentencepiece" },
|
{ name = "sentencepiece" },
|
||||||
{ name = "torch" },
|
{ name = "torch" },
|
||||||
]
|
]
|
||||||
otel = [
|
|
||||||
{ name = "opentelemetry-api" },
|
|
||||||
]
|
|
||||||
pylance = [
|
pylance = [
|
||||||
{ name = "pylance" },
|
{ name = "pylance" },
|
||||||
]
|
]
|
||||||
@@ -1926,7 +1923,6 @@ tests = [
|
|||||||
{ name = "boto3" },
|
{ name = "boto3" },
|
||||||
{ name = "datafusion" },
|
{ name = "datafusion" },
|
||||||
{ name = "duckdb" },
|
{ 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.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 = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" },
|
||||||
{ name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" },
|
{ name = "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 == 'clip'" },
|
||||||
{ name = "open-clip-torch", marker = "extra == 'embeddings'", specifier = ">=2.20.0" },
|
{ name = "open-clip-torch", marker = "extra == 'embeddings'", specifier = ">=2.20.0" },
|
||||||
{ name = "openai", marker = "extra == 'embeddings'", specifier = ">=1.6.1" },
|
{ 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 = "overrides", marker = "python_full_version < '3.12'", specifier = ">=0.7" },
|
||||||
{ name = "packaging", specifier = ">=23.0" },
|
{ name = "packaging", specifier = ">=23.0" },
|
||||||
{ name = "pandas", marker = "extra == 'tests'", specifier = ">=1.4" },
|
{ name = "pandas", marker = "extra == 'tests'", specifier = ">=1.4" },
|
||||||
@@ -2000,7 +1994,7 @@ requires-dist = [
|
|||||||
{ name = "transformers", marker = "extra == 'siglip'", specifier = ">=4.41.0" },
|
{ 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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "lomond"
|
name = "lomond"
|
||||||
@@ -2781,7 +2775,7 @@ name = "nvidia-cudnn-cu13"
|
|||||||
version = "9.19.0.56"
|
version = "9.19.0.56"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
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 = [
|
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" },
|
{ 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"
|
version = "12.0.0.61"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
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 = [
|
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" },
|
{ 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"
|
version = "12.0.4.66"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
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" },
|
||||||
{ 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-cusparse" },
|
||||||
{ 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 = [
|
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" },
|
{ 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"
|
version = "12.6.3.3"
|
||||||
source = { registry = "https://pypi.org/simple" }
|
source = { registry = "https://pypi.org/simple" }
|
||||||
dependencies = [
|
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 = [
|
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" },
|
{ 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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "overrides"
|
name = "overrides"
|
||||||
version = "7.7.0"
|
version = "7.7.0"
|
||||||
|
|||||||
+2
-24
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "lancedb"
|
name = "lancedb"
|
||||||
version = "0.32.0-beta.1"
|
version = "0.31.0-beta.6"
|
||||||
edition.workspace = true
|
edition.workspace = true
|
||||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||||
license.workspace = true
|
license.workspace = true
|
||||||
@@ -44,12 +44,11 @@ lance-io = { workspace = true }
|
|||||||
lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lindera"] }
|
lance-index = { workspace = true, features = ["tokenizer-jieba", "tokenizer-lindera"] }
|
||||||
lance-table = { workspace = true }
|
lance-table = { workspace = true }
|
||||||
lance-linalg = { workspace = true }
|
lance-linalg = { workspace = true }
|
||||||
|
lance-testing = { workspace = true }
|
||||||
lance-encoding = { workspace = true }
|
lance-encoding = { workspace = true }
|
||||||
lance-arrow = { workspace = true }
|
lance-arrow = { workspace = true }
|
||||||
lance-namespace = { workspace = true }
|
lance-namespace = { workspace = true }
|
||||||
lance-namespace-impls = { workspace = true }
|
lance-namespace-impls = { workspace = true }
|
||||||
metrics = { workspace = true, optional = true }
|
|
||||||
metrics-util = { workspace = true, optional = true }
|
|
||||||
moka = { workspace = true }
|
moka = { workspace = true }
|
||||||
pin-project = { workspace = true }
|
pin-project = { workspace = true }
|
||||||
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
|
||||||
@@ -94,7 +93,6 @@ semver = { workspace = true }
|
|||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
anyhow = "1"
|
anyhow = "1"
|
||||||
lance-testing = { workspace = true }
|
|
||||||
tempfile = "3.5.0"
|
tempfile = "3.5.0"
|
||||||
random_word = { version = "0.4.3", features = ["en"] }
|
random_word = { version = "0.4.3", features = ["en"] }
|
||||||
tokio = { version = "1.23", features = ["io-util", "macros", "net", "rt-multi-thread", "sync"] }
|
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"
|
rstest = "0.23.0"
|
||||||
test-log = "0.2"
|
test-log = "0.2"
|
||||||
serial_test = "3"
|
serial_test = "3"
|
||||||
[target.'cfg(unix)'.dev-dependencies]
|
|
||||||
pprof = { version = "0.14", features = ["flamegraph"] }
|
|
||||||
|
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
@@ -130,12 +126,6 @@ azure = [
|
|||||||
"lance-namespace-impls/dir-azure",
|
"lance-namespace-impls/dir-azure",
|
||||||
"lance-namespace-impls/credential-vendor-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 = [
|
huggingface = [
|
||||||
"lance/huggingface",
|
"lance/huggingface",
|
||||||
"lance-io/huggingface",
|
"lance-io/huggingface",
|
||||||
@@ -149,15 +139,6 @@ remote = [
|
|||||||
"lance-namespace-impls/rest",
|
"lance-namespace-impls/rest",
|
||||||
"lance-namespace-impls/rest-adapter",
|
"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"]
|
fp16kernels = ["lance-linalg/fp16kernels"]
|
||||||
s3-test = []
|
s3-test = []
|
||||||
bedrock = ["dep:aws-sdk-bedrockruntime"]
|
bedrock = ["dep:aws-sdk-bedrockruntime"]
|
||||||
@@ -183,9 +164,6 @@ required-features = ["sentence-transformers"]
|
|||||||
name = "bedrock"
|
name = "bedrock"
|
||||||
required-features = ["bedrock"]
|
required-features = ["bedrock"]
|
||||||
|
|
||||||
[[example]]
|
|
||||||
name = "bench_streaming_dataloader"
|
|
||||||
|
|
||||||
[[example]]
|
[[example]]
|
||||||
name = "simple"
|
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(
|
async fn prepare_namespace_root(
|
||||||
uri: &str,
|
uri: &str,
|
||||||
storage_options: &HashMap<String, String>,
|
storage_options: &HashMap<String, String>,
|
||||||
@@ -533,8 +520,6 @@ impl ListingDatabase {
|
|||||||
// will add a trailing '?' to the url
|
// will add a trailing '?' to the url
|
||||||
url.set_query(None);
|
url.set_query(None);
|
||||||
|
|
||||||
let storage_base_uri = Self::storage_base_uri(url.as_str());
|
|
||||||
|
|
||||||
let table_base_uri = if let Some(store) = engine {
|
let table_base_uri = if let Some(store) = engine {
|
||||||
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
|
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
|
||||||
WARN_ONCE.call_once(|| {
|
WARN_ONCE.call_once(|| {
|
||||||
@@ -547,6 +532,8 @@ impl ListingDatabase {
|
|||||||
url.to_string()
|
url.to_string()
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let plain_uri = url.to_string();
|
||||||
|
|
||||||
let session = request
|
let session = request
|
||||||
.session
|
.session
|
||||||
.clone()
|
.clone()
|
||||||
@@ -563,13 +550,13 @@ impl ListingDatabase {
|
|||||||
};
|
};
|
||||||
let (object_store, base_path) = ObjectStore::from_uri_and_params(
|
let (object_store, base_path) = ObjectStore::from_uri_and_params(
|
||||||
session.store_registry(),
|
session.store_registry(),
|
||||||
&storage_base_uri,
|
&plain_uri,
|
||||||
&os_params,
|
&os_params,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
if object_store.is_local() {
|
if object_store.is_local() {
|
||||||
Self::try_create_dir(&storage_base_uri).context(CreateDirSnafu {
|
Self::try_create_dir(&plain_uri).context(CreateDirSnafu {
|
||||||
path: storage_base_uri.clone(),
|
path: plain_uri.clone(),
|
||||||
})?;
|
})?;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -583,7 +570,7 @@ impl ListingDatabase {
|
|||||||
};
|
};
|
||||||
|
|
||||||
let namespace_database = Self::connect_namespace_database(
|
let namespace_database = Self::connect_namespace_database(
|
||||||
&storage_base_uri,
|
&table_base_uri,
|
||||||
options.storage_options.clone(),
|
options.storage_options.clone(),
|
||||||
request.namespace_client_properties.clone(),
|
request.namespace_client_properties.clone(),
|
||||||
request.read_consistency_interval,
|
request.read_consistency_interval,
|
||||||
@@ -1322,60 +1309,6 @@ mod tests {
|
|||||||
(tempdir, db)
|
(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]
|
#[tokio::test]
|
||||||
async fn test_clone_table_basic() {
|
async fn test_clone_table_basic() {
|
||||||
let (_tempdir, db) = setup_database().await;
|
let (_tempdir, db) = setup_database().await;
|
||||||
@@ -2350,24 +2283,6 @@ mod tests {
|
|||||||
assert_eq!(captured.as_deref(), Some("foo=bar"));
|
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
|
/// Regression: connecting via a URL-style URI (which goes through
|
||||||
/// `url::Url::parse` and the `query_pairs_mut()` path) must not
|
/// `url::Url::parse` and the `query_pairs_mut()` path) must not
|
||||||
/// append a trailing `?` to per-table URIs when the input URI has
|
/// append a trailing `?` to per-table URIs when the input URI has
|
||||||
|
|||||||
@@ -391,7 +391,6 @@ mod tests {
|
|||||||
SplitStrategy::Random {
|
SplitStrategy::Random {
|
||||||
seed: Some(42),
|
seed: Some(42),
|
||||||
sizes: SplitSizes::Percentages(vec![0.05, 0.30]),
|
sizes: SplitSizes::Percentages(vec![0.05, 0.30]),
|
||||||
clump_size: None,
|
|
||||||
},
|
},
|
||||||
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