mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-04 12:38:38 +00:00
Compare commits
27 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bff911a65d | |||
| 3a4cdb7aff | |||
| 142ac835d3 | |||
| 3f44f93e92 | |||
| 9dfa43a9de | |||
| 03e895fa5c | |||
| c31e53088e | |||
| 434a5be187 | |||
| 78aa005093 | |||
| 6191542cfe | |||
| 6af3088b91 | |||
| e73d4618d8 | |||
| 3d92106394 | |||
| 5810974b37 | |||
| 8b38500b07 | |||
| fd0a3b97d0 | |||
| b9f33ba1c9 | |||
| d4f4fef3ba | |||
| fbe6a5a3fd | |||
| 127054069a | |||
| b20931b8f7 | |||
| 396d68e490 | |||
| ad37f87387 | |||
| e93476f0e0 | |||
| 2b41fce033 | |||
| 04948fc4f6 | |||
| ff3c7111b9 |
@@ -1,137 +0,0 @@
|
||||
---
|
||||
name: lancedb-branch-ops
|
||||
description: Branch management for LanceDB tables via the REST API. Use this skill whenever someone wants to create, delete, list, or switch branches on a LanceDB table — or needs to make sure a write (metadata update, index build, etc.) lands on a specific branch instead of main. Invoke it even without the word "branch" if context makes clear they want an experimental copy of a table, want to isolate changes, or want to confirm a mutation didn't touch main. Covers: branches/list, branches/create, branches/delete, and passing "branch" in describe/update_field_metadata/create_index to target a non-main version.
|
||||
---
|
||||
|
||||
## Goal
|
||||
|
||||
Manage branches on a LanceDB table: list what exists, create new ones, delete stale ones, and direct read/write operations at a specific branch without touching main.
|
||||
|
||||
## Step 0: Establish the connection
|
||||
|
||||
Use the `lancedb-connect` skill to resolve the base URL and auth headers (`x-api-key`, `x-lancedb-database`). Skip this only if the connection is already known from the current conversation.
|
||||
|
||||
All examples below use `{base_url}` — substitute the resolved endpoint and include the auth headers on every request.
|
||||
|
||||
## The branch model (important)
|
||||
|
||||
LanceDB branches are named snapshots that diverge from the table's current state at creation time. There is **no checkout command** — you never switch the whole table to a branch. Instead, you **pass `"branch": "<name>"` in the request body** of any operation to target that branch. Omitting the key (or sending an empty body) always targets main.
|
||||
|
||||
`branches/list` returns only non-main branches. Main always exists and is not listed.
|
||||
|
||||
## List branches
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/list
|
||||
Content-Type: application/json
|
||||
|
||||
{}
|
||||
```
|
||||
|
||||
Response:
|
||||
```json
|
||||
{
|
||||
"branches": {
|
||||
"experiment-reindex": {"parentVersion": 1, "createAt": 1782506085, "manifestSize": 1029}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
If `branches` is `{}`, the table has no branches besides main.
|
||||
|
||||
## Create a branch
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/create
|
||||
Content-Type: application/json
|
||||
|
||||
{"name": "experiment-reindex"}
|
||||
```
|
||||
|
||||
HTTP 200 with `{}` body = success. The branch is created off the table's current state on main.
|
||||
|
||||
Verify by calling `branches/list` and confirming the new name appears.
|
||||
|
||||
## Delete a branch
|
||||
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/branches/delete
|
||||
Content-Type: application/json
|
||||
|
||||
{"name": "stale-2024"}
|
||||
```
|
||||
|
||||
HTTP 200 with `{}` body = success. Only the branch pointer is removed — main and all row data remain intact.
|
||||
|
||||
Verify by calling `branches/list` (name gone) and `describe` with no branch param (main still responds).
|
||||
|
||||
## Operate on a specific branch
|
||||
|
||||
Pass `"branch": "<name>"` in the body of any operation to scope it to that branch:
|
||||
|
||||
**Read schema on a branch:**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/describe
|
||||
Content-Type: application/json
|
||||
|
||||
{"branch": "wip-branch"}
|
||||
```
|
||||
|
||||
**Write metadata to a branch (not main):**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/update_field_metadata
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"branch": "wip-branch",
|
||||
"updates": [
|
||||
{
|
||||
"path": "category",
|
||||
"metadata": {"lancedb:description": "Product category label."},
|
||||
"replace": false
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Build an index on a branch:**
|
||||
```http
|
||||
POST {base_url}/v1/table/{table_id}/create_index
|
||||
Content-Type: application/json
|
||||
|
||||
{
|
||||
"branch": "wip-branch",
|
||||
"column": "category",
|
||||
"index_type": "BTREE"
|
||||
}
|
||||
```
|
||||
|
||||
## Verifying isolation
|
||||
|
||||
After writing to a branch, always confirm the change did NOT land on main:
|
||||
|
||||
```bash
|
||||
# Should show the new metadata
|
||||
curl -s -X POST {base_url}/v1/table/{table_id}/describe \
|
||||
-H "x-api-key: <key>" -H "x-lancedb-database: <db>" \
|
||||
-H "content-type: application/json" \
|
||||
-d '{"branch": "wip-branch"}'
|
||||
|
||||
# Should NOT show the new metadata
|
||||
curl -s -X POST {base_url}/v1/table/{table_id}/describe \
|
||||
-H "x-api-key: <key>" -H "x-lancedb-database: <db>" \
|
||||
-H "content-type: application/json" \
|
||||
-d '{}'
|
||||
```
|
||||
|
||||
## Quick reference
|
||||
|
||||
| Goal | Endpoint | Body |
|
||||
|------|----------|------|
|
||||
| List all branches | `branches/list` | `{}` |
|
||||
| Create a branch | `branches/create` | `{"name": "..."}` |
|
||||
| Delete a branch | `branches/delete` | `{"name": "..."}` |
|
||||
| Read schema on branch | `describe` | `{"branch": "..."}` |
|
||||
| Write metadata on branch | `update_field_metadata` | `{"branch": "...", "updates": [...]}` |
|
||||
| Build index on branch | `create_index` | `{"branch": "...", "column": ..., "index_type": ...}` |
|
||||
| Target main (default) | any endpoint | omit `"branch"` key |
|
||||
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.31.0-beta.6"
|
||||
current_version = "0.31.0-beta.4"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
# CODEOWNERS
|
||||
#
|
||||
# These owners will be the default owners for everything in the repo.
|
||||
# They will be requested for review when someone opens a pull request.
|
||||
#
|
||||
# See https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
|
||||
|
||||
# Default owners for everything
|
||||
* @jackye1995 @wjones127
|
||||
|
||||
# Release and publish workflows — changes here can affect supply chain security
|
||||
/.github/workflows/ @jackye1995 @wjones127 @Xuanwo
|
||||
|
||||
# Remote client and auth — sensitive networking and auth code
|
||||
/rust/lancedb/src/remote/ @jackye1995 @wjones127
|
||||
|
||||
# Python FFI boundary
|
||||
/python/src/ @jackye1995 @wjones127 @AyushExel
|
||||
|
||||
# NodeJS FFI boundary
|
||||
/nodejs/src/ @jackye1995 @wjones127
|
||||
@@ -25,7 +25,7 @@ jobs:
|
||||
# Only runs on tags that matches the make-release action
|
||||
if: startsWith(github.ref, 'refs/tags/v')
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
workspaces: rust
|
||||
@@ -47,7 +47,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,14 +36,14 @@ jobs:
|
||||
echo "guidelines = ${{ inputs.guidelines }}"
|
||||
|
||||
- name: Checkout Repo
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: ${{ inputs.branch }}
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
# pnpm 11 (used by the nodejs install step below) requires
|
||||
# Node >= 22.13; use 24 since 22 hits EOL in October.
|
||||
@@ -82,7 +82,7 @@ jobs:
|
||||
cache: maven
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v6
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Install Node.js dependencies for TypeScript bindings
|
||||
|
||||
@@ -30,13 +30,13 @@ jobs:
|
||||
echo "tag = ${{ inputs.tag || 'latest' }}"
|
||||
|
||||
- name: Checkout Repo LanceDB
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: true
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ jobs:
|
||||
name: Verify PR title / description conforms to semantic-release
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "18"
|
||||
# These rules are disabled because Github will always ensure there
|
||||
|
||||
@@ -35,7 +35,7 @@ jobs:
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Install dependencies needed for ubuntu
|
||||
run: |
|
||||
sudo apt install -y protobuf-compiler libssl-dev
|
||||
@@ -53,7 +53,7 @@ jobs:
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -e .
|
||||
python -m pip install --extra-index-url https://pypi.fury.io/lance-format/ --extra-index-url https://pypi.fury.io/lancedb/ -r ../docs/requirements.txt
|
||||
- name: Set up node
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 20
|
||||
cache: 'npm'
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Set up Java 8
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
@@ -73,7 +73,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -36,7 +36,7 @@ jobs:
|
||||
working-directory: ./java
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Set up Java 17
|
||||
uses: actions/setup-java@v4
|
||||
with:
|
||||
|
||||
@@ -19,7 +19,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
- name: Install license-header-checker
|
||||
working-directory: /tmp
|
||||
run: |
|
||||
|
||||
@@ -49,7 +49,7 @@ jobs:
|
||||
steps:
|
||||
- name: Output Inputs
|
||||
run: echo "${{ toJSON(github.event.inputs) }}"
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
@@ -38,14 +38,14 @@ jobs:
|
||||
CC: gcc-12
|
||||
CXX: g++-12
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v6
|
||||
- uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October. The library itself still supports Node >= 18
|
||||
@@ -86,14 +86,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v6
|
||||
- uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v4
|
||||
name: Setup Node.js 24 for build
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
@@ -130,7 +130,7 @@ jobs:
|
||||
echo "Run 'pnpm run docs', fix any warnings, and commit the changes."
|
||||
exit 1
|
||||
fi
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v4
|
||||
name: Setup Node.js ${{ matrix.node-version }} for test
|
||||
with:
|
||||
node-version: ${{ matrix.node-version }}
|
||||
@@ -166,14 +166,14 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: pnpm/action-setup@v6
|
||||
- uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
|
||||
@@ -32,7 +32,7 @@ jobs:
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -170,13 +170,13 @@ jobs:
|
||||
run:
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v6
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -190,7 +190,7 @@ jobs:
|
||||
toolchain: stable
|
||||
targets: ${{ matrix.settings.target }}
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v5
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry/index/
|
||||
@@ -244,7 +244,7 @@ jobs:
|
||||
if: ${{ !matrix.settings.docker }}
|
||||
shell: bash
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/*.node
|
||||
@@ -256,7 +256,7 @@ jobs:
|
||||
run: pnpm tsc
|
||||
- name: Upload Generic Artifacts
|
||||
if: ${{ matrix.settings.target == 'aarch64-apple-darwin' }}
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: |
|
||||
@@ -287,13 +287,13 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: nodejs
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v6
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup Node.js 24 for install
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
@@ -303,18 +303,18 @@ jobs:
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Setup Node.js ${{ matrix.node }} for test
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: ${{ matrix.node }}
|
||||
- name: Download artifacts
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: lancedb-${{ matrix.settings.target }}
|
||||
path: nodejs/dist/
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v8
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
@@ -339,13 +339,13 @@ jobs:
|
||||
needs:
|
||||
- test-lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v6
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 11.1.1
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
@@ -353,14 +353,14 @@ jobs:
|
||||
registry-url: "https://registry.npmjs.org"
|
||||
- name: Install dependencies
|
||||
run: pnpm install --frozen-lockfile
|
||||
- uses: actions/download-artifact@v8
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: nodejs-dist
|
||||
path: nodejs/dist
|
||||
# For testing purposes:
|
||||
# run-id: 13982782871
|
||||
# github-token: ${{ secrets.GITHUB_TOKEN }} # token with actions:read permissions on target repo
|
||||
- uses: actions/download-artifact@v8
|
||||
- uses: actions/download-artifact@v4
|
||||
name: Download arch-specific binaries
|
||||
with:
|
||||
pattern: lancedb-*
|
||||
@@ -398,7 +398,7 @@ jobs:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: ./.github/actions/create-failure-issue
|
||||
with:
|
||||
job-results: ${{ toJSON(needs) }}
|
||||
|
||||
@@ -41,7 +41,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -66,7 +66,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -95,7 +95,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -126,7 +126,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -160,7 +160,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -189,7 +189,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -212,7 +212,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: python
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
|
||||
@@ -40,7 +40,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -65,7 +65,7 @@ jobs:
|
||||
timeout-minutes: 10
|
||||
runs-on: ubuntu-24.04
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: EmbarkStudios/cargo-deny-action@v2
|
||||
with:
|
||||
command: check advisories bans licenses sources
|
||||
@@ -78,7 +78,7 @@ jobs:
|
||||
CC: clang
|
||||
CXX: clang++
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
# Building without a lock file often requires the latest Rust version since downstream
|
||||
# dependencies may have updated their minimum Rust version.
|
||||
- uses: actions-rust-lang/setup-rust-toolchain@v1
|
||||
@@ -113,7 +113,7 @@ jobs:
|
||||
CXX: clang++-18
|
||||
GH_TOKEN: ${{ secrets.SOPHON_READ_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -152,7 +152,7 @@ jobs:
|
||||
shell: bash
|
||||
working-directory: rust
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
@@ -181,7 +181,7 @@ jobs:
|
||||
run:
|
||||
working-directory: rust/lancedb
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
- name: Set target
|
||||
run: rustup target add ${{ matrix.target }}
|
||||
- uses: Swatinem/rust-cache@v2
|
||||
@@ -210,7 +210,7 @@ jobs:
|
||||
CC: clang-18
|
||||
CXX: clang++-18
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: true
|
||||
- name: Install dependencies
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
@@ -11,7 +11,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
|
||||
Generated
+78
-241
@@ -1750,7 +1750,7 @@ version = "3.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2039,9 +2039,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-epoch"
|
||||
version = "0.9.20"
|
||||
version = "0.9.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
|
||||
checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
@@ -3014,7 +3014,7 @@ dependencies = [
|
||||
"libc",
|
||||
"option-ext",
|
||||
"redox_users",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3231,7 +3231,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3423,8 +3423,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.4",
|
||||
@@ -3779,30 +3779,6 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "goosefs-sdk"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ae079b88ffe7772d12cfc5c40a5a324babb357893d95b5e3a22ae857f236c5f"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"bytes",
|
||||
"dashmap",
|
||||
"hostname",
|
||||
"prost",
|
||||
"prost-types",
|
||||
"rand 0.9.4",
|
||||
"reqwest 0.12.28",
|
||||
"serde",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
"tracing",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "group"
|
||||
version = "0.12.1"
|
||||
@@ -4014,17 +3990,6 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hostname"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "617aaa3557aef3810a6369d0a99fac8a080891b68bd9f9812a1eeda0c0730cbd"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.4",
|
||||
"libc",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "0.2.12"
|
||||
@@ -4185,19 +4150,6 @@ dependencies = [
|
||||
"webpki-roots 1.0.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-timeout"
|
||||
version = "0.5.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0"
|
||||
dependencies = [
|
||||
"hyper 1.9.0",
|
||||
"hyper-util",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tower-service",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hyper-util"
|
||||
version = "0.1.20"
|
||||
@@ -4514,7 +4466,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46"
|
||||
dependencies = [
|
||||
"hermit-abi",
|
||||
"libc",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4608,7 +4560,7 @@ dependencies = [
|
||||
"portable-atomic-util",
|
||||
"serde_core",
|
||||
"wasm-bindgen",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4774,8 +4726,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4849,8 +4801,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4859,7 +4811,6 @@ dependencies = [
|
||||
"arrow-ord",
|
||||
"arrow-schema",
|
||||
"arrow-select",
|
||||
"bytemuck",
|
||||
"bytes",
|
||||
"futures",
|
||||
"getrandom 0.2.17",
|
||||
@@ -4872,7 +4823,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4886,7 +4837,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -4895,8 +4846,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4906,8 +4857,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4945,8 +4896,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -4976,8 +4927,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -4994,8 +4945,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5004,8 +4955,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5040,8 +4991,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5071,8 +5022,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5137,8 +5088,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-arith",
|
||||
@@ -5179,8 +5130,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5196,8 +5147,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5209,8 +5160,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5233,7 +5184,7 @@ dependencies = [
|
||||
"lance-table",
|
||||
"log",
|
||||
"object_store",
|
||||
"quick-xml 0.40.1",
|
||||
"quick-xml 0.38.4",
|
||||
"rand 0.9.4",
|
||||
"reqwest 0.12.28",
|
||||
"roaring",
|
||||
@@ -5264,8 +5215,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5280,8 +5231,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5320,22 +5271,22 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
"criterion",
|
||||
"lance-arrow",
|
||||
"num-traits",
|
||||
"pprof 0.15.0",
|
||||
"pprof",
|
||||
"rand 0.9.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "9.0.0-beta.19"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.19#8f0e6d3a7c53438275b134c0ac1afbc80600616e"
|
||||
version = "9.0.0-beta.10"
|
||||
source = "git+https://github.com/jackye1995/lance.git?branch=jack%2Fsophon-pr-6325#1c5b5061c60934b4c18dbe86c5e91b4961105989"
|
||||
dependencies = [
|
||||
"icu_segmenter",
|
||||
"jieba-rs",
|
||||
@@ -5348,13 +5299,12 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.31.0-beta.6"
|
||||
version = "0.31.0-beta.4"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
"arrow-cast",
|
||||
"arrow-data",
|
||||
"arrow-ipc",
|
||||
@@ -5410,7 +5360,6 @@ dependencies = [
|
||||
"pin-project",
|
||||
"polars",
|
||||
"polars-arrow",
|
||||
"pprof 0.14.1",
|
||||
"rand 0.9.4",
|
||||
"random_word",
|
||||
"regex",
|
||||
@@ -5434,7 +5383,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.31.0-beta.6"
|
||||
version = "0.31.0-beta.4"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5459,7 +5408,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.34.0-beta.6"
|
||||
version = "0.34.0-beta.4"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -6002,9 +5951,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi"
|
||||
version = "3.10.3"
|
||||
version = "3.9.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0c71997d6f7ad4a756966e452426848ac27d3b37a295302d63afbbcce0270f93"
|
||||
checksum = "b41bda2ac390efb5e8d22025d925ccc3f3807d8c1bea6d19b36127247c4b8f83"
|
||||
dependencies = [
|
||||
"bitflags 2.11.1",
|
||||
"chrono",
|
||||
@@ -6027,9 +5976,9 @@ checksum = "c9c366d2c8c60b86fa632df75f745509b52f9128f91a6bad4c796e44abb505e1"
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive"
|
||||
version = "3.5.9"
|
||||
version = "3.5.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d4ba572deef53e2c386759a8c2014175a62679d74ff83adc205c8bc0e0285727"
|
||||
checksum = "61d66f70256ad5aef58659966064471d0ad90e2897bc36a5a5e0389c85aabc1e"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"ctor 1.0.5",
|
||||
@@ -6041,9 +5990,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "napi-derive-backend"
|
||||
version = "5.1.1"
|
||||
version = "5.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ddd961eb2aa8965e3f29722d754f3a86907eb1984e2fbcbe3fe87b9a02d6bfba"
|
||||
checksum = "81b4b08f15eed7a2a20c3f4c6314013fc3ac890a3afa9892b594485299ebdb2d"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"proc-macro2",
|
||||
@@ -6136,7 +6085,7 @@ version = "0.50.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -6391,9 +6340,7 @@ dependencies = [
|
||||
"opendal-layer-timeout",
|
||||
"opendal-service-azblob",
|
||||
"opendal-service-azdls",
|
||||
"opendal-service-cos",
|
||||
"opendal-service-gcs",
|
||||
"opendal-service-goosefs",
|
||||
"opendal-service-hf",
|
||||
"opendal-service-oss",
|
||||
"opendal-service-s3",
|
||||
@@ -6521,23 +6468,6 @@ dependencies = [
|
||||
"opendal-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opendal-service-cos"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aa8cafe9729213375c7331019b0cb756ad3e1aff7f45cd32c45eae91ebde8901"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"http 1.4.2",
|
||||
"log",
|
||||
"opendal-core",
|
||||
"quick-xml 0.39.4",
|
||||
"reqsign-core",
|
||||
"reqsign-file-read-tokio",
|
||||
"reqsign-tencent-cos",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opendal-service-gcs"
|
||||
version = "0.57.0"
|
||||
@@ -6559,20 +6489,6 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opendal-service-goosefs"
|
||||
version = "0.57.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69e43048bde419947ba826fbdc2f134d6c03f44ebf48bd33a03b72f9fc45fcb4"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"goosefs-sdk",
|
||||
"log",
|
||||
"opendal-core",
|
||||
"serde",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "opendal-service-hf"
|
||||
version = "0.57.0"
|
||||
@@ -7409,28 +7325,6 @@ version = "0.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391"
|
||||
|
||||
[[package]]
|
||||
name = "pprof"
|
||||
version = "0.14.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "afad4d4df7b31280028245f152d5a575083e2abb822d05736f5e47653e77689f"
|
||||
dependencies = [
|
||||
"aligned-vec",
|
||||
"backtrace",
|
||||
"cfg-if 1.0.4",
|
||||
"findshlibs",
|
||||
"inferno",
|
||||
"libc",
|
||||
"log",
|
||||
"nix",
|
||||
"once_cell",
|
||||
"smallvec",
|
||||
"spin 0.10.0",
|
||||
"symbolic-demangle",
|
||||
"tempfile",
|
||||
"thiserror 1.0.69",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pprof"
|
||||
version = "0.15.0"
|
||||
@@ -7506,8 +7400,8 @@ version = "0.14.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7"
|
||||
dependencies = [
|
||||
"heck 0.5.0",
|
||||
"itertools 0.14.0",
|
||||
"heck 0.4.1",
|
||||
"itertools 0.11.0",
|
||||
"log",
|
||||
"multimap",
|
||||
"petgraph",
|
||||
@@ -7526,7 +7420,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"itertools 0.14.0",
|
||||
"itertools 0.11.0",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -7688,6 +7582,15 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-xml"
|
||||
version = "0.38.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b66c2058c55a409d601666cffe35f04333cf1013010882cec174a7467cd4e21c"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-xml"
|
||||
version = "0.39.4"
|
||||
@@ -7698,15 +7601,6 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-xml"
|
||||
version = "0.40.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2474bd2e5029e7ccb6abb2ba48cf2383a333851dedf495901544281590c7da7f"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quinn"
|
||||
version = "0.11.9"
|
||||
@@ -7760,7 +7654,7 @@ dependencies = [
|
||||
"once_cell",
|
||||
"socket2 0.6.3",
|
||||
"tracing",
|
||||
"windows-sys 0.60.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8208,21 +8102,6 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "reqsign-tencent-cos"
|
||||
version = "3.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e128f19525861dbded59e1e7c17653a8ed63d573ca04aed708d552dbef5bb32a"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"http 1.4.2",
|
||||
"log",
|
||||
"percent-encoding",
|
||||
"reqsign-core",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "reqwest"
|
||||
version = "0.12.28"
|
||||
@@ -8515,7 +8394,7 @@ dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -8586,7 +8465,7 @@ dependencies = [
|
||||
"security-framework",
|
||||
"security-framework-sys",
|
||||
"webpki-root-certs",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -9148,7 +9027,7 @@ version = "0.8.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c1c97747dbf44bb1ca44a561ece23508e99cb592e862f22222dcf42f51d1e451"
|
||||
dependencies = [
|
||||
"heck 0.5.0",
|
||||
"heck 0.4.1",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -9160,7 +9039,7 @@ version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "54254b8531cafa275c5e096f62d48c81435d1015405a91198ddb11e967301d40"
|
||||
dependencies = [
|
||||
"heck 0.5.0",
|
||||
"heck 0.4.1",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.117",
|
||||
@@ -9593,7 +9472,7 @@ dependencies = [
|
||||
"getrandom 0.4.2",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -9908,45 +9787,6 @@ dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"h2 0.4.14",
|
||||
"http 1.4.2",
|
||||
"http-body 1.0.1",
|
||||
"http-body-util",
|
||||
"hyper 1.9.0",
|
||||
"hyper-timeout",
|
||||
"hyper-util",
|
||||
"percent-encoding",
|
||||
"pin-project",
|
||||
"socket2 0.6.3",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"prost",
|
||||
"tonic",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
@@ -9955,12 +9795,9 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"indexmap 2.14.0",
|
||||
"pin-project-lite",
|
||||
"slab",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
@@ -10570,7 +10407,7 @@ version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
+14
-15
@@ -13,25 +13,24 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.19", default-features = false, "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.19", "tag" = "v9.0.0-beta.19", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=9.0.0-beta.10", default-features = false, "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-core = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-datagen = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-file = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-io = { "version" = "=9.0.0-beta.10", default-features = false, "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-index = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-linalg = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-namespace = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=9.0.0-beta.10", default-features = false, "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-table = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-testing = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-datafusion = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-encoding = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
lance-arrow = { "version" = "=9.0.0-beta.10", "branch" = "jack/sophon-pr-6325", "git" = "https://github.com/jackye1995/lance.git" }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
arrow = { version = "58.0.0", optional = false }
|
||||
arrow-array = "58.0.0"
|
||||
arrow-buffer = "58.0.0"
|
||||
arrow-data = "58.0.0"
|
||||
arrow-ipc = "58.0.0"
|
||||
arrow-ord = "58.0.0"
|
||||
|
||||
@@ -51,6 +51,18 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0436
|
||||
{ id = "RUSTSEC-2024-0436", reason = "transitive via datafusion; awaiting ecosystem migration" },
|
||||
|
||||
# encoding: unmaintained. Reached through lindera-dictionary, which is
|
||||
# required by the native Lindera tokenizer path. Lindera has not migrated
|
||||
# off this crate yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2021-0153
|
||||
{ id = "RUSTSEC-2021-0153", reason = "transitive via lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# fast-float: unsound and unmaintained. Reached only through polars-arrow
|
||||
# from the optional Polars integration; replacement requires a Polars
|
||||
# dependency upgrade.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2024-0379
|
||||
{ id = "RUSTSEC-2024-0379", reason = "transitive via polars-arrow; waiting on Polars migration" },
|
||||
|
||||
# tantivy: segfault on malformed input due to missing bounds check.
|
||||
# Pulled in via lance for full-text search. We only feed tantivy
|
||||
# documents we construct ourselves, not attacker-controlled bytes.
|
||||
@@ -68,6 +80,18 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0119
|
||||
{ id = "RUSTSEC-2025-0119", reason = "transitive via hf-hub/indicatif; cosmetic formatting crate" },
|
||||
|
||||
# bincode: unmaintained. Reached through lindera and lindera-dictionary,
|
||||
# which are required by the native Lindera tokenizer path. Lindera has not
|
||||
# migrated to another serialization format yet.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2025-0141
|
||||
{ id = "RUSTSEC-2025-0141", reason = "transitive via lindera/lindera-dictionary for native Lindera tokenizer" },
|
||||
|
||||
# lru: soundness issue in IterMut. Reached only through aws-sdk-s3 in
|
||||
# LanceDB's dev-dependency graph; LanceDB does not use that iterator
|
||||
# directly. Clearing this requires the AWS SDK chain to update lru.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0002
|
||||
{ id = "RUSTSEC-2026-0002", reason = "transitive via aws-sdk-s3 dev-dependency; waiting on AWS SDK lru upgrade" },
|
||||
|
||||
# rustls-webpki 0.101.7 (old major line): name-constraint checks for
|
||||
# URI / wildcard names. Pulled in only via the legacy rustls 0.21 chain
|
||||
# from aws-smithy-http-client. The 0.103 line we actively use is patched.
|
||||
@@ -84,23 +108,17 @@ ignore = [
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0104
|
||||
{ id = "RUSTSEC-2026-0104", reason = "only affects rustls-webpki 0.101 from legacy aws-smithy/rustls 0.21 chain" },
|
||||
|
||||
# rand 0.8.5: soundness issue only when ThreadRng reseeds inside a custom
|
||||
# logger. Reached through several transitive chains. LanceDB does not use
|
||||
# rand from a custom logger; upgrade once all pinned chains accept 0.8.6+.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0097
|
||||
{ id = "RUSTSEC-2026-0097", reason = "transitive rand 0.8.5; LanceDB does not call ThreadRng from custom logging" },
|
||||
|
||||
# pyo3 advisories in the Python bindings; tracked pending a patched pyo3 release.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0176
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0177
|
||||
{ id = "RUSTSEC-2026-0176", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
{ id = "RUSTSEC-2026-0177", reason = "pyo3 in Python bindings; awaiting patched pyo3 release" },
|
||||
|
||||
# quick-xml < 0.41.0: quadratic runtime on duplicate attribute names (DoS).
|
||||
# quick-xml < 0.41.0: unbounded namespace-declaration allocation in NsReader (DoS).
|
||||
# Pulled in transitively by inferno (dev-only flame-graph dep), lance-namespace-impls
|
||||
# (git dep from lance), and opendal/reqsign (cloud storage XML parsing). The XML
|
||||
# parsed by opendal/reqsign comes from trusted cloud-storage endpoints (S3, GCS,
|
||||
# Azure), not attacker-controlled input. Clearing requires upstream crates to migrate
|
||||
# to quick-xml >= 0.41.0.
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0194
|
||||
# https://rustsec.org/advisories/RUSTSEC-2026-0195
|
||||
{ id = "RUSTSEC-2026-0194", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
{ id = "RUSTSEC-2026-0195", reason = "transitive via inferno/lance/opendal; XML from trusted cloud endpoints, not attacker-controlled" },
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.31.0-beta.6</version>
|
||||
<version>0.31.0-beta.4</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -518,9 +518,6 @@ x > 5 OR y = 'test'
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND rather
|
||||
than replacing the previous filter.
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
@@ -398,26 +398,6 @@ Drop an index from the table.
|
||||
|
||||
***
|
||||
|
||||
### getLsmWriteSpec()
|
||||
|
||||
```ts
|
||||
abstract getLsmWriteSpec(): Promise<undefined | LsmWriteSpec>
|
||||
```
|
||||
|
||||
Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) currently installed on this table.
|
||||
|
||||
Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
spec has been set, or it was removed with [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)).
|
||||
The returned spec — including its `maintainedIndexes` and
|
||||
`writerConfigDefaults` — mirrors what was passed to
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec).
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`undefined` \| [`LsmWriteSpec`](../interfaces/LsmWriteSpec.md)>
|
||||
|
||||
***
|
||||
|
||||
### indexStats()
|
||||
|
||||
```ts
|
||||
|
||||
@@ -767,9 +767,6 @@ x > 5 OR y = 'test'
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND rather
|
||||
than replacing the previous filter.
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
@@ -8,14 +8,6 @@
|
||||
|
||||
## Properties
|
||||
|
||||
### clumpSize?
|
||||
|
||||
```ts
|
||||
optional clumpSize: number;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### counts?
|
||||
|
||||
```ts
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.31.0-beta.6</version>
|
||||
<version>0.31.0-beta.4</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.31.0-beta.6</version>
|
||||
<version>0.31.0-beta.4</version>
|
||||
<packaging>pom</packaging>
|
||||
<name>${project.artifactId}</name>
|
||||
<description>LanceDB Java SDK Parent POM</description>
|
||||
@@ -28,7 +28,7 @@
|
||||
<properties>
|
||||
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
|
||||
<arrow.version>15.0.0</arrow.version>
|
||||
<lance-core.version>9.0.0-beta.19</lance-core.version>
|
||||
<lance-core.version>9.0.0-beta.10</lance-core.version>
|
||||
<spotless.skip>false</spotless.skip>
|
||||
<spotless.version>2.30.0</spotless.version>
|
||||
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.31.0-beta.6"
|
||||
version = "0.31.0-beta.4"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
@@ -44,6 +44,6 @@ aws-lc-rs = "=1.16.3"
|
||||
napi-build = "2.3.1"
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/goosefs"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -215,20 +215,6 @@ describe("Query orderBy", () => {
|
||||
expect(results[2].score).toBeCloseTo(4.1, 0.001);
|
||||
});
|
||||
|
||||
it("should combine repeated where clauses with AND", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
.where("score > 1.0")
|
||||
.where("score < 3.0")
|
||||
.orderBy({ columnName: "score" })
|
||||
.toArray();
|
||||
// Only rows matching both predicates should be returned, rather than the
|
||||
// second where() silently replacing the first.
|
||||
expect(results.length).toBe(2);
|
||||
expect(results[0].score).toBeCloseTo(1.2, 0.001);
|
||||
expect(results[1].score).toBeCloseTo(2.8, 0.001);
|
||||
});
|
||||
|
||||
it("should support method chaining with limit", async () => {
|
||||
const results = await table
|
||||
.query()
|
||||
|
||||
@@ -2992,56 +2992,6 @@ describe("setLsmWriteSpec / unsetLsmWriteSpec", () => {
|
||||
}),
|
||||
).rejects.toThrow();
|
||||
});
|
||||
|
||||
it("reads back the installed spec via getLsmWriteSpec", async () => {
|
||||
const conn = await connect(tmpDir.name);
|
||||
const table = await makeTable(conn);
|
||||
await table.setUnenforcedPrimaryKey("id");
|
||||
|
||||
// Nothing installed yet.
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// A real scalar index is needed to name it as a maintained index.
|
||||
await table.add([{ id: 1 }, { id: 2 }, { id: 3 }]);
|
||||
await table.createIndex("id");
|
||||
const indexName = (await table.listIndices())[0].name;
|
||||
|
||||
// Bucket spec round-trips, including maintained indexes and writer config
|
||||
// defaults. Lance writer-config keys are canonically snake_case.
|
||||
// biome-ignore lint/style/useNamingConvention: Lance writer-config keys are snake_case
|
||||
const writerConfigDefaults = { durable_write: "false" };
|
||||
await table.setLsmWriteSpec({
|
||||
specType: "bucket",
|
||||
column: "id",
|
||||
numBuckets: 4,
|
||||
maintainedIndexes: [indexName],
|
||||
writerConfigDefaults,
|
||||
});
|
||||
const spec = await table.getLsmWriteSpec();
|
||||
expect(spec).toBeDefined();
|
||||
expect(spec?.specType).toBe("bucket");
|
||||
expect(spec?.column).toBe("id");
|
||||
expect(spec?.numBuckets).toBe(4);
|
||||
expect(spec?.maintainedIndexes).toEqual([indexName]);
|
||||
expect(spec?.writerConfigDefaults).toEqual(writerConfigDefaults);
|
||||
|
||||
// After unset, undefined again.
|
||||
await table.unsetLsmWriteSpec();
|
||||
expect(await table.getLsmWriteSpec()).toBeUndefined();
|
||||
|
||||
// Identity round-trips (column recovered from the schema).
|
||||
await table.setLsmWriteSpec({ specType: "identity", column: "id" });
|
||||
const identity = await table.getLsmWriteSpec();
|
||||
expect(identity?.specType).toBe("identity");
|
||||
expect(identity?.column).toBe("id");
|
||||
await table.unsetLsmWriteSpec();
|
||||
|
||||
// Unsharded round-trips (no routing column).
|
||||
await table.setLsmWriteSpec({ specType: "unsharded" });
|
||||
const unsharded = await table.getLsmWriteSpec();
|
||||
expect(unsharded?.specType).toBe("unsharded");
|
||||
expect(unsharded?.column).toBeFalsy();
|
||||
});
|
||||
});
|
||||
|
||||
describe("LSM merge insert", () => {
|
||||
|
||||
@@ -362,9 +362,6 @@ export class StandardQueryBase<
|
||||
*
|
||||
* Filtering performance can often be improved by creating a scalar index
|
||||
* on the filter column(s).
|
||||
*
|
||||
* Calling this multiple times combines the filters with a logical AND rather
|
||||
* than replacing the previous filter.
|
||||
*/
|
||||
where(predicate: string): this {
|
||||
this.doCall((inner: NativeQueryType) => inner.onlyIf(predicate));
|
||||
|
||||
@@ -585,17 +585,6 @@ export abstract class Table {
|
||||
* @returns {Promise<void>}
|
||||
*/
|
||||
abstract unsetLsmWriteSpec(): Promise<void>;
|
||||
/**
|
||||
* Read the {@link LsmWriteSpec} currently installed on this table.
|
||||
*
|
||||
* Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
|
||||
* spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}).
|
||||
* The returned spec — including its `maintainedIndexes` and
|
||||
* `writerConfigDefaults` — mirrors what was passed to
|
||||
* {@link Table#setLsmWriteSpec}.
|
||||
* @returns {Promise<LsmWriteSpec | undefined>}
|
||||
*/
|
||||
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
|
||||
/**
|
||||
* Drain and close any cached MemWAL shard writers held for this table.
|
||||
*
|
||||
@@ -1102,15 +1091,6 @@ export class LocalTable extends Table {
|
||||
return await this.inner.unsetLsmWriteSpec();
|
||||
}
|
||||
|
||||
async getLsmWriteSpec(): Promise<LsmWriteSpec | undefined> {
|
||||
// The native binding types `specType` as a plain `string`; narrow it back
|
||||
// to the public union. The Rust `From` impl only ever emits one of the
|
||||
// three valid values, so the cast is safe.
|
||||
return ((await this.inner.getLsmWriteSpec()) ?? undefined) as
|
||||
| LsmWriteSpec
|
||||
| undefined;
|
||||
}
|
||||
|
||||
async closeLsmWriters(): Promise<void> {
|
||||
return await this.inner.closeLsmWriters();
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
+2
-2
@@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "@lancedb/lancedb",
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.31.0-beta.6",
|
||||
"version": "0.31.0-beta.4",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
|
||||
+6
-4
@@ -3,7 +3,7 @@
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use lancedb::{ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder};
|
||||
use lancedb::{arrow::IntoArrow, ipc::ipc_file_to_batches, table::merge::MergeInsertBuilder};
|
||||
use napi::bindgen_prelude::*;
|
||||
use napi_derive::napi;
|
||||
|
||||
@@ -66,9 +66,11 @@ impl NativeMergeInsertBuilder {
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn execute(&self, buf: Buffer) -> napi::Result<MergeResult> {
|
||||
let data = ipc_file_to_batches(buf.to_vec()).map_err(|e| {
|
||||
napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e)))
|
||||
})?;
|
||||
let data = ipc_file_to_batches(buf.to_vec())
|
||||
.and_then(IntoArrow::into_arrow)
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!("Failed to read IPC file: {}", convert_error(&e)))
|
||||
})?;
|
||||
|
||||
let this = self.clone();
|
||||
|
||||
|
||||
@@ -16,7 +16,6 @@ pub struct SplitRandomOptions {
|
||||
pub counts: Option<Vec<i64>>,
|
||||
pub fixed: Option<i64>,
|
||||
pub seed: Option<i64>,
|
||||
pub clump_size: Option<i64>,
|
||||
pub split_names: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -126,15 +125,10 @@ impl PermutationBuilder {
|
||||
};
|
||||
|
||||
let seed = options.seed.map(|s| s as u64);
|
||||
let clump_size = options.clump_size.map(|c| c as u64);
|
||||
|
||||
self.modify(|builder| {
|
||||
builder.with_split_strategy(
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
},
|
||||
SplitStrategy::Random { seed, sizes },
|
||||
options.split_names.clone(),
|
||||
)
|
||||
})
|
||||
|
||||
@@ -411,16 +411,6 @@ impl Table {
|
||||
.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn get_lsm_write_spec(&self) -> napi::Result<Option<LsmWriteSpec>> {
|
||||
let spec = self
|
||||
.inner_ref()?
|
||||
.get_lsm_write_spec()
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn close_lsm_writers(&self) -> napi::Result<()> {
|
||||
self.inner_ref()?.close_lsm_writers().await.default_error()
|
||||
@@ -738,47 +728,6 @@ impl TryFrom<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(spec: lancedb::table::LsmWriteSpec) -> Self {
|
||||
use lancedb::table::LsmWriteSpec as Native;
|
||||
match spec {
|
||||
Native::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "bucket".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: Some(num_buckets),
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Identity {
|
||||
column,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "identity".to_string(),
|
||||
column: Some(column),
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
Native::Unsharded {
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => Self {
|
||||
spec_type: "unsharded".to_string(),
|
||||
column: None,
|
||||
num_buckets: None,
|
||||
maintained_indexes: Some(maintained_indexes),
|
||||
writer_config_defaults: Some(writer_config_defaults),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics about a compaction operation.
|
||||
#[napi(object)]
|
||||
#[derive(Clone, Debug)]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.34.0-beta.6"
|
||||
current_version = "0.34.0-beta.4"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.34.0-beta.6"
|
||||
version = "0.34.0-beta.4"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
@@ -47,6 +47,6 @@ pyo3-build-config = { version = "0.28", features = [
|
||||
] }
|
||||
|
||||
[features]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs"]
|
||||
default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface"]
|
||||
fp16kernels = ["lancedb/fp16kernels"]
|
||||
remote = ["lancedb/remote"]
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
"""Benchmark for StreamingDataset throughput.
|
||||
|
||||
Sweeps read_batch_size from 1 to 16384 to show how amortising the per-request
|
||||
overhead scales. Each row at each chunk size is timed via the real
|
||||
StreamingDataset so the numbers reflect production code.
|
||||
|
||||
Run with:
|
||||
cd python
|
||||
uv run --extra tests benchmarks/bench_streaming_dataloader.py
|
||||
|
||||
Optional env vars:
|
||||
BENCH_NUM_ROWS — total rows in the table (default 49152 = 24 × 2048)
|
||||
BENCH_NUM_SPLITS — number of splits (default 24)
|
||||
BENCH_STEPS — round-robin cycles to time per chunk size (default 100)
|
||||
BENCH_ROW_BYTES — approximate bytes per row padded with a binary column
|
||||
(default 4096, mimics a small embedding/image patch)
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import tempfile
|
||||
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
|
||||
from lancedb.streaming import StreamingDataset
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
NUM_SPLITS = int(os.environ.get("BENCH_NUM_SPLITS", 24))
|
||||
# Default: 2048 rows per split so every chunk size up to 16Ki has ≥1 full
|
||||
# chunk (except 16Ki itself which gets a single full-split fetch — still valid).
|
||||
NUM_ROWS = int(os.environ.get("BENCH_NUM_ROWS", NUM_SPLITS * 2048))
|
||||
STEPS = int(os.environ.get("BENCH_STEPS", 100))
|
||||
ROW_BYTES = int(os.environ.get("BENCH_ROW_BYTES", 4096))
|
||||
|
||||
assert NUM_ROWS % NUM_SPLITS == 0, "NUM_ROWS must be divisible by NUM_SPLITS"
|
||||
|
||||
CHUNK_SIZES = [1, 4, 16, 64, 256, 1024, 4096, 16384]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Table helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def make_table(db_path: str) -> lancedb.table.Table:
|
||||
db = lancedb.connect(db_path)
|
||||
payload = b"x" * ROW_BYTES
|
||||
data = pa.table(
|
||||
{
|
||||
"id": pa.array(range(NUM_ROWS), type=pa.int32()),
|
||||
"payload": pa.array([payload] * NUM_ROWS, type=pa.large_binary()),
|
||||
}
|
||||
)
|
||||
return db.create_table("bench", data, mode="overwrite")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def bench_chunk(table, chunk_size: int, steps: int) -> tuple[int, float]:
|
||||
"""Return (rows_drained, elapsed_seconds) for one timed run."""
|
||||
total_rows = steps * NUM_SPLITS
|
||||
ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk_size
|
||||
)
|
||||
count = 0
|
||||
t0 = time.perf_counter()
|
||||
for _ in ds:
|
||||
count += 1
|
||||
if count >= total_rows:
|
||||
break
|
||||
return count, time.perf_counter() - t0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rows_per_split = NUM_ROWS // NUM_SPLITS
|
||||
print("Benchmark config:")
|
||||
print(
|
||||
f" NUM_ROWS={NUM_ROWS} NUM_SPLITS={NUM_SPLITS} "
|
||||
f"rows/split={rows_per_split} STEPS={STEPS} ROW_BYTES={ROW_BYTES}"
|
||||
)
|
||||
print(f" ~{NUM_ROWS * ROW_BYTES / 1024 / 1024:.1f} MB total table size")
|
||||
print()
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
print("Creating table...", flush=True)
|
||||
table = make_table(tmp)
|
||||
|
||||
cols = (
|
||||
f"{'chunk':>6} {'rows':>6} {'elapsed':>8} {'rows/s':>10} {'ms/step':>9}"
|
||||
)
|
||||
print(f"\n{cols}")
|
||||
print("-" * 52)
|
||||
|
||||
for chunk in CHUNK_SIZES:
|
||||
# Warm-up pass (one step's worth of rows)
|
||||
warmup_ds = StreamingDataset(
|
||||
table, num_splits=NUM_SPLITS, shuffle_seed=42, read_batch_size=chunk
|
||||
)
|
||||
warmup_count = 0
|
||||
for _ in warmup_ds:
|
||||
warmup_count += 1
|
||||
if warmup_count >= NUM_SPLITS:
|
||||
break
|
||||
|
||||
drained, elapsed = bench_chunk(table, chunk, STEPS)
|
||||
rows_per_sec = drained / elapsed if elapsed > 0 else float("inf")
|
||||
ms_per_step = elapsed / STEPS * 1000
|
||||
|
||||
print(
|
||||
f"{chunk:>6} {drained:>6} {elapsed:>7.3f}s "
|
||||
f"{rows_per_sec:>10.0f} {ms_per_step:>8.1f}ms"
|
||||
)
|
||||
|
||||
print()
|
||||
print("Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -17,6 +17,17 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||
from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .udf import (
|
||||
udf,
|
||||
table_udf,
|
||||
Udf,
|
||||
JobHandle,
|
||||
JobFailedError,
|
||||
MaterializedView,
|
||||
AsyncJobHandle,
|
||||
AsyncMaterializedView,
|
||||
)
|
||||
from .lineage import Lineage, Node, Edge, FunctionRef
|
||||
from .schema import vector
|
||||
from .table import AsyncTable, Table
|
||||
from ._lancedb import Session
|
||||
@@ -448,6 +459,18 @@ async def connect_async(
|
||||
|
||||
|
||||
__all__ = [
|
||||
"udf",
|
||||
"table_udf",
|
||||
"Udf",
|
||||
"JobHandle",
|
||||
"JobFailedError",
|
||||
"MaterializedView",
|
||||
"AsyncJobHandle",
|
||||
"AsyncMaterializedView",
|
||||
"Lineage",
|
||||
"Node",
|
||||
"Edge",
|
||||
"FunctionRef",
|
||||
"connect",
|
||||
"connect_async",
|
||||
"connect_namespace",
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -54,9 +53,7 @@ class PyExpr:
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
def expr_lit(
|
||||
value: Union[bool, int, float, str, bytes, date, datetime, Decimal],
|
||||
) -> PyExpr: ...
|
||||
def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ...
|
||||
def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ...
|
||||
|
||||
class Session:
|
||||
@@ -229,7 +226,6 @@ class Table:
|
||||
async def set_unenforced_primary_key(self, columns: List[str]) -> None: ...
|
||||
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
|
||||
async def unset_lsm_write_spec(self) -> None: ...
|
||||
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
|
||||
async def close_lsm_writers(self) -> None: ...
|
||||
@property
|
||||
def tags(self) -> Tags: ...
|
||||
@@ -286,23 +282,6 @@ async def connect(
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
oauth_config: Optional[Any] = None,
|
||||
) -> Connection: ...
|
||||
def connect_namespace(
|
||||
namespace_client_impl: str,
|
||||
namespace_client_properties: Dict[str, str],
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
) -> Connection: ...
|
||||
def connect_namespace_client(
|
||||
namespace_client: Any,
|
||||
read_consistency_interval: Optional[float] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
) -> Connection: ...
|
||||
|
||||
class RecordBatchStream:
|
||||
@property
|
||||
|
||||
@@ -65,6 +65,7 @@ if TYPE_CHECKING:
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
from .udf import MaterializedView, AsyncMaterializedView
|
||||
|
||||
from .namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
@@ -562,6 +563,259 @@ class DBConnection(EnforceOverrides):
|
||||
"""
|
||||
raise NotImplementedError("serialize is not supported for this connection type")
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION).
|
||||
|
||||
Pass a ``@udf`` / ``@table_udf``-decorated function (preferred):
|
||||
|
||||
db.create_function(embed)
|
||||
|
||||
or the explicit fields:
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name: str or Udf
|
||||
A decorated UDF object, or the function name.
|
||||
language: str
|
||||
Implementation language (currently "python").
|
||||
return_type: str
|
||||
SQL return type, e.g. "FLOAT", "FLOAT[1536]",
|
||||
"STRUCT(a FLOAT, b VARCHAR)", "TABLE(chunk VARCHAR, idx INT)".
|
||||
body: str
|
||||
Function body: source text, or base64 cloudpickle bytes when
|
||||
options["body_format"] == "cloudpickle".
|
||||
options: dict, optional
|
||||
input_columns, pip, num_gpus, batch_size, timeout,
|
||||
error_policy, docker_image, body_format, ...
|
||||
replace: bool
|
||||
Drop an existing function of the same name first.
|
||||
"""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
LOOP.run(self._conn.create_function(name, language, return_type, body, options))
|
||||
|
||||
def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return LOOP.run(self._conn.list_functions())
|
||||
|
||||
def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
LOOP.run(self._conn.drop_function(name))
|
||||
|
||||
def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "MaterializedView":
|
||||
"""Create a materialized view (CREATE MATERIALIZED VIEW); returns a
|
||||
`MaterializedView` handle (``.wait()`` blocks until it is populated).
|
||||
|
||||
Two ways to specify the view body:
|
||||
|
||||
- ergonomic: pass ``source`` (a table name or table) and ``select``
|
||||
items -- column names, expression strings ("embed(body)"),
|
||||
(alias, expression) tuples, or ``@udf`` / ``@table_udf`` objects.
|
||||
The SELECT is assembled and parsed server-side (one parser, shared
|
||||
with SQL).
|
||||
- raw: pass ``query=`` with a full SELECT, e.g.
|
||||
"SELECT id, embed(body) AS vec FROM articles WHERE id > 1".
|
||||
|
||||
`partition_by` partitions the view's (single) table function on a source
|
||||
column. If that column has an IVF vector index the server partitions by
|
||||
its index clusters (image-dedup style); otherwise it groups by distinct
|
||||
value. (Geneva's `partition_by` and `partition_by_indexed_column` unify
|
||||
here -- the engine picks the strategy from the column.)
|
||||
"""
|
||||
from .udf import build_view_query, MaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
self._drop_view_if_exists(name)
|
||||
job_id = LOOP.run(
|
||||
self._conn.create_materialized_view(
|
||||
name,
|
||||
query=query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
)
|
||||
return MaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def _drop_view_if_exists(self, name: str) -> None:
|
||||
# `replace=True` is "drop if present"; only a not-found error is
|
||||
# benign here. Anything else (perms, server fault) must surface rather
|
||||
# than be masked by a later create failure.
|
||||
try:
|
||||
self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""A `JobHandle` for reconnecting to an inflight job by id -- e.g. an
|
||||
id you stored, or one returned from the SQL / REST surface. Submit
|
||||
methods (`refresh_column`, `MaterializedView.refresh`) already return a
|
||||
handle directly, so you do not need this to wait on a fresh submission."""
|
||||
from .udf import JobHandle
|
||||
|
||||
return JobHandle(self, job_id)
|
||||
|
||||
def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. `direction` is "upstream" | "downstream" | "both" (server
|
||||
default both); `depth` limits column-hops (transitive when omitted)."""
|
||||
# `self._conn` is the AsyncConnection; drive its async `lineage`
|
||||
# (which parses the JSON) on the loop, mirroring create_materialized_view.
|
||||
return LOOP.run(
|
||||
self._conn.lineage(table, column, direction=direction, depth=depth)
|
||||
)
|
||||
|
||||
def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a materialized-view refresh, return the job id.
|
||||
The public surface is ``MaterializedView.refresh()`` (which returns a
|
||||
`JobHandle`); this stays private so refresh is only reached through the
|
||||
handle.
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._conn._refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
)
|
||||
|
||||
def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH). Returns a
|
||||
plan with .has_work / .source_version / .last_refreshed_version /
|
||||
.full_refresh / .rebuild / .units_total. `full=True` plans a full
|
||||
rebuild (incremental planning needs stable row IDs on the source)."""
|
||||
return LOOP.run(
|
||||
self._conn.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
)
|
||||
|
||||
def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options (ALTER MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.alter_materialized_view(name, auto_refresh=auto_refresh))
|
||||
|
||||
def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition (DROP MATERIALIZED VIEW)."""
|
||||
LOOP.run(self._conn.drop_materialized_view(name))
|
||||
|
||||
def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return LOOP.run(self._conn.list_materialized_views())
|
||||
|
||||
def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return LOOP.run(self._conn.list_jobs())
|
||||
|
||||
def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
|
||||
Passing ``table`` (the job's table) lets the server answer with an O(1)
|
||||
single-node read instead of scanning the database's active jobs.
|
||||
Returns the job's status, or None if it's unknown or no longer active.
|
||||
"""
|
||||
return LOOP.run(self._conn.get_job(job_id, table))
|
||||
|
||||
def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False if none was inflight (already finished or
|
||||
unknown id) -- cancellation is best-effort.
|
||||
"""
|
||||
return LOOP.run(self._conn.cancel_job(job_id))
|
||||
|
||||
def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Pass ``job_id`` to narrow to a single job. Unlike :meth:`list_jobs`
|
||||
(live, inflight) these are the terminal records.
|
||||
"""
|
||||
return LOOP.run(self._conn.job_history(job_id))
|
||||
|
||||
def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS),
|
||||
optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return LOOP.run(self._conn.errors(job_id, table))
|
||||
|
||||
|
||||
class LanceDBConnection(DBConnection):
|
||||
"""
|
||||
@@ -1787,6 +2041,200 @@ class AsyncConnection(object):
|
||||
)
|
||||
return AsyncTable(table)
|
||||
|
||||
# -- Derived compute: functions, materialized views, jobs -------------
|
||||
# Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
# connections raise NotImplementedError for now.
|
||||
|
||||
async def create_function(
|
||||
self,
|
||||
name,
|
||||
language: str = "python",
|
||||
return_type: Optional[str] = None,
|
||||
body: Optional[str] = None,
|
||||
options: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
replace: bool = False,
|
||||
):
|
||||
"""Register a UDF (CREATE FUNCTION). Accepts a ``@udf``/``@table_udf``
|
||||
object (preferred) or the explicit (name, language, return_type, body,
|
||||
options)."""
|
||||
from .udf import Udf
|
||||
|
||||
if isinstance(name, Udf):
|
||||
req = name.create_request()
|
||||
name, language, return_type, body, options = (
|
||||
req["name"],
|
||||
req["language"],
|
||||
req["return_type"],
|
||||
req["body"],
|
||||
req["options"],
|
||||
)
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_function(name)
|
||||
except Exception:
|
||||
pass
|
||||
await self._inner.create_function(name, language, return_type, body, options)
|
||||
|
||||
async def list_functions(self):
|
||||
"""List registered functions (SHOW FUNCTIONS)."""
|
||||
return await self._inner.list_functions()
|
||||
|
||||
async def drop_function(self, name: str):
|
||||
"""Drop a registered function (DROP FUNCTION)."""
|
||||
await self._inner.drop_function(name)
|
||||
|
||||
async def create_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
source=None,
|
||||
select=None,
|
||||
*,
|
||||
query: Optional[str] = None,
|
||||
where: Optional[str] = None,
|
||||
auto_refresh: bool = False,
|
||||
with_no_data: bool = False,
|
||||
replace: bool = False,
|
||||
partition_by: Optional[str] = None,
|
||||
) -> "AsyncMaterializedView":
|
||||
"""Create a materialized view; returns an `AsyncMaterializedView`
|
||||
handle (``.wait()`` blocks until populated). Pass either ``query=`` (a
|
||||
full SELECT) or ``source`` + ``select`` items; `partition_by`
|
||||
partitions the view's table function on a source column (index-cluster
|
||||
if the column is IVF-indexed, else distinct-value). See the sync
|
||||
method for the select grammar."""
|
||||
from .udf import build_view_query, AsyncMaterializedView
|
||||
|
||||
if query is None:
|
||||
if source is None or select is None:
|
||||
raise ValueError(
|
||||
"create_materialized_view needs either query= or both "
|
||||
"source and select"
|
||||
)
|
||||
query = build_view_query(source, select)
|
||||
if where:
|
||||
query += f" WHERE {where}"
|
||||
if replace:
|
||||
try:
|
||||
await self.drop_materialized_view(name)
|
||||
except Exception as e:
|
||||
msg = str(e).lower()
|
||||
if "not found" not in msg and "does not exist" not in msg:
|
||||
raise
|
||||
job_id = await self._inner.create_materialized_view(
|
||||
name,
|
||||
query,
|
||||
auto_refresh=auto_refresh,
|
||||
with_no_data=with_no_data,
|
||||
partition_by=partition_by,
|
||||
)
|
||||
return AsyncMaterializedView(self, name, job_id=job_id)
|
||||
|
||||
def job(self, job_id: str):
|
||||
"""An `AsyncJobHandle` for reconnecting to an inflight job by id (a
|
||||
stored id, or one from the SQL / REST surface). Submit methods already
|
||||
return a handle, so this is only needed to re-attach to an existing
|
||||
job."""
|
||||
from .udf import AsyncJobHandle
|
||||
|
||||
return AsyncJobHandle(self, job_id)
|
||||
|
||||
async def lineage(
|
||||
self,
|
||||
table: str,
|
||||
column: Optional[str] = None,
|
||||
*,
|
||||
direction: Optional[str] = None,
|
||||
depth: Optional[int] = None,
|
||||
):
|
||||
"""Derived-compute lineage of a table/view (or column). See the sync
|
||||
`Connection.lineage`. Returns a `Lineage`."""
|
||||
from .lineage import Lineage
|
||||
|
||||
raw = await self._inner.table_lineage(table, column, direction, depth)
|
||||
return Lineage.from_json(raw)
|
||||
|
||||
async def _refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Internal: submit a refresh, return the job id. The public surface is
|
||||
``AsyncMaterializedView.refresh()`` (returns an `AsyncJobHandle`).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh.
|
||||
"""
|
||||
return await self._inner.refresh_materialized_view(
|
||||
name,
|
||||
full=full,
|
||||
src_version=src_version,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
)
|
||||
|
||||
async def explain_refresh_materialized_view(
|
||||
self,
|
||||
name: str,
|
||||
*,
|
||||
full: bool = False,
|
||||
src_version: Optional[int] = None,
|
||||
):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return await self._inner.explain_refresh_materialized_view(
|
||||
name, full=full, src_version=src_version
|
||||
)
|
||||
|
||||
async def alter_materialized_view(self, name: str, *, auto_refresh: bool):
|
||||
"""Update a materialized view's options."""
|
||||
await self._inner.alter_materialized_view(name, auto_refresh)
|
||||
|
||||
async def drop_materialized_view(self, name: str):
|
||||
"""Drop a materialized view definition."""
|
||||
await self._inner.drop_materialized_view(name)
|
||||
|
||||
async def list_materialized_views(self):
|
||||
"""List registered materialized view definitions."""
|
||||
return await self._inner.list_materialized_views()
|
||||
|
||||
async def list_jobs(self):
|
||||
"""List inflight server-side jobs across the database's tables."""
|
||||
return await self._inner.list_jobs()
|
||||
|
||||
async def get_job(self, job_id: str, table: "str | None" = None):
|
||||
"""Look up one server-side job by id (the wait()/status poll path).
|
||||
``table`` (the job's table) enables an O(1) server-side lookup.
|
||||
Returns the job's status, or None if unknown / no longer active."""
|
||||
return await self._inner.get_job(job_id, table)
|
||||
|
||||
async def cancel_job(self, job_id: str) -> bool:
|
||||
"""Cancel an inflight server-side job by id (CANCEL JOB).
|
||||
|
||||
Returns True if a matching inflight job was found and flagged for
|
||||
cancellation, False otherwise (best-effort).
|
||||
"""
|
||||
return await self._inner.cancel_job(job_id)
|
||||
|
||||
async def job_history(self, job_id: "str | None" = None):
|
||||
"""Durable history of completed server-side jobs (SHOW JOB HISTORY).
|
||||
|
||||
Reads each table's durable job-history store. Pass ``job_id`` to narrow
|
||||
to a single job. Unlike :meth:`list_jobs` (live, inflight) these are the
|
||||
terminal records, with created/updated/completed timestamps.
|
||||
"""
|
||||
return await self._inner.job_history(job_id)
|
||||
|
||||
async def errors(self, job_id: "str | None" = None, table: "str | None" = None):
|
||||
"""Per-row UDF errors recorded by ``error_policy=skip`` (SHOW ERRORS).
|
||||
|
||||
Optionally filtered by ``job_id`` and/or ``table``.
|
||||
"""
|
||||
return await self._inner.errors(job_id, table)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
cur_name: str,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Union
|
||||
from typing import TYPE_CHECKING, List, Optional, Sequence, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -56,16 +56,6 @@ class OllamaEmbeddings(TextEmbeddingFunction):
|
||||
embeddings = self._compute_embedding(texts)
|
||||
return list(embeddings)
|
||||
|
||||
def __getstate__(self) -> dict[str, Any]:
|
||||
state = super().__getstate__()
|
||||
state["__dict__"] = {
|
||||
k: v for k, v in state["__dict__"].items() if k != "_ollama_client"
|
||||
}
|
||||
return state
|
||||
|
||||
def __setstate__(self, state: dict[str, Any]) -> None:
|
||||
super().__setstate__(state)
|
||||
|
||||
@cached_property
|
||||
def _ollama_client(self) -> "ollama.Client":
|
||||
ollama = attempt_import_or_raise("ollama")
|
||||
|
||||
@@ -19,8 +19,6 @@ operators::
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import date, datetime
|
||||
from decimal import Decimal
|
||||
from typing import Iterable, Union
|
||||
|
||||
import pyarrow as pa
|
||||
@@ -65,7 +63,7 @@ def _coerce(value: "ExprLike") -> "Expr":
|
||||
|
||||
|
||||
# Type alias used in annotations.
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes, date, datetime, Decimal]
|
||||
ExprLike = Union["Expr", bool, int, float, str, bytes]
|
||||
|
||||
|
||||
class Expr:
|
||||
@@ -120,18 +118,10 @@ class Expr:
|
||||
"""Logical AND (``expr_a & expr_b``)."""
|
||||
return Expr(self._inner.and_(_coerce(other)._inner))
|
||||
|
||||
def __rand__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical AND (``True & expr``)."""
|
||||
return Expr(_coerce(other)._inner.and_(self._inner))
|
||||
|
||||
def __or__(self, other: "Expr") -> "Expr":
|
||||
"""Logical OR (``expr_a | expr_b``)."""
|
||||
return Expr(self._inner.or_(_coerce(other)._inner))
|
||||
|
||||
def __ror__(self, other: ExprLike) -> "Expr":
|
||||
"""Right-hand logical OR (``False | expr``)."""
|
||||
return Expr(_coerce(other)._inner.or_(self._inner))
|
||||
|
||||
def __invert__(self) -> "Expr":
|
||||
"""Logical NOT (``~expr``)."""
|
||||
return Expr(self._inner.not_())
|
||||
@@ -276,14 +266,13 @@ def col(name: str) -> Expr:
|
||||
return Expr(expr_col(name))
|
||||
|
||||
|
||||
def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr:
|
||||
def lit(value: Union[bool, int, float, str, bytes]) -> Expr:
|
||||
"""Create a literal (constant) value expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value:
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``,
|
||||
``datetime``, or ``Decimal``.
|
||||
A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -291,9 +280,6 @@ def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) ->
|
||||
>>> col("price") * lit(1.1)
|
||||
Expr((price * 1.1))
|
||||
"""
|
||||
if not isinstance(value, (bool, int, float, str, bytes, date, datetime, Decimal)):
|
||||
raise TypeError(f"Unsupported literal type: {type(value).__name__}")
|
||||
|
||||
return Expr(expr_lit(value))
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""Client-side model of derived-compute lineage.
|
||||
|
||||
`Connection.lineage()` / `Table.lineage()` / `MaterializedView.lineage()` return
|
||||
a `Lineage`: the graph of what a column or materialized view derives from
|
||||
(upstream), what derives from it (downstream), and -- for each derived column --
|
||||
the function that produced it, the version it was produced with, and whether
|
||||
that is stale relative to the function the registry now holds.
|
||||
|
||||
The server returns this as JSON (the wire contract); these classes deserialize
|
||||
it. Nothing here talks to the server.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionRef:
|
||||
"""The function that produced a derived column, with version + location."""
|
||||
|
||||
name: str
|
||||
#: Version that produced the data (stamped at compute time), if known.
|
||||
as_computed_version: Optional[str] = None
|
||||
#: Version the registry currently holds for this function name.
|
||||
current_version: Optional[str] = None
|
||||
#: True when the column was produced by an older function than the registry
|
||||
#: now holds -- i.e. silently stale; re-refresh to catch up.
|
||||
stale_vs_current: bool = False
|
||||
language: Optional[str] = None
|
||||
docker_image: Optional[str] = None
|
||||
env_digest: Optional[str] = None
|
||||
code_uri: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "FunctionRef":
|
||||
return cls(
|
||||
name=d["name"],
|
||||
as_computed_version=d.get("as_computed_version"),
|
||||
current_version=d.get("current_version"),
|
||||
stale_vs_current=d.get("stale_vs_current", False),
|
||||
language=d.get("language"),
|
||||
docker_image=d.get("docker_image"),
|
||||
env_digest=d.get("env_digest"),
|
||||
code_uri=d.get("code_uri"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
"""A lineage node: a table, view, column, or function."""
|
||||
|
||||
kind: str # "table" | "view" | "column" | "function"
|
||||
id: str # "table", "table.column", or "fn:name@version"
|
||||
table: Optional[str] = None
|
||||
function: Optional[FunctionRef] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Node":
|
||||
fn = d.get("function")
|
||||
return cls(
|
||||
kind=d["kind"],
|
||||
id=d["id"],
|
||||
table=d.get("table"),
|
||||
function=FunctionRef._from(fn) if fn else None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Edge:
|
||||
"""`downstream` depends on `upstream`, produced by `via` (a function name,
|
||||
or None for a passthrough)."""
|
||||
|
||||
downstream: str
|
||||
upstream: str
|
||||
via: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def _from(cls, d: dict) -> "Edge":
|
||||
return cls(downstream=d["downstream"], upstream=d["upstream"], via=d.get("via"))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Lineage:
|
||||
"""A derived-compute lineage graph (nodes + labeled edges)."""
|
||||
|
||||
target: str
|
||||
nodes: List[Node] = field(default_factory=list)
|
||||
edges: List[Edge] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, raw: Union[str, bytes, dict]) -> "Lineage":
|
||||
d = json.loads(raw) if isinstance(raw, (str, bytes)) else raw
|
||||
return cls(
|
||||
target=d.get("target", ""),
|
||||
nodes=[Node._from(n) for n in d.get("nodes", [])],
|
||||
edges=[Edge._from(e) for e in d.get("edges", [])],
|
||||
)
|
||||
|
||||
def functions(self) -> List[FunctionRef]:
|
||||
"""The function nodes in the graph."""
|
||||
return [n.function for n in self.nodes if n.function is not None]
|
||||
|
||||
def stale(self) -> List[FunctionRef]:
|
||||
"""Functions whose as-computed version is behind the current registry
|
||||
version -- the columns they produced are silently out of date."""
|
||||
return [f for f in self.functions() if f.stale_vs_current]
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
def prune(d: dict) -> dict:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
return {
|
||||
"target": self.target,
|
||||
"nodes": [
|
||||
prune(
|
||||
{
|
||||
"kind": n.kind,
|
||||
"id": n.id,
|
||||
"table": n.table,
|
||||
"function": prune(vars(n.function)) if n.function else None,
|
||||
}
|
||||
)
|
||||
for n in self.nodes
|
||||
],
|
||||
"edges": [prune(vars(e)) for e in self.edges],
|
||||
}
|
||||
|
||||
def to_graphviz(self) -> str:
|
||||
"""Graphviz DOT for the lineage DAG: columns/tables as nodes, function
|
||||
names on edges, drift edges dashed + red."""
|
||||
stale_names = {f.name for f in self.stale()}
|
||||
out = [
|
||||
"digraph lineage {",
|
||||
" rankdir=LR;",
|
||||
' node [fontname="monospace"];',
|
||||
]
|
||||
for n in self.nodes:
|
||||
if n.kind == "function":
|
||||
continue
|
||||
shape = "ellipse" if n.kind in ("table", "view") else "box"
|
||||
out.append(f' "{n.id}" [shape={shape}];')
|
||||
for e in self.edges:
|
||||
attrs = ""
|
||||
if e.via:
|
||||
if e.via in stale_names:
|
||||
attrs = f' [label="{e.via}" color=red style=dashed]'
|
||||
else:
|
||||
attrs = f' [label="{e.via}"]'
|
||||
out.append(f' "{e.upstream}" -> "{e.downstream}"{attrs};')
|
||||
out.append("}")
|
||||
return "\n".join(out)
|
||||
|
||||
def _repr_html_(self) -> str:
|
||||
warn = ""
|
||||
drift = self.stale()
|
||||
if drift:
|
||||
names = ", ".join(sorted({f.name for f in drift}))
|
||||
warn = (
|
||||
f'<p style="color:#b00000"><b>stale vs current:</b> {names} '
|
||||
"(re-refresh to catch up)</p>"
|
||||
)
|
||||
rows = "".join(
|
||||
f"<tr><td><code>{e.downstream}</code></td>"
|
||||
f"<td>← {e.via or ''}</td>"
|
||||
f"<td><code>{e.upstream}</code></td></tr>"
|
||||
for e in self.edges
|
||||
)
|
||||
return (
|
||||
f"<b>lineage: <code>{self.target}</code></b>{warn}"
|
||||
"<table><tr><th>derived</th><th>via</th><th>from</th></tr>"
|
||||
f"{rows}</table>"
|
||||
)
|
||||
@@ -51,15 +51,6 @@ class LanceMergeInsertBuilder(object):
|
||||
If there are multiple matches then the behavior is undefined.
|
||||
Currently this causes multiple copies of the row to be created
|
||||
but that behavior is subject to change.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: Optional[str], default None
|
||||
An optional filter to limit which rows are updated. Column
|
||||
references in this expression must be prefixed with "target."
|
||||
to refer to the existing table data. For example, to only
|
||||
update rows where the existing color is red, use:
|
||||
``where="target.color = 'red'"``
|
||||
"""
|
||||
self._when_matched_update_all = True
|
||||
self._when_matched_update_all_condition = where
|
||||
|
||||
+178
-225
@@ -38,13 +38,15 @@ from lance_namespace_urllib3_client.models.query_table_request_vector import (
|
||||
QueryTableRequestVector,
|
||||
)
|
||||
from lance_namespace_urllib3_client.models.string_fts_query import StringFtsQuery
|
||||
from lance_namespace.errors import NamespaceNotEmptyError, TableNotFoundError
|
||||
from lancedb._lancedb import (
|
||||
connect_namespace as _connect_namespace,
|
||||
connect_namespace_client as _connect_namespace_client,
|
||||
)
|
||||
from lance_namespace.errors import TableNotFoundError
|
||||
from lancedb._lancedb import connect_namespace_client as _connect_namespace_client
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import AsyncConnection, DBConnection
|
||||
from lancedb.namespace_utils import (
|
||||
_normalize_create_namespace_mode,
|
||||
_normalize_drop_namespace_mode,
|
||||
_normalize_drop_namespace_behavior,
|
||||
)
|
||||
from lance_namespace import (
|
||||
LanceNamespace,
|
||||
connect as namespace_connect,
|
||||
@@ -53,6 +55,13 @@ from lance_namespace import (
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
ListTablesRequest,
|
||||
DescribeNamespaceRequest,
|
||||
DropTableRequest,
|
||||
RenameTableRequest,
|
||||
ListNamespacesRequest,
|
||||
CreateNamespaceRequest,
|
||||
DropNamespaceRequest,
|
||||
)
|
||||
from lancedb.table import AsyncTable, LanceTable, Table
|
||||
from lancedb.util import validate_table_name
|
||||
@@ -377,10 +386,6 @@ def _builds_namespace_natively(
|
||||
return namespace_client_impl == "rest" and bool(namespace_client_properties)
|
||||
|
||||
|
||||
def _supports_native_namespace(namespace_client_impl: str) -> bool:
|
||||
return namespace_client_impl in {"dir", "rest"}
|
||||
|
||||
|
||||
class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
A LanceDB connection that uses a namespace for table management.
|
||||
@@ -391,7 +396,7 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace_client: Optional[LanceNamespace] = None,
|
||||
namespace_client: LanceNamespace,
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
@@ -399,7 +404,6 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
):
|
||||
"""
|
||||
Initialize a namespace-based LanceDB connection.
|
||||
@@ -441,36 +445,30 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# When the namespace connection or client is built natively in Rust, the
|
||||
# underlying Rust table performs QueryTable pushdown through the
|
||||
# read-freshness context provider, which the pure-Python ``query_table``
|
||||
# path bypasses.
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
# When the namespace client is built natively (see Rust
|
||||
# ``build_namespace_natively``), the underlying Rust table performs
|
||||
# QueryTable pushdown through the read-freshness context provider, which
|
||||
# the pure-Python ``query_table`` path bypasses.
|
||||
self._route_pushdown_to_rust = _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
else:
|
||||
if namespace_client is None:
|
||||
raise ValueError("namespace_client is required without a native _inner")
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
self._uri = self._inner.uri
|
||||
)
|
||||
|
||||
@override
|
||||
def serialize(self) -> str:
|
||||
@@ -516,11 +514,11 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
)
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
self._inner.table_names(
|
||||
namespace_path=namespace_path, start_after=page_token, limit=limit
|
||||
)
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return response.tables if response.tables else []
|
||||
|
||||
@override
|
||||
def create_table(
|
||||
@@ -591,8 +589,8 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
)
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if "Table not found" in str(e) or "was not found" in str(e):
|
||||
except RuntimeError as e:
|
||||
if "Table not found" in str(e):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
|
||||
raise
|
||||
@@ -614,9 +612,12 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
|
||||
@override
|
||||
def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
|
||||
# Use namespace drop_table directly
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
LOOP.run(self._inner.drop_table(name, namespace_path=namespace_path))
|
||||
table_id = namespace_path + [name]
|
||||
request = DropTableRequest(id=table_id)
|
||||
self._namespace_client.drop_table(request)
|
||||
|
||||
@override
|
||||
def rename_table(
|
||||
@@ -630,19 +631,14 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
cur_namespace_path = []
|
||||
if new_namespace_path is None:
|
||||
new_namespace_path = []
|
||||
try:
|
||||
LOOP.run(
|
||||
self._inner.rename_table(
|
||||
cur_name,
|
||||
new_name,
|
||||
cur_namespace_path=cur_namespace_path,
|
||||
new_namespace_path=new_namespace_path,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "rename_table not implemented" in str(e):
|
||||
raise NotImplementedError("rename_table not implemented") from e
|
||||
raise
|
||||
cur_table_id = cur_namespace_path + [cur_name]
|
||||
new_namespace_id = new_namespace_path if new_namespace_path else None
|
||||
request = RenameTableRequest(
|
||||
id=cur_table_id,
|
||||
new_table_name=new_name,
|
||||
new_namespace_id=new_namespace_id,
|
||||
)
|
||||
self._namespace_client.rename_table(request)
|
||||
|
||||
@override
|
||||
def drop_database(self):
|
||||
@@ -654,7 +650,8 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
def drop_all_tables(self, namespace_path: Optional[List[str]] = None):
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
LOOP.run(self._inner.drop_all_tables(namespace_path=namespace_path))
|
||||
for table_name in self.table_names(namespace_path=namespace_path):
|
||||
self.drop_table(table_name, namespace_path=namespace_path)
|
||||
|
||||
@override
|
||||
def list_namespaces(
|
||||
@@ -684,10 +681,13 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
request = ListNamespacesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_namespaces(request)
|
||||
return ListNamespacesResponse(
|
||||
namespaces=response.namespaces if response.namespaces else [],
|
||||
page_token=response.page_token,
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -715,12 +715,14 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
CreateNamespaceResponse
|
||||
Response containing the properties of the created namespace.
|
||||
"""
|
||||
return LOOP.run(
|
||||
self._inner.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
properties=properties,
|
||||
)
|
||||
request = CreateNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_create_namespace_mode(mode),
|
||||
properties=properties,
|
||||
)
|
||||
response = self._namespace_client.create_namespace(request)
|
||||
return CreateNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
@override
|
||||
@@ -748,18 +750,20 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
DropNamespaceResponse
|
||||
Response containing properties and transaction_id if applicable.
|
||||
"""
|
||||
try:
|
||||
return LOOP.run(
|
||||
self._inner.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
request = DropNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_drop_namespace_mode(mode),
|
||||
behavior=_normalize_drop_namespace_behavior(behavior),
|
||||
)
|
||||
response = self._namespace_client.drop_namespace(request)
|
||||
return DropNamespaceResponse(
|
||||
properties=(
|
||||
response.properties if hasattr(response, "properties") else None
|
||||
),
|
||||
transaction_id=(
|
||||
response.transaction_id if hasattr(response, "transaction_id") else None
|
||||
),
|
||||
)
|
||||
|
||||
@override
|
||||
def describe_namespace(
|
||||
@@ -778,7 +782,11 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
DescribeNamespaceResponse
|
||||
Response containing the namespace properties.
|
||||
"""
|
||||
return LOOP.run(self._inner.describe_namespace(namespace_path))
|
||||
request = DescribeNamespaceRequest(id=namespace_path)
|
||||
response = self._namespace_client.describe_namespace(request)
|
||||
return DescribeNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
@override
|
||||
def list_tables(
|
||||
@@ -808,10 +816,13 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return LOOP.run(
|
||||
self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return ListTablesResponse(
|
||||
tables=response.tables if response.tables else [],
|
||||
page_token=response.page_token,
|
||||
)
|
||||
|
||||
def _lance_table_from_uri(
|
||||
@@ -867,18 +878,6 @@ class LanceNamespaceDBConnection(DBConnection):
|
||||
LanceNamespace
|
||||
The namespace client for this connection.
|
||||
"""
|
||||
if self._namespace_client is None:
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot construct a Python namespace client without "
|
||||
"namespace implementation properties"
|
||||
)
|
||||
self._namespace_client = namespace_connect(
|
||||
self._namespace_client_impl, self._namespace_client_properties
|
||||
)
|
||||
return self._namespace_client
|
||||
|
||||
|
||||
@@ -892,7 +891,7 @@ class AsyncLanceNamespaceDBConnection:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
namespace_client: Optional[LanceNamespace] = None,
|
||||
namespace_client: LanceNamespace,
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
@@ -900,7 +899,6 @@ class AsyncLanceNamespaceDBConnection:
|
||||
namespace_client_pushdown_operations: Optional[List[str]] = None,
|
||||
namespace_client_impl: Optional[str] = None,
|
||||
namespace_client_properties: Optional[Dict[str, str]] = None,
|
||||
_inner: Optional[AsyncConnection] = None,
|
||||
):
|
||||
"""
|
||||
Initialize an async namespace-based LanceDB connection.
|
||||
@@ -942,35 +940,29 @@ class AsyncLanceNamespaceDBConnection:
|
||||
)
|
||||
self._namespace_client_impl = namespace_client_impl
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
# See LanceNamespaceDBConnection: when Rust owns the namespace
|
||||
# connection/client, its table performs QueryTable pushdown through the
|
||||
# read-freshness provider, so defer to it rather than the urllib3 client
|
||||
# path (which omits x-lancedb-min-timestamp).
|
||||
self._route_pushdown_to_rust = _inner is not None or _builds_namespace_natively(
|
||||
# See LanceNamespaceDBConnection: when built natively the Rust table runs
|
||||
# QueryTable pushdown through the read-freshness provider, so defer to it
|
||||
# rather than the urllib3 client (which omits x-lancedb-min-timestamp).
|
||||
self._route_pushdown_to_rust = _builds_namespace_natively(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
if _inner is not None:
|
||||
self._inner = _inner
|
||||
else:
|
||||
if namespace_client is None:
|
||||
raise ValueError("namespace_client is required without a native _inner")
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
self._inner = AsyncConnection(
|
||||
_connect_namespace_client(
|
||||
namespace_client,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=self.storage_options or None,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=(
|
||||
list(self._namespace_client_pushdown_operations)
|
||||
),
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
)
|
||||
)
|
||||
|
||||
async def table_names(
|
||||
self,
|
||||
@@ -994,9 +986,11 @@ class AsyncLanceNamespaceDBConnection:
|
||||
)
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.table_names(
|
||||
namespace_path=namespace_path, start_after=page_token, limit=limit
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return response.tables if response.tables else []
|
||||
|
||||
async def create_table(
|
||||
self,
|
||||
@@ -1059,8 +1053,8 @@ class AsyncLanceNamespaceDBConnection:
|
||||
storage_options=storage_options,
|
||||
index_cache_size=index_cache_size,
|
||||
)
|
||||
except (RuntimeError, ValueError) as e:
|
||||
if "Table not found" in str(e) or "was not found" in str(e):
|
||||
except RuntimeError as e:
|
||||
if "Table not found" in str(e):
|
||||
table_id = namespace_path + [name]
|
||||
raise TableNotFoundError(f"Table not found: {'$'.join(table_id)}")
|
||||
raise
|
||||
@@ -1081,7 +1075,9 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""Drop a table from the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
await self._inner.drop_table(name, namespace_path=namespace_path)
|
||||
table_id = namespace_path + [name]
|
||||
request = DropTableRequest(id=table_id)
|
||||
self._namespace_client.drop_table(request)
|
||||
|
||||
async def rename_table(
|
||||
self,
|
||||
@@ -1095,17 +1091,14 @@ class AsyncLanceNamespaceDBConnection:
|
||||
cur_namespace_path = []
|
||||
if new_namespace_path is None:
|
||||
new_namespace_path = []
|
||||
try:
|
||||
await self._inner.rename_table(
|
||||
cur_name,
|
||||
new_name,
|
||||
cur_namespace_path=cur_namespace_path,
|
||||
new_namespace_path=new_namespace_path,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "rename_table not implemented" in str(e):
|
||||
raise NotImplementedError("rename_table not implemented") from e
|
||||
raise
|
||||
cur_table_id = cur_namespace_path + [cur_name]
|
||||
new_namespace_id = new_namespace_path if new_namespace_path else None
|
||||
request = RenameTableRequest(
|
||||
id=cur_table_id,
|
||||
new_table_name=new_name,
|
||||
new_namespace_id=new_namespace_id,
|
||||
)
|
||||
self._namespace_client.rename_table(request)
|
||||
|
||||
async def drop_database(self):
|
||||
"""Deprecated method."""
|
||||
@@ -1117,7 +1110,9 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""Drop all tables in the namespace."""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
await self._inner.drop_all_tables(namespace_path=namespace_path)
|
||||
table_names = await self.table_names(namespace_path=namespace_path)
|
||||
for table_name in table_names:
|
||||
await self.drop_table(table_name, namespace_path=namespace_path)
|
||||
|
||||
async def list_namespaces(
|
||||
self,
|
||||
@@ -1146,8 +1141,13 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.list_namespaces(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
request = ListNamespacesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_namespaces(request)
|
||||
return ListNamespacesResponse(
|
||||
namespaces=response.namespaces if response.namespaces else [],
|
||||
page_token=response.page_token,
|
||||
)
|
||||
|
||||
async def create_namespace(
|
||||
@@ -1174,11 +1174,15 @@ class AsyncLanceNamespaceDBConnection:
|
||||
CreateNamespaceResponse
|
||||
Response containing the properties of the created namespace.
|
||||
"""
|
||||
return await self._inner.create_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
request = CreateNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_create_namespace_mode(mode),
|
||||
properties=properties,
|
||||
)
|
||||
response = self._namespace_client.create_namespace(request)
|
||||
return CreateNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
async def drop_namespace(
|
||||
self,
|
||||
@@ -1204,16 +1208,20 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DropNamespaceResponse
|
||||
Response containing properties and transaction_id if applicable.
|
||||
"""
|
||||
try:
|
||||
return await self._inner.drop_namespace(
|
||||
namespace_path=namespace_path,
|
||||
mode=mode,
|
||||
behavior=behavior,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Namespace not empty" in str(e):
|
||||
raise NamespaceNotEmptyError(str(e)) from e
|
||||
raise
|
||||
request = DropNamespaceRequest(
|
||||
id=namespace_path,
|
||||
mode=_normalize_drop_namespace_mode(mode),
|
||||
behavior=_normalize_drop_namespace_behavior(behavior),
|
||||
)
|
||||
response = self._namespace_client.drop_namespace(request)
|
||||
return DropNamespaceResponse(
|
||||
properties=(
|
||||
response.properties if hasattr(response, "properties") else None
|
||||
),
|
||||
transaction_id=(
|
||||
response.transaction_id if hasattr(response, "transaction_id") else None
|
||||
),
|
||||
)
|
||||
|
||||
async def describe_namespace(
|
||||
self, namespace_path: List[str]
|
||||
@@ -1231,7 +1239,11 @@ class AsyncLanceNamespaceDBConnection:
|
||||
DescribeNamespaceResponse
|
||||
Response containing the namespace properties.
|
||||
"""
|
||||
return await self._inner.describe_namespace(namespace_path)
|
||||
request = DescribeNamespaceRequest(id=namespace_path)
|
||||
response = self._namespace_client.describe_namespace(request)
|
||||
return DescribeNamespaceResponse(
|
||||
properties=response.properties if hasattr(response, "properties") else None
|
||||
)
|
||||
|
||||
async def list_tables(
|
||||
self,
|
||||
@@ -1260,8 +1272,13 @@ class AsyncLanceNamespaceDBConnection:
|
||||
"""
|
||||
if namespace_path is None:
|
||||
namespace_path = []
|
||||
return await self._inner.list_tables(
|
||||
namespace_path=namespace_path, page_token=page_token, limit=limit
|
||||
request = ListTablesRequest(
|
||||
id=namespace_path, page_token=page_token, limit=limit
|
||||
)
|
||||
response = self._namespace_client.list_tables(request)
|
||||
return ListTablesResponse(
|
||||
tables=response.tables if response.tables else [],
|
||||
page_token=response.page_token,
|
||||
)
|
||||
|
||||
async def namespace_client(self) -> LanceNamespace:
|
||||
@@ -1275,18 +1292,6 @@ class AsyncLanceNamespaceDBConnection:
|
||||
LanceNamespace
|
||||
The namespace client for this connection.
|
||||
"""
|
||||
if self._namespace_client is None:
|
||||
if (
|
||||
self._namespace_client_impl is None
|
||||
or self._namespace_client_properties is None
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot construct a Python namespace client without "
|
||||
"namespace implementation properties"
|
||||
)
|
||||
self._namespace_client = namespace_connect(
|
||||
self._namespace_client_impl, self._namespace_client_properties
|
||||
)
|
||||
return self._namespace_client
|
||||
|
||||
|
||||
@@ -1337,32 +1342,6 @@ def connect_namespace(
|
||||
LanceNamespaceDBConnection
|
||||
A namespace-based connection to LanceDB
|
||||
"""
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
)
|
||||
)
|
||||
return LanceNamespaceDBConnection(
|
||||
namespace_client=None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
@@ -1438,32 +1417,6 @@ def connect_namespace_async(
|
||||
... tables = await db.table_names()
|
||||
... table = await db.create_table("my_table", schema=schema)
|
||||
"""
|
||||
if _supports_native_namespace(namespace_client_impl):
|
||||
inner = AsyncConnection(
|
||||
_connect_namespace(
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=(
|
||||
read_consistency_interval.total_seconds()
|
||||
if read_consistency_interval is not None
|
||||
else None
|
||||
),
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
)
|
||||
)
|
||||
return AsyncLanceNamespaceDBConnection(
|
||||
namespace_client=None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
namespace_client_pushdown_operations=namespace_client_pushdown_operations,
|
||||
namespace_client_impl=namespace_client_impl,
|
||||
namespace_client_properties=namespace_client_properties,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
namespace_client = namespace_connect(
|
||||
namespace_client_impl, namespace_client_properties
|
||||
)
|
||||
|
||||
@@ -11,7 +11,7 @@ import pyarrow as pa
|
||||
from ._lancedb import async_permutation_builder, PermutationReader
|
||||
from .table import LanceTable, Table
|
||||
from .background_loop import LOOP
|
||||
from .util import batch_to_tensor, batch_to_tensor_dict, batch_to_tensor_rows
|
||||
from .util import batch_to_tensor, batch_to_tensor_rows
|
||||
from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -65,7 +65,6 @@ class PermutationBuilder:
|
||||
counts: Optional[list[int]] = None,
|
||||
fixed: Optional[int] = None,
|
||||
seed: Optional[int] = None,
|
||||
clump_size: Optional[int] = None,
|
||||
split_names: Optional[list[str]] = None,
|
||||
) -> "PermutationBuilder":
|
||||
"""
|
||||
@@ -88,9 +87,6 @@ class PermutationBuilder:
|
||||
Rows will be randomly assigned to splits. The optional seed can be provided to
|
||||
make the assignment deterministic.
|
||||
|
||||
If clump_size is provided, rows are shuffled as contiguous groups of that size,
|
||||
preserving I/O locality while still randomising the split assignment.
|
||||
|
||||
The optional split_names can be provided to name the splits. If not provided,
|
||||
the splits can only be referenced by their index.
|
||||
"""
|
||||
@@ -99,7 +95,6 @@ class PermutationBuilder:
|
||||
counts=counts,
|
||||
fixed=fixed,
|
||||
seed=seed,
|
||||
clump_size=clump_size,
|
||||
split_names=split_names,
|
||||
)
|
||||
return self
|
||||
@@ -946,7 +941,6 @@ class Permutation:
|
||||
"pandas",
|
||||
"arrow",
|
||||
"torch",
|
||||
"torch_row",
|
||||
"torch_col",
|
||||
"polars",
|
||||
],
|
||||
@@ -962,19 +956,15 @@ class Permutation:
|
||||
- "python_col" - the batch will be a dict of lists (one entry per column)
|
||||
- "pandas" - the batch will be a pandas DataFrame
|
||||
- "arrow" - the batch will be a pyarrow RecordBatch
|
||||
- "torch" - the batch will be a list of per-row dicts mapping column
|
||||
name to a 0-D torch tensor. Works with the default
|
||||
``torch.utils.data.DataLoader`` collate, which stacks the per-row
|
||||
dicts back into a dict of batched tensors.
|
||||
- "torch_row" - the batch will be a list of tensors, one per row
|
||||
- "torch" - the batch will be a list of tensors, one per row
|
||||
- "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns)
|
||||
- "polars" - the batch will be a polars DataFrame
|
||||
|
||||
Conversion may or may not involve a data copy. Lance uses Arrow internally
|
||||
and so it is able to zero-copy to the arrow and polars formats.
|
||||
|
||||
Conversion to torch and torch_col will be zero-copy but will only support a
|
||||
subset of data types (numeric types).
|
||||
Conversion to torch_col will be zero-copy but will only support a subset of data
|
||||
types (numeric types).
|
||||
|
||||
Conversion to numpy and/or pandas will typically be zero-copy for numeric
|
||||
types. Conversion of strings, lists, and structs will require creating python
|
||||
@@ -995,8 +985,6 @@ class Permutation:
|
||||
elif format == "arrow":
|
||||
return self.with_transform(Transforms.arrow2arrow)
|
||||
elif format == "torch":
|
||||
return self.with_transform(batch_to_tensor_dict)
|
||||
elif format == "torch_row":
|
||||
return self.with_transform(batch_to_tensor_rows)
|
||||
elif format == "torch_col":
|
||||
return self.with_transform(batch_to_tensor)
|
||||
|
||||
@@ -119,27 +119,6 @@ def _filter_to_sql(filter: Optional[Union[str, Expr]]) -> Optional[str]:
|
||||
return filter
|
||||
|
||||
|
||||
def _combine_where(
|
||||
existing: Optional[Union[str, Expr]], new: Union[str, Expr]
|
||||
) -> Union[str, Expr]:
|
||||
"""Combine a new filter with an existing one using a logical AND.
|
||||
|
||||
Calling ``where`` more than once composes the filters with AND instead of
|
||||
replacing the previous filter. Two :class:`~lancedb.expr.Expr` filters are
|
||||
combined as an expression; otherwise both filters are lowered to SQL strings
|
||||
and combined as SQL.
|
||||
"""
|
||||
if existing is None:
|
||||
return new
|
||||
existing_is_expr = isinstance(existing, Expr)
|
||||
new_is_expr = isinstance(new, Expr)
|
||||
if existing_is_expr and new_is_expr:
|
||||
return existing & new
|
||||
existing_sql = existing.to_sql() if existing_is_expr else existing
|
||||
new_sql = new.to_sql() if new_is_expr else new
|
||||
return f"({existing_sql}) AND ({new_sql})"
|
||||
|
||||
|
||||
def _projection_to_scanner_kwargs(
|
||||
columns: Optional[
|
||||
Union[
|
||||
@@ -1169,13 +1148,8 @@ class LanceQueryBuilder(ABC):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = _combine_where(self._where, where)
|
||||
self._where = where
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
|
||||
@@ -1719,13 +1693,8 @@ class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||||
-------
|
||||
LanceQueryBuilder
|
||||
The LanceQueryBuilder object.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
self._where = _combine_where(self._where, where)
|
||||
self._where = where
|
||||
if prefilter is not None:
|
||||
self._postfilter = not prefilter
|
||||
return self
|
||||
@@ -2925,9 +2894,6 @@ class AsyncStandardQuery(AsyncQueryBase):
|
||||
|
||||
Filtering performance can often be improved by creating a scalar index
|
||||
on the filter column(s).
|
||||
|
||||
Calling this multiple times combines the filters with a logical AND
|
||||
rather than replacing the previous filter.
|
||||
"""
|
||||
if isinstance(predicate, Expr):
|
||||
self._inner.where_expr(predicate._inner)
|
||||
|
||||
@@ -13,10 +13,14 @@ from typing import (
|
||||
Iterable,
|
||||
List,
|
||||
Optional,
|
||||
TYPE_CHECKING,
|
||||
Union,
|
||||
Literal,
|
||||
overload,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..udf import JobHandle
|
||||
import warnings
|
||||
|
||||
from lancedb import __version__
|
||||
@@ -884,8 +888,142 @@ class RemoteTable(Table):
|
||||
def count_rows(self, filter: Optional[str] = None) -> int:
|
||||
return LOOP.run(self._table.count_rows(filter))
|
||||
|
||||
def add_columns(self, transforms: Dict[str, str]) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
def add_columns(
|
||||
self,
|
||||
transforms: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
computed: Optional[Dict[str, tuple]] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
LOOP.run(self._table.add_columns(computed=computed))
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "JobHandle":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `JobHandle` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("c").wait()``). Server-backed feature
|
||||
(LanceDB Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from ..udf import JobHandle
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return JobHandle(self._job_conn(), job_id)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Derived-compute lineage of this table, or one of its columns:
|
||||
upstream sources, downstream dependents, and the function version +
|
||||
location that produced each derived column (with a drift flag). Returns
|
||||
a `Lineage`. See `Connection.lineage`."""
|
||||
return self._job_conn().lineage(
|
||||
self._name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
def _job_conn(self):
|
||||
"""A client connection for polling jobs this table spawns. Built lazily
|
||||
from the table's serialized connection state and cached (not pickled --
|
||||
a forked/unpickled table rebuilds it on next use)."""
|
||||
from lancedb import deserialize_conn
|
||||
|
||||
conn = getattr(self, "_job_conn_cache", None)
|
||||
if conn is None:
|
||||
conn = deserialize_conn(self._serialized_connection_state())
|
||||
self._job_conn_cache = conn
|
||||
return conn
|
||||
|
||||
def load_columns(
|
||||
self,
|
||||
source: Union[str, Iterable[str]],
|
||||
pk: str,
|
||||
columns: Union[Iterable[str], Dict[str, str]],
|
||||
*,
|
||||
source_format: str = "parquet",
|
||||
source_pk: Optional[str] = None,
|
||||
on_missing: str = "carry",
|
||||
source_storage_options: Optional[Dict[str, str]] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
commit_granularity: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Fill existing columns from an external source by primary-key join.
|
||||
|
||||
The distributed-job equivalent of Geneva's ``Table.load_columns()``:
|
||||
imports precomputed values (e.g. embeddings) from Parquet/Lance/IPC into
|
||||
this table, matching on a primary key. Returns the load job id.
|
||||
Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source: str | list[str]
|
||||
One source URI or a list of URIs.
|
||||
pk: str
|
||||
Destination primary-key column. Also the source key unless
|
||||
``source_pk`` is given.
|
||||
columns: list[str] | dict[str, str]
|
||||
Value columns to load. A list loads same-named columns; a dict maps
|
||||
``{target: source}``.
|
||||
source_format: str
|
||||
``"parquet"`` (default), ``"lance"``, or ``"ipc"``.
|
||||
source_pk: str, optional
|
||||
Source primary-key column when it differs from ``pk``.
|
||||
on_missing: str
|
||||
Behavior for destination rows with no source match:
|
||||
``"carry"`` (default, keep existing), ``"null"``, or ``"error"``.
|
||||
"""
|
||||
if isinstance(source, str):
|
||||
source = [source]
|
||||
if isinstance(columns, dict):
|
||||
mappings = [(target, src) for target, src in columns.items()]
|
||||
else:
|
||||
mappings = [(c, None) for c in columns]
|
||||
return LOOP.run(
|
||||
self._table.load_columns(
|
||||
list(source),
|
||||
source_format,
|
||||
pk,
|
||||
mappings,
|
||||
source_key=source_pk,
|
||||
source_storage_options=source_storage_options,
|
||||
on_missing=on_missing,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
commit_granularity=commit_granularity,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
@@ -912,10 +1050,6 @@ class RemoteTable(Table):
|
||||
"""Not supported on LanceDB Cloud."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""No-op on LanceDB Cloud (no local shard writers)."""
|
||||
return LOOP.run(self._table.close_lsm_writers())
|
||||
|
||||
@@ -156,16 +156,9 @@ class MRRReranker(Reranker):
|
||||
reciprocal_rank = 1.0 / rank
|
||||
mrr_score_map[result_id].append(reciprocal_rank)
|
||||
|
||||
# MRR averages the reciprocal rank across *all* ranking systems, treating
|
||||
# a system in which a document does not appear as a reciprocal rank of 0.
|
||||
# We therefore divide by the total number of systems, not by the number of
|
||||
# systems the document happens to appear in -- otherwise a document found
|
||||
# by a single ranking would outrank one ranked highly by every system,
|
||||
# defeating the purpose of fusing the rankings.
|
||||
num_systems = len(vector_results)
|
||||
final_mrr_scores = {}
|
||||
for result_id, reciprocal_ranks in mrr_score_map.items():
|
||||
mean_rr = float(np.sum(reciprocal_ranks)) / num_systems
|
||||
mean_rr = np.mean(reciprocal_ranks)
|
||||
final_mrr_scores[result_id] = mean_rr
|
||||
|
||||
combined = pa.concat_tables(vector_results, **self._concat_tables_args)
|
||||
|
||||
@@ -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)
|
||||
+236
-52
@@ -651,16 +651,6 @@ def _append_vector_columns(
|
||||
col_data = func.compute_source_embeddings_with_retry(
|
||||
batch[conf.source_column]
|
||||
)
|
||||
# Replace vectors with wrong length (including empty lists
|
||||
# returned for inputs like empty strings) with None so that
|
||||
# _handle_bad_vectors can process them according to the
|
||||
# on_bad_vectors policy instead of crashing when PyArrow
|
||||
# tries to cast them into a fixed-size list array.
|
||||
expected_ndims = conf.function.ndims()
|
||||
col_data = [
|
||||
v if v is not None and len(v) == expected_ndims else None
|
||||
for v in col_data
|
||||
]
|
||||
if no_vector_column:
|
||||
batch = batch.append_column(
|
||||
schema.field(vector_column),
|
||||
@@ -712,6 +702,24 @@ def _normalize_progress(progress):
|
||||
return progress, False
|
||||
|
||||
|
||||
def _computed_groups(computed):
|
||||
"""Group computed columns by expression, preserving declaration order
|
||||
(struct-returning functions need their columns adjacent so schema order
|
||||
matches field order). Accepts the ergonomic forms -- `fn("col")` values
|
||||
and tuple keys for struct fan-out -- via `_normalize_computed`."""
|
||||
from .udf import _normalize_computed
|
||||
|
||||
groups = []
|
||||
for name, (sql_type, expression) in _normalize_computed(computed).items():
|
||||
for expr, cols in groups:
|
||||
if expr == expression:
|
||||
cols.append((name, sql_type))
|
||||
break
|
||||
else:
|
||||
groups.append((expression, [(name, sql_type)]))
|
||||
return groups
|
||||
|
||||
|
||||
class Table(ABC):
|
||||
"""
|
||||
A Table is a collection of Records in a LanceDB Database.
|
||||
@@ -817,6 +825,59 @@ class Table(ABC):
|
||||
"""The number of rows in this Table"""
|
||||
return self.count_rows(None)
|
||||
|
||||
def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF -- no compute happens
|
||||
here (the agent fills them lazily, or refresh_column() triggers a run).
|
||||
|
||||
.. deprecated::
|
||||
A computed column is an expression over a registered function, so
|
||||
bind it as one: ``add_columns(computed={"vec": embed("data")})``.
|
||||
``embed("data")`` applies the function to the `data` column and
|
||||
infers the type from the function's return signature -- the
|
||||
function never couples to a particular column. Prefer that form.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"vec": embed("data")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
else:
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
self.add_columns(computed=computed)
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def embedding_functions(self) -> Dict[str, EmbeddingFunctionConfig]:
|
||||
@@ -2152,19 +2213,12 @@ class LanceTable(Table):
|
||||
|
||||
branch = self.current_branch()
|
||||
version = None if branch is not None else self.version
|
||||
namespace_client = self._namespace_client
|
||||
if namespace_client is None:
|
||||
conn_uri = getattr(self._conn, "uri", "")
|
||||
if get_uri_scheme(conn_uri) == "namespace":
|
||||
namespace_client = self._conn.namespace_client()
|
||||
self._namespace_client = namespace_client
|
||||
|
||||
if namespace_client is not None:
|
||||
if self._namespace_client is not None:
|
||||
table_id = self._namespace_path + [self.name]
|
||||
ds = lance.dataset(
|
||||
version=version,
|
||||
storage_options=self._conn.storage_options,
|
||||
namespace_client=namespace_client,
|
||||
namespace_client=self._namespace_client,
|
||||
table_id=table_id,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -3727,9 +3781,68 @@ class LanceTable(Table):
|
||||
return LOOP.run(self._table.index_stats(index_name))
|
||||
|
||||
def add_columns(
|
||||
self, transforms: Dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
return LOOP.run(self._table.add_columns(transforms))
|
||||
self,
|
||||
transforms: Dict[str, str]
|
||||
| pa.field
|
||||
| List[pa.field]
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
result = None
|
||||
if transforms is not None:
|
||||
result = LOOP.run(self._table.add_columns(transforms))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
result_unused = LOOP.run(self._table.add_columns(computed=computed))
|
||||
del result_unused
|
||||
return result
|
||||
|
||||
def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> "JobHandle":
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
|
||||
The expression is resolved server-side from each column's stored
|
||||
binding; columns bound to the same struct-returning function
|
||||
refresh together. Returns a `JobHandle` to wait on, poll, or cancel
|
||||
(``tbl.refresh_column("col").wait()``) -- mirrors
|
||||
`MaterializedView.refresh()`. Server-backed feature (LanceDB
|
||||
Enterprise / Cloud).
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh) and override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill).
|
||||
"""
|
||||
from .udf import JobHandle
|
||||
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
job_id = LOOP.run(
|
||||
self._table.refresh_column(
|
||||
list(columns),
|
||||
where=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
)
|
||||
return JobHandle(self._conn, job_id, table=self.name)
|
||||
|
||||
def alter_columns(
|
||||
self, *alterations: Iterable[Dict[str, str]]
|
||||
@@ -3759,11 +3872,6 @@ class LanceTable(Table):
|
||||
[`AsyncTable.unset_lsm_write_spec`][lancedb.AsyncTable.unset_lsm_write_spec]."""
|
||||
return LOOP.run(self._table.unset_lsm_write_spec())
|
||||
|
||||
def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the installed LsmWriteSpec, or ``None``. See
|
||||
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
|
||||
return LOOP.run(self._table.get_lsm_write_spec())
|
||||
|
||||
def close_lsm_writers(self) -> None:
|
||||
"""Close cached MemWAL shard writers. See
|
||||
[`AsyncTable.close_lsm_writers`][lancedb.AsyncTable.close_lsm_writers]."""
|
||||
@@ -4035,16 +4143,7 @@ def _handle_bad_vector_column(
|
||||
dim = _infer_vector_dim(vec_arr)
|
||||
if dim is None:
|
||||
return data
|
||||
|
||||
is_null = pc.is_null(vec_arr)
|
||||
# pc.list_value_length returns null for null list entries, so
|
||||
# pc.not_equal(null, dim) also returns null. Use or_kleene so that
|
||||
# True OR null = True (Kleene three-valued logic), ensuring null vectors
|
||||
# are counted as wrong-dim.
|
||||
has_wrong_dim = pc.or_kleene(
|
||||
is_null,
|
||||
pc.not_equal(pc.list_value_length(vec_arr), dim),
|
||||
)
|
||||
has_wrong_dim = pc.not_equal(pc.list_value_length(vec_arr), dim)
|
||||
|
||||
has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py()
|
||||
|
||||
@@ -4422,17 +4521,6 @@ class AsyncTable:
|
||||
"""
|
||||
await self._inner.unset_lsm_write_spec()
|
||||
|
||||
async def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
|
||||
"""Read the LsmWriteSpec currently installed on this table.
|
||||
|
||||
Returns ``None`` when the MemWAL LSM write path is not enabled (no
|
||||
spec has been set, or it was removed with `unset_lsm_write_spec`).
|
||||
The returned spec — including its ``maintained_indexes`` and
|
||||
``writer_config_defaults`` — mirrors what was passed to
|
||||
`set_lsm_write_spec`.
|
||||
"""
|
||||
return await self._inner.get_lsm_write_spec()
|
||||
|
||||
async def close_lsm_writers(self) -> None:
|
||||
"""Drain and close any cached MemWAL shard writers for this table.
|
||||
|
||||
@@ -5432,9 +5520,44 @@ class AsyncTable:
|
||||
|
||||
return await self._inner.update(updates_sql, where)
|
||||
|
||||
async def refresh_column(
|
||||
self,
|
||||
columns,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
num_workers: Optional[int] = None,
|
||||
max_workers: Optional[int] = None,
|
||||
batch_size: Optional[int] = None,
|
||||
priority: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Trigger recompute of computed columns (REFRESH COLUMN).
|
||||
Returns the refresh job id. Server-backed feature.
|
||||
|
||||
num_workers / max_workers / batch_size / priority are per-refresh
|
||||
scheduling knobs (how to run THIS refresh); they override any default
|
||||
the function carries. `priority` is a Kueue tier
|
||||
(training | interactive | backfill)."""
|
||||
if isinstance(columns, str):
|
||||
columns = [columns]
|
||||
return await self._inner.refresh_column(
|
||||
list(columns),
|
||||
where_clause=where,
|
||||
num_workers=num_workers,
|
||||
max_workers=max_workers,
|
||||
batch_size=batch_size,
|
||||
priority=priority,
|
||||
)
|
||||
|
||||
async def add_columns(
|
||||
self, transforms: dict[str, str] | pa.field | List[pa.field] | pa.Schema
|
||||
) -> AddColumnsResult:
|
||||
self,
|
||||
transforms: dict[str, str]
|
||||
| pa.field
|
||||
| List[pa.field]
|
||||
| pa.Schema
|
||||
| None = None,
|
||||
*,
|
||||
computed: Optional[Dict] = None,
|
||||
) -> Optional[AddColumnsResult]:
|
||||
"""
|
||||
Add new columns with defined values.
|
||||
|
||||
@@ -5453,6 +5576,7 @@ class AsyncTable:
|
||||
version: the new version number of the table after adding columns.
|
||||
|
||||
"""
|
||||
result = None
|
||||
if isinstance(transforms, pa.Field):
|
||||
transforms = [transforms]
|
||||
if isinstance(transforms, list) and all(
|
||||
@@ -5460,9 +5584,69 @@ class AsyncTable:
|
||||
):
|
||||
transforms = pa.schema(transforms)
|
||||
if isinstance(transforms, pa.Schema):
|
||||
return await self._inner.add_columns_with_schema(transforms)
|
||||
result = await self._inner.add_columns_with_schema(transforms)
|
||||
elif transforms is not None:
|
||||
result = await self._inner.add_columns(list(transforms.items()))
|
||||
if computed:
|
||||
# computed binds an expression over a registered function to a
|
||||
# column: {col: fn("input_col")} -- fn("input_col") yields the
|
||||
# expression and carries the inferred type; a tuple key fans a
|
||||
# STRUCT return out to several columns. Declares the binding only;
|
||||
# the server fills the values (server-backed). The legacy
|
||||
# {col: (sql_type, expression)} tuple form is still accepted.
|
||||
for expression, cols in _computed_groups(computed):
|
||||
await self._inner.add_computed_columns(cols, expression)
|
||||
return result
|
||||
|
||||
async def add_computed_column(
|
||||
self,
|
||||
columns,
|
||||
fn,
|
||||
args: Optional[List[str]] = None,
|
||||
types=None,
|
||||
) -> None:
|
||||
"""Declare computed column(s) bound to a UDF (async).
|
||||
|
||||
.. deprecated::
|
||||
Use ``add_columns(computed={"col": fn("input_col")})`` -- a computed
|
||||
column is an expression over a registered function, so bind it that
|
||||
way instead of coupling the UDF to the column here.
|
||||
"""
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"add_computed_column is deprecated; use add_columns(computed="
|
||||
'{"col": fn("input_col")}).',
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
from .udf import Udf, struct_field_types
|
||||
|
||||
multi = isinstance(columns, (tuple, list))
|
||||
if isinstance(fn, Udf):
|
||||
expr = fn.expression(*(args or []))
|
||||
if types is None:
|
||||
if multi:
|
||||
if not fn.returns.upper().startswith("STRUCT"):
|
||||
raise ValueError(
|
||||
"several columns need a STRUCT-returning function"
|
||||
)
|
||||
types = struct_field_types(fn.returns)
|
||||
else:
|
||||
types = fn.returns
|
||||
else:
|
||||
return await self._inner.add_columns(list(transforms.items()))
|
||||
if types is None:
|
||||
raise ValueError("pass types= when fn is a name string")
|
||||
expr = f"{fn}({', '.join(args or [])})"
|
||||
if multi:
|
||||
if len(types) != len(columns):
|
||||
raise ValueError(
|
||||
f"{len(columns)} columns but {len(types)} output types"
|
||||
)
|
||||
computed = {c: (t, expr) for c, t in zip(columns, types)}
|
||||
else:
|
||||
computed = {columns: (types, expr)}
|
||||
await self.add_columns(computed=computed)
|
||||
|
||||
async def alter_columns(
|
||||
self, *alterations: Iterable[dict[str, Any]]
|
||||
|
||||
@@ -0,0 +1,753 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""UDF authoring for LanceDB derived compute (server-backed).
|
||||
|
||||
`@udf` / `@table_udf` turn a plain Python function into a registrable
|
||||
server-side UDF: a cloudpickled (or source) body, a SQL signature inferred
|
||||
from type hints, and the runtime options (pip deps, GPUs, batching, ...).
|
||||
Register and use them through the existing connection/table API:
|
||||
|
||||
import lancedb
|
||||
from lancedb import udf, table_udf
|
||||
|
||||
db = lancedb.connect("db://my_db", api_key="...", host_override="...")
|
||||
|
||||
@udf(pip=["torch>=2.0"], num_gpus=1)
|
||||
def embed(text: str) -> list[float]:
|
||||
return model.encode(text).tolist()
|
||||
|
||||
db.create_function(embed) # CREATE FUNCTION (once)
|
||||
tbl = db.open_table("docs")
|
||||
tbl.add_columns(computed={"vec": embed("text")}) # bind embed(text) -> vec
|
||||
tbl.refresh_column("vec").wait() # materialize (returns a JobHandle)
|
||||
view = db.create_materialized_view("chunks", tbl, ["id", chunk_fn])
|
||||
|
||||
`embed("text")` applies the registered function to the `text` column and yields
|
||||
the expression `embed(text)`; the function itself stays decoupled from any
|
||||
column, so the same `embed` works on any column or table.
|
||||
|
||||
These operations are server-backed (LanceDB Enterprise / Cloud); the
|
||||
decorator itself works locally (define + call), only registration needs a
|
||||
remote connection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import dataclasses
|
||||
import functools
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
import typing
|
||||
|
||||
# -- type hints -> SQL type strings -------------------------------------
|
||||
|
||||
_SCALARS = {
|
||||
int: "BIGINT",
|
||||
# Pragmatic default for ML workloads: python float maps to FLOAT
|
||||
# (Float32). Use an explicit `returns=` for DOUBLE.
|
||||
float: "FLOAT",
|
||||
str: "VARCHAR",
|
||||
bool: "BOOLEAN",
|
||||
bytes: "BLOB",
|
||||
}
|
||||
|
||||
|
||||
class TypeInferenceError(TypeError):
|
||||
pass
|
||||
|
||||
|
||||
def sql_type(hint) -> str:
|
||||
"""SQL type string for a python type hint."""
|
||||
if hint in _SCALARS:
|
||||
return _SCALARS[hint]
|
||||
origin = typing.get_origin(hint)
|
||||
if origin in (list, typing.List):
|
||||
(item,) = typing.get_args(hint) or (None,)
|
||||
if item in _SCALARS:
|
||||
return f"{_SCALARS[item]}[]"
|
||||
raise TypeInferenceError(
|
||||
f"unsupported list item type {item!r}; use an explicit returns="
|
||||
)
|
||||
fields = _struct_fields(hint)
|
||||
if fields is not None:
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"STRUCT({inner})"
|
||||
raise TypeInferenceError(
|
||||
f"cannot infer a SQL type for {hint!r}; pass an explicit type string"
|
||||
)
|
||||
|
||||
|
||||
def _struct_fields(hint):
|
||||
"""(name, hint) pairs for a TypedDict or dataclass, else None."""
|
||||
if dataclasses.is_dataclass(hint):
|
||||
return [(f.name, f.type) for f in dataclasses.fields(hint)]
|
||||
# TypedDict detection: a dict subclass with __annotations__.
|
||||
if (
|
||||
isinstance(hint, type)
|
||||
and issubclass(hint, dict)
|
||||
and typing.get_type_hints(hint)
|
||||
):
|
||||
return list(typing.get_type_hints(hint).items())
|
||||
return None
|
||||
|
||||
|
||||
def return_type(fn, override: "str | None", table: bool) -> str:
|
||||
"""SQL return type for a function: explicit override wins, else the
|
||||
return annotation. Table functions render as TABLE(...) and accept
|
||||
struct-shaped hints (TypedDict/dataclass, optionally list-wrapped)."""
|
||||
if override is not None:
|
||||
s = override.strip()
|
||||
if table and not s.upper().startswith("TABLE"):
|
||||
if s.upper().startswith("STRUCT"):
|
||||
return "TABLE" + s[len("STRUCT") :]
|
||||
raise TypeInferenceError(
|
||||
"a table function's returns= must be TABLE(...) or STRUCT(...)"
|
||||
)
|
||||
return s
|
||||
|
||||
hints = typing.get_type_hints(fn)
|
||||
ret = hints.get("return")
|
||||
if ret is None:
|
||||
raise TypeInferenceError(
|
||||
f"function {fn.__name__!r} needs a return annotation or returns="
|
||||
)
|
||||
if table:
|
||||
# Accept list[Row] / Row where Row is a TypedDict or dataclass.
|
||||
if typing.get_origin(ret) in (list, typing.List):
|
||||
(ret,) = typing.get_args(ret)
|
||||
fields = _struct_fields(ret)
|
||||
if fields is None:
|
||||
raise TypeInferenceError(
|
||||
"a table function must return rows shaped as a TypedDict or "
|
||||
"dataclass (optionally list-wrapped); or pass returns=..."
|
||||
)
|
||||
inner = ", ".join(f"{name} {sql_type(h)}" for name, h in fields)
|
||||
return f"TABLE({inner})"
|
||||
return sql_type(ret)
|
||||
|
||||
|
||||
def param_types(fn) -> "list[tuple[str, str]]":
|
||||
"""(name, sql type) per parameter, from annotations. Each UDF
|
||||
parameter binds to a source column of the same name by default."""
|
||||
hints = typing.get_type_hints(fn)
|
||||
out = []
|
||||
for name, p in inspect.signature(fn).parameters.items():
|
||||
if p.kind in (p.VAR_POSITIONAL, p.VAR_KEYWORD):
|
||||
raise TypeInferenceError("*args/**kwargs are not supported in UDFs")
|
||||
hint = hints.get(name)
|
||||
if hint is None:
|
||||
raise TypeInferenceError(
|
||||
f"parameter {name!r} of {fn.__name__!r} needs a type annotation"
|
||||
)
|
||||
out.append((name, sql_type(hint)))
|
||||
return out
|
||||
|
||||
|
||||
# -- column expressions -------------------------------------------------
|
||||
|
||||
|
||||
class ColumnExpr(str):
|
||||
"""A computed-column expression produced by applying a registered
|
||||
function to column names, e.g. ``embed("data") -> "embed(data)"``.
|
||||
|
||||
It IS the expression string everywhere a string is expected (views, SQL,
|
||||
logging), and additionally carries the function's declared return type so
|
||||
``add_columns(computed=...)`` can declare the column without a hand-written
|
||||
type. ``field_types`` holds the per-field SQL types of a STRUCT return, for
|
||||
fanning one expression out to several columns.
|
||||
"""
|
||||
|
||||
data_type: "str | None"
|
||||
field_types: "list[str] | None"
|
||||
|
||||
def __new__(cls, expr: str, data_type=None, field_types=None):
|
||||
obj = super().__new__(cls, expr)
|
||||
obj.data_type = data_type
|
||||
obj.field_types = field_types
|
||||
return obj
|
||||
|
||||
|
||||
def _normalize_computed(computed: dict) -> dict:
|
||||
"""Normalize the user-facing ``computed=`` mapping to the canonical
|
||||
``{name: (sql_type, expression)}`` form.
|
||||
|
||||
Accepts, per entry:
|
||||
- value is a `ColumnExpr` (from ``fn("col")``): the column's SQL type
|
||||
comes from the function's return type -- no hand-written type needed. A
|
||||
tuple key (``("chunk", "idx")``) fans a STRUCT return out to one
|
||||
(type, expression) entry per field, in declared order.
|
||||
- value is a legacy ``(sql_type, expression)`` tuple: passed through (the
|
||||
escape hatch, e.g. bare-name function strings).
|
||||
"""
|
||||
out: dict = {}
|
||||
for key, val in computed.items():
|
||||
if isinstance(val, ColumnExpr):
|
||||
expr = str(val)
|
||||
if isinstance(key, (tuple, list)):
|
||||
if not val.field_types:
|
||||
raise ValueError(
|
||||
f"columns {tuple(key)} need a STRUCT-returning function; "
|
||||
f"{expr} returns a single value"
|
||||
)
|
||||
if len(val.field_types) != len(key):
|
||||
raise ValueError(
|
||||
f"{len(key)} columns but {len(val.field_types)} struct fields "
|
||||
f"in {expr}"
|
||||
)
|
||||
for name, t in zip(key, val.field_types):
|
||||
out[name] = (t, expr)
|
||||
else:
|
||||
if val.data_type is None:
|
||||
raise ValueError(f"cannot infer a type for {expr}; pass types=")
|
||||
out[key] = (val.data_type, expr)
|
||||
else:
|
||||
out[key] = val
|
||||
return out
|
||||
|
||||
|
||||
# -- the @udf / @table_udf decorators -----------------------------------
|
||||
|
||||
|
||||
class Udf:
|
||||
def __init__(
|
||||
self,
|
||||
fn,
|
||||
*,
|
||||
returns: "str | None" = None,
|
||||
table: bool = False,
|
||||
name: "str | None" = None,
|
||||
pip: "list[str] | None" = None,
|
||||
pip_index_url: "str | None" = None,
|
||||
pip_extra_index_urls: "list[str] | None" = None,
|
||||
find_links: "list[str] | None" = None,
|
||||
requirements: "str | list[str] | None" = None,
|
||||
conda: "list[str] | None" = None,
|
||||
conda_channels: "list[str] | None" = None,
|
||||
env: "dict[str, str] | list[str] | None" = None,
|
||||
num_cpus: "int | None" = None,
|
||||
num_gpus: "int | None" = None,
|
||||
batch_size: "int | None" = None,
|
||||
timeout: "float | None" = None,
|
||||
error_policy: "str | None" = None,
|
||||
max_skip_ratio: "float | None" = None,
|
||||
retries: "int | None" = None,
|
||||
docker_image: "str | None" = None,
|
||||
description: "str | None" = None,
|
||||
prefer_source: bool = False,
|
||||
):
|
||||
functools.update_wrapper(self, fn)
|
||||
self.fn = fn
|
||||
self.name = name or fn.__name__
|
||||
self.table = table
|
||||
self.params = param_types(fn)
|
||||
self.returns = return_type(fn, returns, table)
|
||||
self.prefer_source = prefer_source
|
||||
self.options: "dict[str, str]" = {}
|
||||
if conda and (pip or requirements):
|
||||
raise ValueError("pass conda or pip/requirements, not both")
|
||||
if conda_channels and not conda:
|
||||
raise ValueError("conda_channels requires conda")
|
||||
if pip:
|
||||
self.options["pip"] = ",".join(pip)
|
||||
if pip_extra_index_urls:
|
||||
self.options["pip_extra_index_urls"] = ",".join(pip_extra_index_urls)
|
||||
if find_links:
|
||||
self.options["find_links"] = ",".join(find_links)
|
||||
if requirements:
|
||||
self.options["requirements"] = _format_requirements(requirements)
|
||||
if conda:
|
||||
self.options["conda"] = ",".join(conda)
|
||||
if conda_channels:
|
||||
self.options["conda_channels"] = ",".join(conda_channels)
|
||||
if env:
|
||||
self.options["env"] = _format_env(env)
|
||||
for key, val in [
|
||||
("pip_index_url", pip_index_url),
|
||||
("num_cpus", num_cpus),
|
||||
("num_gpus", num_gpus),
|
||||
("batch_size", batch_size),
|
||||
("timeout", timeout),
|
||||
("error_policy", error_policy),
|
||||
("max_skip_ratio", max_skip_ratio),
|
||||
("retries", retries),
|
||||
("docker_image", docker_image),
|
||||
]:
|
||||
if val is not None:
|
||||
self.options[key] = str(val)
|
||||
# Keep the source in the description (when available) so the
|
||||
# catalog stays inspectable even for pickled bodies.
|
||||
if description is not None:
|
||||
self.options["description"] = description
|
||||
else:
|
||||
try:
|
||||
self.options["description"] = textwrap.dedent(inspect.getsource(fn))
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
"""Call with real values to run locally; call with column-name
|
||||
strings to build an expression for backfills and views, e.g.
|
||||
``embed("data")`` -> the expression ``embed(data)`` (a `ColumnExpr`
|
||||
carrying the function's return type for `add_columns(computed=...)`)."""
|
||||
if args and all(isinstance(a, str) for a in args) and not kwargs:
|
||||
return self.expression(*args)
|
||||
return self.fn(*args, **kwargs)
|
||||
|
||||
def expression(self, *columns: str) -> ColumnExpr:
|
||||
"""The expression applying this function to `columns` (default: the
|
||||
function's own parameter names). Returns a `ColumnExpr` -- a string
|
||||
that also carries the declared return type (and struct field types)."""
|
||||
cols = columns or [p for p, _ in self.params]
|
||||
expr = f"{self.name}({', '.join(cols)})"
|
||||
field_types = None
|
||||
if self.returns.upper().startswith("STRUCT"):
|
||||
field_types = struct_field_types(self.returns)
|
||||
return ColumnExpr(expr, data_type=self.returns, field_types=field_types)
|
||||
|
||||
def _body(self) -> "tuple[str, str]":
|
||||
"""(body literal, body_format). Source when requested and
|
||||
retrievable; cloudpickle otherwise (handles closures)."""
|
||||
if self.prefer_source:
|
||||
try:
|
||||
src = textwrap.dedent(inspect.getsource(self.fn))
|
||||
# Strip the decorator line(s) so the stored body is a
|
||||
# plain function definition.
|
||||
lines = src.splitlines(keepends=True)
|
||||
while lines and lines[0].lstrip().startswith("@"):
|
||||
lines.pop(0)
|
||||
return "".join(lines), "source"
|
||||
except (OSError, TypeError):
|
||||
pass
|
||||
import cloudpickle
|
||||
|
||||
raw = cloudpickle.dumps(self.fn)
|
||||
return base64.b64encode(raw).decode("ascii"), "cloudpickle"
|
||||
|
||||
def _body_and_options(self) -> "tuple[str, dict[str, str]]":
|
||||
"""The body literal plus the finalized options (body_format /
|
||||
python_version / cloudpickle-pip bookkeeping for a non-source
|
||||
body)."""
|
||||
body, body_format = self._body()
|
||||
options = dict(self.options)
|
||||
if body_format != "source":
|
||||
options["body_format"] = body_format
|
||||
# Pickled code objects only load under the same interpreter
|
||||
# minor version; record ours so the worker can fail with a
|
||||
# clear message instead of a bytecode error.
|
||||
options["python_version"] = self.pickle_environment()
|
||||
# The worker deserializes the body with cloudpickle; make sure
|
||||
# the job's pip environment provides it. Conda bakes inject
|
||||
# cloudpickle server-side, so do not create an invalid pip+conda
|
||||
# declaration here.
|
||||
if "conda" not in options:
|
||||
pip = [d for d in options.get("pip", "").split(",") if d]
|
||||
if not any(d.startswith("cloudpickle") for d in pip):
|
||||
pip.append("cloudpickle")
|
||||
options["pip"] = ",".join(pip)
|
||||
return body, options
|
||||
|
||||
def create_request(self) -> dict:
|
||||
"""Keyword arguments for `connection.create_function`."""
|
||||
body, options = self._body_and_options()
|
||||
return {
|
||||
"name": self.name,
|
||||
"language": "python",
|
||||
"return_type": self.returns,
|
||||
"body": body,
|
||||
"options": options,
|
||||
}
|
||||
|
||||
def create_statement(self) -> str:
|
||||
"""The equivalent `CREATE FUNCTION` SQL (for SQL-surface callers)."""
|
||||
params = ", ".join(f"{n} {t}" for n, t in self.params)
|
||||
body, options = self._body_and_options()
|
||||
with_clause = ""
|
||||
if options:
|
||||
rendered = ", ".join(
|
||||
f"{k} = '{_escape(v)}'" for k, v in sorted(options.items())
|
||||
)
|
||||
with_clause = f" WITH ({rendered})"
|
||||
return (
|
||||
f"CREATE FUNCTION {self.name}({params}) RETURNS {self.returns} "
|
||||
f"LANGUAGE python AS '{_escape_body(body)}'{with_clause}"
|
||||
)
|
||||
|
||||
def pickle_environment(self) -> str:
|
||||
"""Python version the body pickles under -- workers should match
|
||||
the minor version for cloudpickle compatibility."""
|
||||
return f"{sys.version_info.major}.{sys.version_info.minor}"
|
||||
|
||||
|
||||
def _escape(s: str) -> str:
|
||||
return str(s).replace("'", "''")
|
||||
|
||||
|
||||
def _format_requirements(requirements: "str | list[str]") -> str:
|
||||
if isinstance(requirements, str):
|
||||
return requirements
|
||||
return "\n".join(str(req) for req in requirements)
|
||||
|
||||
|
||||
def _format_env(env: "dict[str, str] | list[str]") -> str:
|
||||
if isinstance(env, dict):
|
||||
return "; ".join(f"{key}={value}" for key, value in env.items())
|
||||
return "; ".join(str(entry) for entry in env)
|
||||
|
||||
|
||||
def _escape_body(body: str) -> str:
|
||||
# The server unescapes \n / \t in single-quoted bodies; encode real
|
||||
# newlines accordingly and escape quotes.
|
||||
return (
|
||||
body.replace("\\", "\\\\")
|
||||
.replace("'", "''")
|
||||
.replace("\n", "\\n")
|
||||
.replace("\t", "\\t")
|
||||
)
|
||||
|
||||
|
||||
def udf(fn=None, **kwargs):
|
||||
"""Decorate a function as a scalar (or struct-returning) UDF.
|
||||
|
||||
@udf
|
||||
def doubled(val: int) -> float: ...
|
||||
|
||||
@udf(pip=["torch>=2"], num_gpus=1)
|
||||
def embed(body: str) -> list[float]: ...
|
||||
"""
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
def table_udf(fn=None, **kwargs):
|
||||
"""Decorate a table function (UDTF): each input row may emit zero or
|
||||
more output rows. Only usable in materialized views.
|
||||
|
||||
class Chunk(TypedDict):
|
||||
chunk: str
|
||||
chunk_idx: int
|
||||
|
||||
@table_udf
|
||||
def chunker(body: str) -> list[Chunk]: ...
|
||||
"""
|
||||
kwargs["table"] = True
|
||||
if fn is not None:
|
||||
return Udf(fn, **kwargs)
|
||||
return lambda f: Udf(f, **kwargs)
|
||||
|
||||
|
||||
# -- view / job handles (thin references over a connection) -------------
|
||||
|
||||
|
||||
def struct_field_types(returns: str) -> "list[str]":
|
||||
"""Field type strings of a STRUCT(...) SQL type, in declared order."""
|
||||
inner = returns.strip()[len("STRUCT(") : -1]
|
||||
fields, depth, start = [], 0, 0
|
||||
for i, c in enumerate(inner):
|
||||
if c in "([":
|
||||
depth += 1
|
||||
elif c in ")]":
|
||||
depth -= 1
|
||||
elif c == "," and depth == 0:
|
||||
fields.append(inner[start:i].strip())
|
||||
start = i + 1
|
||||
fields.append(inner[start:].strip())
|
||||
# Each field is "name TYPE"; drop the name.
|
||||
return [f.split(None, 1)[1] for f in fields]
|
||||
|
||||
|
||||
def build_view_query(source, select) -> str:
|
||||
"""Assemble a view SELECT from a source (name or table) and select
|
||||
items: a column name, an expression string, a (alias, expression)
|
||||
tuple, or a @udf/@table_udf object."""
|
||||
src = source.name if hasattr(source, "name") else source
|
||||
items = []
|
||||
for item in select:
|
||||
if isinstance(item, Udf):
|
||||
items.append(item.expression())
|
||||
elif isinstance(item, tuple):
|
||||
alias, expr = item
|
||||
expr = expr.expression() if isinstance(expr, Udf) else expr
|
||||
items.append(f"{expr} AS {alias}")
|
||||
else:
|
||||
items.append(item)
|
||||
return f"SELECT {', '.join(items)} FROM {src}"
|
||||
|
||||
|
||||
def _job_id_matches(handle_id: str, listed_id: str) -> bool:
|
||||
# The refresh/backfill endpoints return the submission id (a uuid), but
|
||||
# the agent names the manifest job "<table>-<type>-<first 8 of the
|
||||
# submission id>" -- which is what list_jobs and cancel report. Match the
|
||||
# canonical id directly, or by that submission prefix.
|
||||
if listed_id == handle_id:
|
||||
return True
|
||||
prefix = handle_id[:8]
|
||||
return len(prefix) >= 4 and prefix in listed_id
|
||||
|
||||
|
||||
class MaterializedView:
|
||||
"""A reference to a materialized view (name + connection). Operations are
|
||||
server-backed connection calls bound to the name.
|
||||
|
||||
``create_materialized_view`` returns one of these; ``job_id`` is the
|
||||
initial-population job (None when the view was created with no data), so
|
||||
``db.create_materialized_view(...).wait()`` blocks until it is populated.
|
||||
"""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return JobHandle(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
def refresh(self, full: bool = False) -> "JobHandle":
|
||||
"""Refresh the materialized view; returns a `JobHandle` to wait on,
|
||||
poll, or cancel (``view.refresh().wait()``).
|
||||
|
||||
``full=True`` forces a full rebuild (recompute and replace every row)
|
||||
instead of the default incremental refresh. A full rebuild preserves
|
||||
the view's indexes -- they are reindexed by the distributed indexer.
|
||||
"""
|
||||
job_id = self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return JobHandle(self.conn, job_id, table=self.name)
|
||||
|
||||
def explain_refresh(self, full: bool = False):
|
||||
"""Plan a refresh without running it (EXPLAIN REFRESH)."""
|
||||
return self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
def alter(self, auto_refresh: bool) -> None:
|
||||
self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
def drop(self) -> None:
|
||||
self.conn.drop_materialized_view(self.name)
|
||||
|
||||
# A materialized view is a first-class table: it can be indexed and
|
||||
# searched like any other. These open the materialized dataset by name and
|
||||
# delegate. Indexes declared this way are recorded against the view, so the
|
||||
# engine re-applies them after a full refresh rebuilds the dataset (a full
|
||||
# refresh overwrites the dataset, which would otherwise drop its indices).
|
||||
def _table(self):
|
||||
return self.conn.open_table(self.name)
|
||||
|
||||
def create_index(self, *args, **kwargs):
|
||||
"""Build an index on the materialized view (see Table.create_index)."""
|
||||
return self._table().create_index(*args, **kwargs)
|
||||
|
||||
def create_scalar_index(self, *args, **kwargs):
|
||||
"""Build a scalar index on the materialized view."""
|
||||
return self._table().create_scalar_index(*args, **kwargs)
|
||||
|
||||
def create_fts_index(self, *args, **kwargs):
|
||||
"""Build a full-text-search index on the materialized view."""
|
||||
return self._table().create_fts_index(*args, **kwargs)
|
||||
|
||||
def search(self, *args, **kwargs):
|
||||
"""Search the materialized view (vector / FTS / hybrid)."""
|
||||
return self._table().search(*args, **kwargs)
|
||||
|
||||
def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or one of its columns). Delegates
|
||||
to the backing table; the server already includes the view's sources
|
||||
and downstream dependents. Returns a `Lineage`."""
|
||||
return self._table().lineage(column, direction=direction, depth=depth)
|
||||
|
||||
|
||||
_PROGRESS = re.compile(r"(\d+)/(\d+)")
|
||||
|
||||
|
||||
class JobFailedError(RuntimeError):
|
||||
"""Raised by ``JobHandle.wait()`` when the server reports the job ``failed``.
|
||||
|
||||
Carries the server-side error so a doomed backfill (e.g. a multi-column
|
||||
``REFRESH COLUMN`` of a scalar UDF) surfaces its real cause promptly,
|
||||
instead of the caller blocking until ``wait()``'s timeout.
|
||||
"""
|
||||
|
||||
def __init__(self, job_id: str, error: "str | None"):
|
||||
self.job_id = job_id
|
||||
self.error = error
|
||||
super().__init__(f"job {job_id} failed: {error or 'unknown error'}")
|
||||
|
||||
|
||||
class JobHandle:
|
||||
"""A reference to an inflight server-side job, with polling helpers."""
|
||||
|
||||
#: How long an unseen job is treated as still materializing (submission
|
||||
#: -> agent cycle -> manifest write is async).
|
||||
GRACE_SECONDS = 20.0
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.id = job_id
|
||||
#: The job's table, when known (refresh_column / MV refresh). Lets the
|
||||
#: server resolve this job with an O(1) single-node read; without it the
|
||||
#: lookup scans the database's active jobs (still correct).
|
||||
self.table = table
|
||||
self._created = time.monotonic()
|
||||
self._seen = False
|
||||
|
||||
def _job(self):
|
||||
# Poll by id (one job), not list_jobs (every active job): the server
|
||||
# matches the submission/manifest id and reads just this table's node.
|
||||
return self.conn.get_job(self.id, self.table)
|
||||
|
||||
def status(self) -> str:
|
||||
"""pending / running / cancelling / stale, or 'finished' once the
|
||||
job has left the inflight listing."""
|
||||
job = self._job()
|
||||
if job is not None:
|
||||
self._seen = True
|
||||
return job.state
|
||||
if not self._seen and time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "finished"
|
||||
|
||||
def progress(self) -> "tuple[int, int] | None":
|
||||
"""(units_done, units_total) while running, else None."""
|
||||
job = self._job()
|
||||
if job is not None and job.units_total is not None:
|
||||
return job.units_done or 0, job.units_total
|
||||
return None
|
||||
|
||||
def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
state = self.status()
|
||||
if state in ("finished", "stale"):
|
||||
return state
|
||||
if state == "failed":
|
||||
# Terminal failure -- surface the server error now, don't block
|
||||
# until `timeout`. `finalize` wrote it to the job's status node.
|
||||
job = self._job()
|
||||
raise JobFailedError(self.id, job.error if job is not None else None)
|
||||
if state == "pending":
|
||||
time.sleep(min(poll, 0.5))
|
||||
continue
|
||||
job = self._job()
|
||||
if job is not None and job.committed:
|
||||
return "finished"
|
||||
time.sleep(poll)
|
||||
raise TimeoutError(f"job {self.id} still {self.status()} after {timeout}s")
|
||||
|
||||
def cancel(self) -> None:
|
||||
# Cancel by the canonical manifest id (what cancel matches), found
|
||||
# via the submission prefix; fall back to the raw id.
|
||||
job = self._job()
|
||||
self.conn.cancel_job(job.job_id if job is not None else self.id)
|
||||
|
||||
|
||||
class AsyncMaterializedView:
|
||||
"""Async reference to a materialized view (name + async connection)."""
|
||||
|
||||
def __init__(self, conn, name: str, job_id: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.name = name
|
||||
#: initial-population job id from create, or None (with_no_data).
|
||||
self.job_id = job_id
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
"""Block until the initial-population job (from create) finishes.
|
||||
A no-op when the view was created with no data."""
|
||||
if self.job_id is None:
|
||||
return "finished"
|
||||
return await AsyncJobHandle(self.conn, self.job_id, table=self.name).wait(
|
||||
timeout=timeout, poll=poll
|
||||
)
|
||||
|
||||
async def refresh(self, full: bool = False) -> "AsyncJobHandle":
|
||||
"""Refresh the materialized view; returns an `AsyncJobHandle` to wait
|
||||
on, poll, or cancel.
|
||||
|
||||
``full=True`` forces a full rebuild instead of an incremental refresh
|
||||
(indexes are preserved and reindexed by the distributed indexer).
|
||||
"""
|
||||
job_id = await self.conn._refresh_materialized_view(self.name, full=full)
|
||||
return AsyncJobHandle(self.conn, job_id, table=self.name)
|
||||
|
||||
async def explain_refresh(self, full: bool = False):
|
||||
return await self.conn.explain_refresh_materialized_view(self.name, full=full)
|
||||
|
||||
async def alter(self, auto_refresh: bool) -> None:
|
||||
await self.conn.alter_materialized_view(self.name, auto_refresh=auto_refresh)
|
||||
|
||||
async def drop(self) -> None:
|
||||
await self.conn.drop_materialized_view(self.name)
|
||||
|
||||
async def lineage(self, column=None, *, direction=None, depth=None):
|
||||
"""Lineage of the materialized view (or column). Returns a `Lineage`."""
|
||||
return await self.conn.lineage(
|
||||
self.name, column, direction=direction, depth=depth
|
||||
)
|
||||
|
||||
|
||||
class AsyncJobHandle:
|
||||
"""Async reference to an inflight server-side job, with polling helpers."""
|
||||
|
||||
GRACE_SECONDS = 20.0
|
||||
|
||||
def __init__(self, conn, job_id: str, table: "str | None" = None):
|
||||
self.conn = conn
|
||||
self.id = job_id
|
||||
#: See JobHandle.table -- enables an O(1) by-id lookup when known.
|
||||
self.table = table
|
||||
self._created = time.monotonic()
|
||||
self._seen = False
|
||||
|
||||
async def _job(self):
|
||||
# Poll by id, not list_jobs (see JobHandle._job).
|
||||
return await self.conn.get_job(self.id, self.table)
|
||||
|
||||
async def status(self) -> str:
|
||||
job = await self._job()
|
||||
if job is not None:
|
||||
self._seen = True
|
||||
return job.state
|
||||
if not self._seen and time.monotonic() - self._created < self.GRACE_SECONDS:
|
||||
return "pending"
|
||||
return "finished"
|
||||
|
||||
async def progress(self) -> "tuple[int, int] | None":
|
||||
job = await self._job()
|
||||
if job is not None and job.units_total is not None:
|
||||
return job.units_done or 0, job.units_total
|
||||
return None
|
||||
|
||||
async def wait(self, timeout: float = 3600.0, poll: float = 2.0) -> str:
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
state = await self.status()
|
||||
if state in ("finished", "stale"):
|
||||
return state
|
||||
if state == "failed":
|
||||
# Terminal failure -- surface the server error now, don't block
|
||||
# until `timeout`. `finalize` wrote it to the job's status node.
|
||||
job = await self._job()
|
||||
raise JobFailedError(self.id, job.error if job is not None else None)
|
||||
if state == "pending":
|
||||
await asyncio.sleep(min(poll, 0.5))
|
||||
continue
|
||||
job = await self._job()
|
||||
if job is not None and job.committed:
|
||||
return "finished"
|
||||
await asyncio.sleep(poll)
|
||||
raise TimeoutError(
|
||||
f"job {self.id} still {await self.status()} after {timeout}s"
|
||||
)
|
||||
|
||||
async def cancel(self) -> None:
|
||||
job = await self._job()
|
||||
await self.conn.cancel_job(job.job_id if job is not None else self.id)
|
||||
@@ -177,10 +177,7 @@ def flatten_columns(tbl: pa.Table, flatten: Optional[Union[int, bool]] = None):
|
||||
continue
|
||||
else:
|
||||
break
|
||||
# `bool` is a subclass of `int`, so guard against it explicitly: `flatten=False`
|
||||
# (and `None`) must mean "do not flatten" rather than falling into the integer
|
||||
# branch and raising on the `flatten <= 0` check.
|
||||
elif isinstance(flatten, int) and not isinstance(flatten, bool):
|
||||
elif isinstance(flatten, int):
|
||||
if flatten <= 0:
|
||||
raise ValueError(
|
||||
"Please specify a positive integer for flatten or the boolean "
|
||||
@@ -518,38 +515,3 @@ def batch_to_tensor_rows(batch: pa.RecordBatch):
|
||||
stacked = torch.tensor(numpy.column_stack(columns))
|
||||
rows = list(stacked.unbind(dim=0))
|
||||
return rows
|
||||
|
||||
|
||||
def batch_to_tensor_dict(batch: pa.RecordBatch):
|
||||
"""
|
||||
Convert a PyArrow RecordBatch to a list of per-row dicts of PyTorch Tensors.
|
||||
|
||||
Each column is first converted to a 1-D tensor (zero-copy via DLPack), then
|
||||
sliced per row. The result is a list whose length is ``batch.num_rows`` and
|
||||
whose items are dicts keyed by column name. This shape composes directly
|
||||
with PyTorch's default ``DataLoader`` collate, which stacks the per-row
|
||||
dicts back into a dict of batched tensors.
|
||||
|
||||
Fails if torch is not installed.
|
||||
Fails if a column's data type is not supported by PyTorch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
batch : pa.RecordBatch
|
||||
The record batch to convert.
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[dict[str, torch.Tensor]]
|
||||
One per-row dict per row in the batch. Each dict maps column name to a
|
||||
0-D tensor view into the column.
|
||||
"""
|
||||
torch = attempt_import_or_raise("torch", "torch")
|
||||
tensors = {
|
||||
name: torch.from_dlpack(col)
|
||||
for name, col in zip(batch.schema.names, batch.columns)
|
||||
}
|
||||
return [
|
||||
{name: tensor[i] for name, tensor in tensors.items()}
|
||||
for i in range(batch.num_rows)
|
||||
]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,7 +2,6 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from typing import List, Optional, Union
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -243,49 +242,6 @@ def test_embedding_with_bad_results(tmp_path):
|
||||
assert tbl["vector"].null_count == 1
|
||||
|
||||
|
||||
def test_embedding_with_empty_output_vectors(tmp_path):
|
||||
"""Regression test for issue #1672.
|
||||
|
||||
When an embedding function returns an empty list (e.g. for empty-string
|
||||
inputs), _append_vector_columns used to crash because PyArrow cannot cast
|
||||
[] into a fixed-size list element. The fix replaces wrong-length vectors
|
||||
with None before building the Arrow array so that _handle_bad_vectors can
|
||||
process them normally.
|
||||
"""
|
||||
|
||||
@register("empty-vec-embedding")
|
||||
class EmptyVecEmbeddingFunction(TextEmbeddingFunction):
|
||||
def ndims(self):
|
||||
return 128
|
||||
|
||||
def generate_embeddings(self, texts: Union[List[str], np.ndarray]) -> list:
|
||||
# Simulate a model that returns an empty list for blank inputs
|
||||
return [
|
||||
[] if text.strip() == "" else np.random.randn(self.ndims()).tolist()
|
||||
for text in texts
|
||||
]
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
registry = EmbeddingFunctionRegistry.get_instance()
|
||||
model = registry.get("empty-vec-embedding").create()
|
||||
|
||||
class Schema(LanceModel):
|
||||
text: str = model.SourceField()
|
||||
vector: Vector(model.ndims()) = model.VectorField()
|
||||
|
||||
table = db.create_table("test_empty_vec", schema=Schema, mode="overwrite")
|
||||
|
||||
# Should not crash; the row with the empty string should be dropped
|
||||
table.add(
|
||||
[{"text": "hello world"}, {"text": ""}, {"text": "foo"}],
|
||||
on_bad_vectors="drop",
|
||||
)
|
||||
|
||||
assert len(table) == 2
|
||||
texts = table.to_arrow()["text"].to_pylist()
|
||||
assert "" not in texts
|
||||
|
||||
|
||||
def test_with_existing_vectors(tmp_path):
|
||||
@register("mock-embedding")
|
||||
class MockEmbeddingFunction(TextEmbeddingFunction):
|
||||
@@ -592,22 +548,6 @@ def test_openai_no_retry_on_401(mock_sleep):
|
||||
assert mock_sleep.call_count == 0
|
||||
|
||||
|
||||
def test_ollama_embeddings_pickle():
|
||||
"""OllamaEmbeddings must pickle even after the cached client is created."""
|
||||
registry = get_registry()
|
||||
model = registry.get("ollama").create(name="nomic-embed-text")
|
||||
|
||||
# Simulate accessing the cached client, which stores it on the instance.
|
||||
model.__dict__["_ollama_client"] = MagicMock()
|
||||
|
||||
pickled = pickle.dumps(model)
|
||||
restored = pickle.loads(pickled)
|
||||
|
||||
assert restored.name == "nomic-embed-text"
|
||||
assert restored.host == "http://localhost:11434"
|
||||
assert "_ollama_client" not in restored.__dict__
|
||||
|
||||
|
||||
def test_url_retrieve_downloads_image():
|
||||
"""
|
||||
Embedding functions like open-clip, siglip, and jinaai use url_retrieve()
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
"""JobHandle.wait() terminal-state handling.
|
||||
|
||||
Regression coverage for the cluster backfill-failure hang: the server reports a
|
||||
doomed job as ``state="failed"`` within seconds, but ``wait()`` used to ignore
|
||||
``failed`` and block until its (default 3600s) timeout. These tests pin that a
|
||||
``failed`` job raises ``JobFailedError`` promptly, carrying the server error.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from lancedb.udf import JobHandle, AsyncJobHandle, JobFailedError
|
||||
|
||||
|
||||
class FakeJobInfo:
|
||||
"""Mirror of the pyo3 builtins.JobInfo fields wait()/status() read."""
|
||||
|
||||
def __init__(self, state, error=None, committed=False, units_total=None):
|
||||
self.state = state
|
||||
self.error = error
|
||||
self.committed = committed
|
||||
self.units_total = units_total
|
||||
self.units_done = None
|
||||
self.job_id = "job-1"
|
||||
|
||||
|
||||
class FakeConn:
|
||||
"""get_job() walks a scripted list of JobInfo (or None) snapshots, holding
|
||||
the last one once exhausted, so wait() polls a deterministic timeline."""
|
||||
|
||||
def __init__(self, snapshots):
|
||||
self._snaps = list(snapshots)
|
||||
self.calls = 0
|
||||
|
||||
def get_job(self, job_id, table=None):
|
||||
snap = self._snaps[min(self.calls, len(self._snaps) - 1)]
|
||||
self.calls += 1
|
||||
return snap
|
||||
|
||||
|
||||
class AsyncFakeConn(FakeConn):
|
||||
async def get_job(self, job_id, table=None):
|
||||
return FakeConn.get_job(self, job_id, table)
|
||||
|
||||
|
||||
def test_wait_raises_on_failed_promptly():
|
||||
# pending -> failed: wait() must raise the server error, not TimeoutError.
|
||||
conn = FakeConn(
|
||||
[None, FakeJobInfo("failed", error="multi-column backfill needs a STRUCT")]
|
||||
)
|
||||
jh = JobHandle(conn, "job-1", table="t")
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
jh.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5 # prompt, nowhere near the 30s timeout
|
||||
assert "STRUCT" in str(exc.value)
|
||||
assert exc.value.error == "multi-column backfill needs a STRUCT"
|
||||
assert exc.value.job_id == "job-1"
|
||||
|
||||
|
||||
def test_wait_returns_finished_on_success():
|
||||
# running -> finished (job left the inflight listing) returns normally.
|
||||
conn = FakeConn([FakeJobInfo("running", units_total=2), None])
|
||||
jh = JobHandle(conn, "job-1", table="t")
|
||||
jh._seen = True # already observed, so a None now means "finished" not grace
|
||||
assert jh.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
|
||||
def test_wait_returns_finished_on_committed():
|
||||
# A committed job that is still listed resolves to finished.
|
||||
conn = FakeConn([FakeJobInfo("running", committed=True, units_total=2)])
|
||||
jh = JobHandle(conn, "job-1", table="t")
|
||||
jh._seen = True
|
||||
assert jh.wait(timeout=30, poll=0.01) == "finished"
|
||||
|
||||
|
||||
def test_async_wait_raises_on_failed_promptly():
|
||||
conn = AsyncFakeConn([None, FakeJobInfo("failed", error="boom")])
|
||||
jh = AsyncJobHandle(conn, "job-1", table="t")
|
||||
|
||||
async def run():
|
||||
t0 = time.monotonic()
|
||||
with pytest.raises(JobFailedError) as exc:
|
||||
await jh.wait(timeout=30, poll=0.01)
|
||||
assert time.monotonic() - t0 < 5
|
||||
assert exc.value.error == "boom"
|
||||
|
||||
asyncio.run(run())
|
||||
@@ -11,7 +11,6 @@ import lancedb
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb._lancedb import LsmWriteSpec
|
||||
from lancedb.index import BTree
|
||||
|
||||
SCHEMA = pa.schema(
|
||||
[
|
||||
@@ -137,76 +136,3 @@ def test_lsm_write_spec_identity_and_writer_config_defaults():
|
||||
s = s.with_writer_config_defaults({"durable_write": "false"})
|
||||
assert s.writer_config_defaults == {"durable_write": "false"}
|
||||
assert "durable_write" in repr(s)
|
||||
|
||||
|
||||
def test_get_lsm_write_spec(tmp_path):
|
||||
_db, table = _make_table(tmp_path)
|
||||
table.set_unenforced_primary_key("id")
|
||||
|
||||
# None when nothing is installed.
|
||||
assert table.get_lsm_write_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
table.create_index("id", config=BTree())
|
||||
idx_name = table.list_indices()[0].name
|
||||
|
||||
# Bucket spec round-trips, including maintained indexes and writer config
|
||||
# defaults.
|
||||
table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 4)
|
||||
.with_maintained_indexes([idx_name])
|
||||
.with_writer_config_defaults({"durable_write": "false"})
|
||||
)
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
assert spec.num_buckets == 4
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
assert spec.writer_config_defaults == {"durable_write": "false"}
|
||||
|
||||
# After unset, None again.
|
||||
table.unset_lsm_write_spec()
|
||||
assert table.get_lsm_write_spec() is None
|
||||
|
||||
# Identity round-trips (column recovered from the schema).
|
||||
table.set_lsm_write_spec(LsmWriteSpec.identity("id"))
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec.spec_type == "identity"
|
||||
assert spec.column == "id"
|
||||
table.unset_lsm_write_spec()
|
||||
|
||||
# Unsharded round-trips (no routing column).
|
||||
table.set_lsm_write_spec(LsmWriteSpec.unsharded())
|
||||
spec = table.get_lsm_write_spec()
|
||||
assert spec.spec_type == "unsharded"
|
||||
assert spec.column is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_get_lsm_write_spec(tmp_path):
|
||||
db = await lancedb.connect_async(
|
||||
tmp_path, read_consistency_interval=timedelta(seconds=0)
|
||||
)
|
||||
table = await db.create_table(
|
||||
"t",
|
||||
pa.RecordBatchReader.from_batches(SCHEMA, [_batch(["seed"], [0])]),
|
||||
)
|
||||
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
|
||||
# A real scalar index is needed to name it as a maintained index.
|
||||
await table.create_index("id", config=BTree())
|
||||
idx_name = (await table.list_indices())[0].name
|
||||
|
||||
await table.set_lsm_write_spec(
|
||||
LsmWriteSpec.bucket("id", 8).with_maintained_indexes([idx_name])
|
||||
)
|
||||
spec = await table.get_lsm_write_spec()
|
||||
assert spec is not None
|
||||
assert spec.spec_type == "bucket"
|
||||
assert spec.column == "id"
|
||||
assert spec.num_buckets == 8
|
||||
assert spec.maintained_indexes == [idx_name]
|
||||
await table.unset_lsm_write_spec()
|
||||
assert await table.get_lsm_write_spec() is None
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
|
||||
import tempfile
|
||||
import shutil
|
||||
import importlib
|
||||
import pytest
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
@@ -104,40 +103,6 @@ class TestNamespaceConnection:
|
||||
assert isinstance(db, lancedb.LanceNamespaceDBConnection)
|
||||
assert len(list(db.table_names())) == 0
|
||||
|
||||
def test_sync_builtin_namespace_uses_rust_without_python_client(self, monkeypatch):
|
||||
"""Built-in sync namespace connections should not construct or call the
|
||||
Python namespace client for normal namespace/table management."""
|
||||
namespace_module = importlib.import_module("lancedb.namespace")
|
||||
|
||||
def fail_namespace_connect(*args, **kwargs):
|
||||
raise AssertionError("Python namespace client should not be constructed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
namespace_module, "namespace_connect", fail_namespace_connect
|
||||
)
|
||||
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
assert isinstance(db, lancedb.LanceNamespaceDBConnection)
|
||||
assert db._namespace_client is None
|
||||
assert db._route_pushdown_to_rust is True
|
||||
|
||||
db.create_namespace(["test_ns"])
|
||||
assert "test_ns" in db.list_namespaces().namespaces
|
||||
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = db.create_table("test_table", schema=schema, namespace_path=["test_ns"])
|
||||
assert table.namespace == ["test_ns"]
|
||||
assert "test_table" in db.table_names(namespace_path=["test_ns"])
|
||||
assert "test_table" in db.list_tables(namespace_path=["test_ns"]).tables
|
||||
|
||||
opened = db.open_table("test_table", namespace_path=["test_ns"])
|
||||
assert opened.namespace == ["test_ns"]
|
||||
|
||||
db.drop_table("test_table", namespace_path=["test_ns"])
|
||||
assert db.list_tables(namespace_path=["test_ns"]).tables == []
|
||||
db.drop_namespace(["test_ns"])
|
||||
assert "test_ns" not in db.list_namespaces().namespaces
|
||||
|
||||
def test_create_table_through_namespace(self):
|
||||
"""Test creating a table through namespace."""
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
@@ -599,61 +564,6 @@ class TestAsyncNamespaceConnection:
|
||||
table_names = await db.table_names()
|
||||
assert len(list(table_names)) == 0
|
||||
|
||||
async def test_async_builtin_namespace_uses_rust_without_python_client(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""Built-in async namespace connections should not construct or call the
|
||||
Python namespace client for normal namespace/table management."""
|
||||
namespace_module = importlib.import_module("lancedb.namespace")
|
||||
|
||||
def fail_namespace_connect(*args, **kwargs):
|
||||
raise AssertionError("Python namespace client should not be constructed")
|
||||
|
||||
monkeypatch.setattr(
|
||||
namespace_module, "namespace_connect", fail_namespace_connect
|
||||
)
|
||||
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert isinstance(db, lancedb.AsyncLanceNamespaceDBConnection)
|
||||
assert db._namespace_client is None
|
||||
assert db._route_pushdown_to_rust is True
|
||||
|
||||
await db.create_namespace(["test_ns"])
|
||||
assert "test_ns" in (await db.list_namespaces()).namespaces
|
||||
|
||||
schema = pa.schema([pa.field("id", pa.int64())])
|
||||
table = await db.create_table(
|
||||
"test_table", schema=schema, namespace_path=["test_ns"]
|
||||
)
|
||||
assert table._namespace_path == ["test_ns"]
|
||||
assert table._namespace_client is None
|
||||
assert table._route_pushdown_to_rust is True
|
||||
assert "test_table" in await db.table_names(namespace_path=["test_ns"])
|
||||
assert "test_table" in (await db.list_tables(namespace_path=["test_ns"])).tables
|
||||
|
||||
opened = await db.open_table("test_table", namespace_path=["test_ns"])
|
||||
assert opened._namespace_path == ["test_ns"]
|
||||
|
||||
await db.drop_table("test_table", namespace_path=["test_ns"])
|
||||
assert (await db.list_tables(namespace_path=["test_ns"])).tables == []
|
||||
await db.drop_namespace(["test_ns"])
|
||||
assert "test_ns" not in (await db.list_namespaces()).namespaces
|
||||
|
||||
async def test_async_namespace_client_is_lazy(self):
|
||||
"""namespace_client() should still return the backing client on demand."""
|
||||
pytest.importorskip("lance")
|
||||
from lance.namespace import DirectoryNamespace
|
||||
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert db._namespace_client is None
|
||||
|
||||
ns_client = await db.namespace_client()
|
||||
|
||||
assert isinstance(ns_client, DirectoryNamespace)
|
||||
namespace_id = ns_client.namespace_id().replace("\\\\", "\\")
|
||||
assert str(self.temp_dir) in namespace_id
|
||||
assert db._namespace_client is ns_client
|
||||
|
||||
# Async connect via namespace helper is not enabled yet.
|
||||
|
||||
async def test_create_table_async(self):
|
||||
@@ -908,11 +818,10 @@ class TestPushdownOperations:
|
||||
)
|
||||
assert db._route_pushdown_to_rust is True
|
||||
|
||||
def test_route_pushdown_to_rust_for_native_dir(self):
|
||||
"""The sync dir connection is natively built and defers QueryTable
|
||||
pushdown to Rust."""
|
||||
def test_route_pushdown_to_rust_false_for_dir(self):
|
||||
"""A non-native (dir) connection keeps the Python pushdown path."""
|
||||
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
|
||||
assert db._route_pushdown_to_rust is True
|
||||
assert db._route_pushdown_to_rust is False
|
||||
|
||||
def test_async_route_pushdown_to_rust_for_native_rest(self):
|
||||
"""The async connection must not silently bypass the read-freshness fix:
|
||||
@@ -925,11 +834,10 @@ class TestPushdownOperations:
|
||||
)
|
||||
assert db._route_pushdown_to_rust is True
|
||||
|
||||
def test_async_route_pushdown_to_rust_for_native_dir(self):
|
||||
"""The async dir connection is natively built and defers QueryTable
|
||||
pushdown to Rust."""
|
||||
def test_async_route_pushdown_to_rust_false_for_dir(self):
|
||||
"""The async non-native (dir) connection keeps the Python pushdown path."""
|
||||
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
|
||||
assert db._route_pushdown_to_rust is True
|
||||
assert db._route_pushdown_to_rust is False
|
||||
|
||||
def test_lance_table_to_arrow_uses_query_pushdown(self):
|
||||
namespace_client = _NamespaceClient()
|
||||
|
||||
@@ -935,41 +935,14 @@ def test_transform_fn(mem_db):
|
||||
try:
|
||||
import torch
|
||||
|
||||
# "torch" returns a list of per-row dicts. Default DataLoader collate
|
||||
# stacks the per-row dicts back into a dict of batched tensors.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
torch_batch = list(torch_perm.iter(10, skip_last_batch=False))[0]
|
||||
assert isinstance(torch_batch, list)
|
||||
assert len(torch_batch) == 10
|
||||
assert isinstance(torch_batch[0], dict)
|
||||
assert set(torch_batch[0].keys()) == {"id", "value"}
|
||||
assert isinstance(torch_batch[0]["id"], torch.Tensor)
|
||||
assert torch_batch[0]["id"].dtype == torch.int64
|
||||
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert set(rows[0].keys()) == {"id", "value"}
|
||||
assert isinstance(rows[0]["id"], torch.Tensor)
|
||||
|
||||
# "torch_row" returns a list of tensors, one per row.
|
||||
torch_rows = list(
|
||||
permutation.with_format("torch_row").iter(10, skip_last_batch=False)
|
||||
torch_result = list(
|
||||
permutation.with_format("torch").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_rows, list)
|
||||
assert len(torch_rows) == 10
|
||||
assert isinstance(torch_rows[0], torch.Tensor)
|
||||
assert torch_rows[0].shape == (2,)
|
||||
assert torch_rows[0].dtype == torch.int64
|
||||
|
||||
# "torch_col" stacks columns into a single 2D tensor.
|
||||
torch_col = list(
|
||||
permutation.with_format("torch_col").iter(10, skip_last_batch=False)
|
||||
)[0]
|
||||
assert isinstance(torch_col, torch.Tensor)
|
||||
assert torch_col.shape == (2, 10)
|
||||
assert torch_col.dtype == torch.int64
|
||||
assert isinstance(torch_result, list)
|
||||
assert len(torch_result) == 10
|
||||
assert isinstance(torch_result[0], torch.Tensor)
|
||||
assert torch_result[0].shape == (2,)
|
||||
assert torch_result[0].dtype == torch.int64
|
||||
except ImportError:
|
||||
# Skip check if torch is not installed
|
||||
pass
|
||||
|
||||
@@ -502,61 +502,6 @@ def test_with_row_id(table: lancedb.table.Table):
|
||||
assert rs["_rowid"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_where_repeated_combines_with_and(table: lancedb.table.Table):
|
||||
# Calling where() more than once should AND the filters together instead of
|
||||
# silently replacing the previous one (regression test for #2649).
|
||||
builder = table.search().where("id >= 1").where("id < 2")
|
||||
assert builder._where == "(id >= 1) AND (id < 2)"
|
||||
|
||||
ids = [row["id"] for row in builder.limit(10).to_list()]
|
||||
assert ids == [1]
|
||||
|
||||
|
||||
def test_where_repeated_combines_expr(table: lancedb.table.Table):
|
||||
from lancedb.expr import col, lit
|
||||
|
||||
builder = table.search().where(col("id") >= lit(1)).where(col("id") < lit(2))
|
||||
ids = [row["id"] for row in builder.limit(10).to_list()]
|
||||
assert ids == [1]
|
||||
|
||||
|
||||
def test_where_mixed_filter_kinds_combines(table: lancedb.table.Table):
|
||||
# Mixing a SQL string filter with an expression filter lowers the
|
||||
# expression to SQL and combines them as SQL strings.
|
||||
from lancedb.expr import col, lit
|
||||
|
||||
builder = table.search().where("id >= 1").where(col("id") < lit(2))
|
||||
ids = [row["id"] for row in builder.limit(10).to_list()]
|
||||
assert ids == [1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_where_repeated_combines_with_and_async(table_async: AsyncTable):
|
||||
ids = [
|
||||
row["id"]
|
||||
for row in (
|
||||
await table_async.query().where("id >= 1").where("id < 2").to_list()
|
||||
)
|
||||
]
|
||||
assert ids == [1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_where_mixed_filter_kinds_combines_async(table_async: AsyncTable):
|
||||
from lancedb.expr import col, lit
|
||||
|
||||
ids = [
|
||||
row["id"]
|
||||
for row in (
|
||||
await table_async.query()
|
||||
.where("id >= 1")
|
||||
.where(col("id") < lit(2))
|
||||
.to_list()
|
||||
)
|
||||
]
|
||||
assert ids == [1]
|
||||
|
||||
|
||||
def test_distance_range(table: lancedb.table.Table):
|
||||
q = [0, 0]
|
||||
rs = table.search(q).to_arrow()
|
||||
|
||||
@@ -412,12 +412,10 @@ def test_remote_permutation_is_picklable():
|
||||
content_len = int(request.headers.get("Content-Length"))
|
||||
body = json.loads(request.rfile.read(content_len))
|
||||
if "filter" in body:
|
||||
match = re.search(
|
||||
r"_rowoffset\s+in\s+\((.*?)\)", body["filter"], re.IGNORECASE
|
||||
)
|
||||
offsets = [int(o.strip()) for o in match.group(1).split(",")]
|
||||
match = re.search(r"_rowoffset in \((.*?)\)", body["filter"])
|
||||
offsets = [int(offset.strip()) for offset in match.group(1).split(",")]
|
||||
else:
|
||||
offsets = list(range(len(rows)))
|
||||
offsets = rows
|
||||
table = pa.table({"a": [rows[offset] for offset in offsets]})
|
||||
|
||||
request.send_response(200)
|
||||
|
||||
@@ -350,38 +350,6 @@ def test_mrr_reranker_empty_input():
|
||||
reranker.rerank_multivector([])
|
||||
|
||||
|
||||
def test_mrr_multivector_rewards_consensus():
|
||||
# Reciprocal ranks must be averaged across *all* ranking systems, treating a
|
||||
# missing system as 0. A document ranked first by every system must outrank a
|
||||
# document ranked first by only one of them.
|
||||
reranker = MRRReranker()
|
||||
|
||||
def ranking(row_ids):
|
||||
return pa.table({"_rowid": pa.array(row_ids, type=pa.int64())})
|
||||
|
||||
# Doc 1 is rank 1 in only the first system; doc 2 is rank 1 in two systems
|
||||
# and rank 2 in the third (strong cross-system consensus).
|
||||
rs1 = ranking([1, 2, 3])
|
||||
rs2 = ranking([2, 3, 4])
|
||||
rs3 = ranking([2, 5, 6])
|
||||
|
||||
result = reranker.rerank_multivector([rs1, rs2, rs3])
|
||||
scores = {
|
||||
row_id: score
|
||||
for row_id, score in zip(
|
||||
result["_rowid"].to_pylist(),
|
||||
result["_relevance_score"].to_pylist(),
|
||||
)
|
||||
}
|
||||
|
||||
# sum of reciprocal ranks / number of systems
|
||||
assert scores[1] == pytest.approx(1.0 / 3)
|
||||
assert scores[2] == pytest.approx((0.5 + 1.0 + 1.0) / 3)
|
||||
assert scores[2] > scores[1]
|
||||
# The consensus document ranks first overall.
|
||||
assert result["_rowid"].to_pylist()[0] == 2
|
||||
|
||||
|
||||
def test_rrf_reranker_distance():
|
||||
data = pa.table(
|
||||
{
|
||||
|
||||
@@ -1137,16 +1137,6 @@ def test_namespace_open_table_with_branch_version(tmp_path):
|
||||
assert db.open_table("t", namespace_path=["ns1"], branch="exp").count_rows() == 3
|
||||
|
||||
|
||||
def test_namespace_root_table_to_lance_uses_namespace_client(tmp_path):
|
||||
pytest.importorskip("lance") # "dir" impl is lance.namespace.DirectoryNamespace
|
||||
db = lancedb.connect_namespace("dir", {"root": str(tmp_path)})
|
||||
table = db.create_table("t", [{"i": 0}])
|
||||
|
||||
assert table._namespace_client is None
|
||||
assert table.to_lance().count_rows() == 1
|
||||
assert table._namespace_client is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_namespace_open_table_with_branch_version(tmp_path):
|
||||
pytest.importorskip("lance") # "dir" impl is lance.namespace.DirectoryNamespace
|
||||
|
||||
@@ -148,47 +148,15 @@ def test_permutation_dataloader(mem_db):
|
||||
for batch in dataloader:
|
||||
assert batch["a"].size(0) == 10
|
||||
|
||||
# "torch" produces a list of per-row dicts per batch. The default
|
||||
# DataLoader collate stacks the per-row dicts back into a batched dict.
|
||||
torch_perm = permutation.with_format("torch")
|
||||
batch = next(torch_perm.iter(10, skip_last_batch=False))
|
||||
assert isinstance(batch, list)
|
||||
assert len(batch) == 10
|
||||
assert isinstance(batch[0], dict)
|
||||
assert isinstance(batch[0]["a"], torch.Tensor)
|
||||
rows = torch_perm.__getitems__([0, 1, 2])
|
||||
assert isinstance(rows, list)
|
||||
assert len(rows) == 3
|
||||
assert isinstance(rows[0], dict)
|
||||
assert isinstance(rows[0]["a"], torch.Tensor)
|
||||
dataloader = torch.utils.data.DataLoader(torch_perm, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
# Spawn-based workers exercise the pickle round-trip path: the new
|
||||
# transform-as-list shape must survive pickling so workers produce the
|
||||
# same per-row dicts the parent does.
|
||||
spawn_loader = torch.utils.data.DataLoader(
|
||||
torch_perm,
|
||||
batch_size=10,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
for batch in spawn_loader:
|
||||
assert isinstance(batch, dict)
|
||||
assert batch["a"].shape == (10,)
|
||||
|
||||
# "torch_row" returns a list of row tensors. Works with the default
|
||||
# DataLoader collate (stacks rows into 2D).
|
||||
row_perm = permutation.with_format("torch_row")
|
||||
dataloader = torch.utils.data.DataLoader(row_perm, batch_size=10, shuffle=True)
|
||||
permutation = permutation.with_format("torch")
|
||||
dataloader = torch.utils.data.DataLoader(permutation, batch_size=10, shuffle=True)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 10
|
||||
assert batch.size(1) == 1
|
||||
|
||||
col_perm = permutation.with_format("torch_col")
|
||||
permutation = permutation.with_format("torch_col")
|
||||
dataloader = torch.utils.data.DataLoader(
|
||||
col_perm, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
permutation, collate_fn=lambda x: x, batch_size=10, shuffle=True
|
||||
)
|
||||
for batch in dataloader:
|
||||
assert batch.size(0) == 1
|
||||
|
||||
@@ -25,42 +25,10 @@ import pandas as pd
|
||||
import polars as pl
|
||||
import pytest
|
||||
import lancedb
|
||||
from lancedb.util import flatten_columns, get_uri_scheme, join_uri, value_to_sql
|
||||
from lancedb.util import get_uri_scheme, join_uri, value_to_sql
|
||||
from utils import exception_output
|
||||
|
||||
|
||||
def _struct_table() -> pa.Table:
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
"nested": pa.array([{"a": 1, "b": 2}, {"a": 3, "b": 4}]),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_flatten_columns():
|
||||
tbl = _struct_table()
|
||||
|
||||
# None / False mean "do not flatten": the struct column is preserved.
|
||||
# `False` is a regression guard: because bool is a subclass of int it used
|
||||
# to fall into the integer branch and raise ValueError (see issue).
|
||||
for no_flatten in (None, False):
|
||||
result = flatten_columns(tbl, no_flatten)
|
||||
assert result.column_names == ["id", "nested"]
|
||||
|
||||
# True flattens all nested levels.
|
||||
flattened = flatten_columns(tbl, True)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# A positive integer flattens up to that depth.
|
||||
flattened = flatten_columns(tbl, 1)
|
||||
assert flattened.column_names == ["id", "nested.a", "nested.b"]
|
||||
|
||||
# Non-positive integers are still rejected.
|
||||
with pytest.raises(ValueError):
|
||||
flatten_columns(tbl, 0)
|
||||
|
||||
|
||||
def test_normalize_uri():
|
||||
uris = [
|
||||
"relative/path",
|
||||
|
||||
+392
-41
@@ -18,7 +18,10 @@ use lancedb::{
|
||||
connection::Connection as LanceConnection,
|
||||
connection::NamespaceClientPushdownOperation,
|
||||
database::namespace::LanceNamespaceDatabase,
|
||||
database::{CreateTableMode, Database, ReadConsistency},
|
||||
database::{
|
||||
CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode, Database,
|
||||
ReadConsistency, RefreshMaterializedViewRequest, TableLineageRequest,
|
||||
},
|
||||
};
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
@@ -27,6 +30,92 @@ use pyo3::{
|
||||
types::{PyDict, PyDictMethods},
|
||||
};
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A registered materialized view definition.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
pub projection: Vec<String>,
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// One inflight server-side job.
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// One durable, completed/terminal server-side job record (SHOW JOB HISTORY).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobHistoryEntry {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// One per-row UDF error recorded by `error_policy=skip` (SHOW ERRORS).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct JobErrorEntry {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a REFRESH MATERIALIZED VIEW would execute (EXPLAIN REFRESH).
|
||||
#[pyclass(get_all)]
|
||||
#[derive(Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
pub rebuild: bool,
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct Connection {
|
||||
inner: Option<LanceConnection>,
|
||||
@@ -310,6 +399,308 @@ impl Connection {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, language, return_type, body, options=None))]
|
||||
pub fn create_function(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: Option<HashMap<String, String>>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_function(CreateFunctionRequest {
|
||||
name,
|
||||
language,
|
||||
return_type,
|
||||
body,
|
||||
options: options.unwrap_or_default(),
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let functions = inner.list_functions().await.infer_error()?;
|
||||
Ok(functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_function(self_: PyRef<'_, Self>, name: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_function(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, query, auto_refresh=false, with_no_data=false, partition_by=None))]
|
||||
pub fn create_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
partition_by: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.create_materialized_view(CreateMaterializedViewRequest {
|
||||
name,
|
||||
query,
|
||||
auto_refresh,
|
||||
with_no_data,
|
||||
partition_by,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None, num_workers=None, max_workers=None))]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_materialized_view(RefreshMaterializedViewRequest {
|
||||
name,
|
||||
full,
|
||||
src_version,
|
||||
num_workers,
|
||||
max_workers,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), returned as the
|
||||
/// server's lineage JSON string (the Python layer parses it).
|
||||
pub fn table_lineage(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
column: Option<String>,
|
||||
direction: Option<String>,
|
||||
depth: Option<u32>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.table_lineage(TableLineageRequest {
|
||||
name,
|
||||
column,
|
||||
direction,
|
||||
depth,
|
||||
})
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (name, full=false, src_version=None))]
|
||||
pub fn explain_refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let p = inner
|
||||
.explain_refresh_materialized_view(&name, full, src_version)
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: p.table_name,
|
||||
has_work: p.has_work,
|
||||
source_version: p.source_version,
|
||||
last_refreshed_version: p.last_refreshed_version,
|
||||
full_refresh: p.full_refresh,
|
||||
rebuild: p.rebuild,
|
||||
units_total: p.units_total,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn alter_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
auto_refresh: bool,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.alter_materialized_view(&name, auto_refresh)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn drop_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
name: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.drop_materialized_view(&name).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_materialized_views(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let views = inner.list_materialized_views().await.infer_error()?;
|
||||
Ok(views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let jobs = inner.list_jobs().await.infer_error()?;
|
||||
Ok(jobs
|
||||
.into_iter()
|
||||
.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cancel_job(self_: PyRef<'_, Self>, job_id: String) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.cancel_job(&job_id).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id, table=None))]
|
||||
pub fn get_job(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: String,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let job = inner
|
||||
.get_job(&job_id, table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(job.map(|j| JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
}))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None))]
|
||||
pub fn job_history(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner.job_history(job_id.as_deref()).await.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|r| JobHistoryEntry {
|
||||
table: r.table,
|
||||
job_id: r.job_id,
|
||||
job_type: r.job_type,
|
||||
state: r.state,
|
||||
column: r.column,
|
||||
created_ms: r.created_ms,
|
||||
updated_ms: r.updated_ms,
|
||||
completed_ms: r.completed_ms,
|
||||
rows_processed: r.rows_processed,
|
||||
rows_skipped: r.rows_skipped,
|
||||
error: r.error,
|
||||
events: r.events,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (job_id=None, table=None))]
|
||||
pub fn errors(
|
||||
self_: PyRef<'_, Self>,
|
||||
job_id: Option<String>,
|
||||
table: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.get_inner()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let rows = inner
|
||||
.errors(job_id.as_deref(), table.as_deref())
|
||||
.await
|
||||
.infer_error()?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|e| JobErrorEntry {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
})
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (cur_name, new_name, cur_namespace_path=None, new_namespace_path=None))]
|
||||
pub fn rename_table(
|
||||
self_: PyRef<'_, Self>,
|
||||
@@ -655,46 +1046,6 @@ pub fn connect_namespace_client(
|
||||
)))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (
|
||||
namespace_client_impl,
|
||||
namespace_client_properties,
|
||||
read_consistency_interval=None,
|
||||
storage_options=None,
|
||||
session=None,
|
||||
namespace_client_pushdown_operations=None,
|
||||
))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn connect_namespace(
|
||||
namespace_client_impl: String,
|
||||
namespace_client_properties: HashMap<String, String>,
|
||||
read_consistency_interval: Option<f64>,
|
||||
storage_options: Option<HashMap<String, String>>,
|
||||
session: Option<crate::session::Session>,
|
||||
namespace_client_pushdown_operations: Option<Vec<String>>,
|
||||
) -> PyResult<Connection> {
|
||||
let read_consistency_interval = read_consistency_interval.map(Duration::from_secs_f64);
|
||||
let namespace_client_pushdown_operations =
|
||||
parse_namespace_client_pushdown_operations(namespace_client_pushdown_operations)?;
|
||||
|
||||
let mut builder =
|
||||
lancedb::connect_namespace(&namespace_client_impl, namespace_client_properties)
|
||||
.pushdown_operations(namespace_client_pushdown_operations);
|
||||
if let Some(storage_options) = storage_options {
|
||||
builder = builder.storage_options(storage_options);
|
||||
}
|
||||
if let Some(read_consistency_interval) = read_consistency_interval {
|
||||
builder = builder.read_consistency_interval(read_consistency_interval);
|
||||
}
|
||||
if let Some(session) = session {
|
||||
builder = builder.session(session.inner.clone());
|
||||
}
|
||||
|
||||
Ok(Connection::new(
|
||||
crate::runtime::block_on(builder.execute()).infer_error()?,
|
||||
))
|
||||
}
|
||||
|
||||
/// Whether to build the namespace natively (from impl + properties) instead of
|
||||
/// wrapping a pre-built client. Native construction is required for the
|
||||
/// read-freshness provider to be installed
|
||||
|
||||
+8
-60
@@ -7,14 +7,12 @@
|
||||
//! build type-safe filter / projection expressions that map directly to
|
||||
//! DataFusion [`Expr`] nodes, bypassing SQL string parsing.
|
||||
|
||||
use std::ops::{Add, Div, Mul, Not, Sub};
|
||||
|
||||
use arrow::{datatypes::DataType, pyarrow::PyArrowType};
|
||||
use datafusion_common::ScalarValue;
|
||||
use lancedb::expr::{
|
||||
DfExpr, col as ldb_col, contains, expr_cast, is_in, lit as df_lit, lower, upper,
|
||||
};
|
||||
use pyo3::types::{PyBytes, PyDate, PyDateTime};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, PyAny, PyResult, exceptions::PyValueError, prelude::*, pyfunction};
|
||||
|
||||
/// A type-safe DataFusion expression.
|
||||
@@ -65,30 +63,30 @@ impl PyExpr {
|
||||
Self(self.0.clone().or(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Logical NOT.
|
||||
fn not_(&self) -> Self {
|
||||
use std::ops::Not;
|
||||
Self(self.0.clone().not())
|
||||
}
|
||||
|
||||
// ── arithmetic ───────────────────────────────────────────────────────────
|
||||
|
||||
/// Add expressions.
|
||||
fn add(&self, other: &Self) -> Self {
|
||||
use std::ops::Add;
|
||||
Self(self.0.clone().add(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Subtract expressions.
|
||||
fn sub(&self, other: &Self) -> Self {
|
||||
use std::ops::Sub;
|
||||
Self(self.0.clone().sub(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Multiply expressions.
|
||||
fn mul(&self, other: &Self) -> Self {
|
||||
use std::ops::Mul;
|
||||
Self(self.0.clone().mul(other.0.clone()))
|
||||
}
|
||||
|
||||
/// Divide expressions.
|
||||
fn div(&self, other: &Self) -> Self {
|
||||
use std::ops::Div;
|
||||
Self(self.0.clone().div(other.0.clone()))
|
||||
}
|
||||
|
||||
@@ -155,8 +153,7 @@ pub fn expr_col(name: &str) -> PyExpr {
|
||||
|
||||
/// Create a literal value expression.
|
||||
///
|
||||
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`, `date`,
|
||||
/// `datetime`, `Decimal`.
|
||||
/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`.
|
||||
#[pyfunction]
|
||||
pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
// bool must be checked before int because bool is a subclass of int in Python
|
||||
@@ -166,19 +163,6 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
if let Ok(i) = value.extract::<i64>() {
|
||||
return Ok(PyExpr(df_lit(i)));
|
||||
}
|
||||
// Decimal must be checked before f64: Python's Decimal implements __float__,
|
||||
// so value.extract::<f64>() would succeed and silently truncate the value to
|
||||
// f64, losing precision. Build a Decimal128 scalar to preserve it instead.
|
||||
if value.get_type().name()? == "Decimal" {
|
||||
let s = value.call_method0("__str__")?.extract::<String>()?;
|
||||
// Parse the decimal string into an i128 value, precision, and scale.
|
||||
let (val, precision, scale) = parse_decimal(&s)?;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Decimal128(
|
||||
Some(val),
|
||||
precision,
|
||||
scale,
|
||||
))));
|
||||
}
|
||||
if let Ok(f) = value.extract::<f64>() {
|
||||
return Ok(PyExpr(df_lit(f)));
|
||||
}
|
||||
@@ -189,48 +173,12 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult<PyExpr> {
|
||||
let bytes = value.extract::<Vec<u8>>()?;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Binary(Some(bytes)))));
|
||||
}
|
||||
|
||||
// datetime.datetime is a subclass of datetime.date, so it must be checked first.
|
||||
if let Ok(dt) = value.cast::<PyDateTime>() {
|
||||
let ts: f64 = dt.call_method0("timestamp")?.extract()?;
|
||||
let micros = (ts * 1_000_000.0).round() as i64;
|
||||
return Ok(PyExpr(df_lit(ScalarValue::TimestampMicrosecond(
|
||||
Some(micros),
|
||||
None,
|
||||
))));
|
||||
}
|
||||
if let Ok(d) = value.cast::<PyDate>() {
|
||||
let ordinal: i32 = d.call_method0("toordinal")?.extract()?;
|
||||
let days = ordinal - 719163; // Unix epoch is 1970-01-01
|
||||
return Ok(PyExpr(df_lit(ScalarValue::Date32(Some(days)))));
|
||||
}
|
||||
|
||||
Err(PyValueError::new_err(format!(
|
||||
"unsupported literal type: {}. Supported: bool, int, float, str, bytes, date, datetime, Decimal",
|
||||
"unsupported literal type: {}. Supported: bool, int, float, str, bytes",
|
||||
value.get_type().name()?
|
||||
)))
|
||||
}
|
||||
|
||||
fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> {
|
||||
let s = s.trim();
|
||||
let dot_pos = s.find('.');
|
||||
let scale = if let Some(pos) = dot_pos {
|
||||
(s.len() - pos - 1) as i8
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let digits = s.replace('.', "");
|
||||
let val = digits
|
||||
.parse::<i128>()
|
||||
.map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?;
|
||||
|
||||
// Precision is total number of digits
|
||||
let precision = digits.trim_start_matches('-').len() as u8;
|
||||
|
||||
Ok((val, precision, scale))
|
||||
}
|
||||
|
||||
/// Call an arbitrary registered SQL function by name.
|
||||
///
|
||||
/// See `lancedb::expr::func` for the list of supported function names.
|
||||
|
||||
+6
-2
@@ -2,7 +2,7 @@
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use arrow::RecordBatchStream;
|
||||
use connection::{Connection, connect, connect_namespace, connect_namespace_client};
|
||||
use connection::{Connection, connect, connect_namespace_client};
|
||||
use env_logger::Env;
|
||||
use expr::{PyExpr, expr_col, expr_func, expr_lit};
|
||||
use index::IndexConfig;
|
||||
@@ -41,6 +41,11 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
.write_style("LANCEDB_LOG_STYLE");
|
||||
env_logger::init_from_env(env);
|
||||
m.add_class::<Connection>()?;
|
||||
m.add_class::<connection::FunctionInfo>()?;
|
||||
m.add_class::<connection::MaterializedViewInfo>()?;
|
||||
m.add_class::<connection::JobInfo>()?;
|
||||
m.add_class::<connection::JobHistoryEntry>()?;
|
||||
m.add_class::<connection::JobErrorEntry>()?;
|
||||
m.add_class::<Session>()?;
|
||||
m.add_class::<Table>()?;
|
||||
m.add_class::<IndexConfig>()?;
|
||||
@@ -62,7 +67,6 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<PyPermutationReader>()?;
|
||||
m.add_class::<PyExpr>()?;
|
||||
m.add_function(wrap_pyfunction!(connect, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(connect_namespace_client, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?;
|
||||
|
||||
@@ -79,14 +79,13 @@ impl PyAsyncPermutationBuilder {
|
||||
|
||||
#[pymethods]
|
||||
impl PyAsyncPermutationBuilder {
|
||||
#[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, clump_size=None, split_names=None))]
|
||||
#[pyo3(signature = (*, ratios=None, counts=None, fixed=None, seed=None, split_names=None))]
|
||||
pub fn split_random(
|
||||
slf: PyRefMut<'_, Self>,
|
||||
ratios: Option<Vec<f64>>,
|
||||
counts: Option<Vec<u64>>,
|
||||
fixed: Option<u64>,
|
||||
seed: Option<u64>,
|
||||
clump_size: Option<u64>,
|
||||
split_names: Option<Vec<String>>,
|
||||
) -> PyResult<Self> {
|
||||
// Check that exactly one split type is provided
|
||||
@@ -112,14 +111,7 @@ impl PyAsyncPermutationBuilder {
|
||||
};
|
||||
|
||||
slf.modify(|builder| {
|
||||
builder.with_split_strategy(
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
},
|
||||
split_names,
|
||||
)
|
||||
builder.with_split_strategy(SplitStrategy::Random { seed, sizes }, split_names)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+79
-16
@@ -17,8 +17,8 @@ use arrow::{
|
||||
pyarrow::{FromPyArrow, PyArrowType, ToPyArrow},
|
||||
};
|
||||
use lancedb::table::{
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform,
|
||||
OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, LoadColumnsRequest,
|
||||
NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable,
|
||||
};
|
||||
use pyo3::{
|
||||
Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python,
|
||||
@@ -322,12 +322,6 @@ impl From<LsmWriteSpec> for lancedb::table::LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<lancedb::table::LsmWriteSpec> for LsmWriteSpec {
|
||||
fn from(inner: lancedb::table::LsmWriteSpec) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(get_all, from_py_object)]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct AddColumnsResult {
|
||||
@@ -1035,14 +1029,6 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_lsm_write_spec(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let spec = inner.get_lsm_write_spec().await.infer_error()?;
|
||||
Ok(spec.map(LsmWriteSpec::from))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn close_lsm_writers(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -1074,6 +1060,83 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_computed_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<(String, String)>,
|
||||
expression: String,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.add_computed_columns(&columns, &expression)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (columns, where_clause=None, num_workers=None, max_workers=None, batch_size=None, priority=None))]
|
||||
pub fn refresh_column(
|
||||
self_: PyRef<'_, Self>,
|
||||
columns: Vec<String>,
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
inner
|
||||
.refresh_column(
|
||||
&columns,
|
||||
where_clause,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
priority,
|
||||
)
|
||||
.await
|
||||
.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
#[pyo3(signature = (source_uris, source_format, target_key, columns, source_key=None, source_storage_options=None, on_missing=None, num_workers=None, max_workers=None, batch_size=None, commit_granularity=None, priority=None))]
|
||||
pub fn load_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
source_uris: Vec<String>,
|
||||
source_format: String,
|
||||
target_key: String,
|
||||
columns: Vec<(String, Option<String>)>,
|
||||
source_key: Option<String>,
|
||||
source_storage_options: Option<std::collections::HashMap<String, String>>,
|
||||
on_missing: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
commit_granularity: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
let request = LoadColumnsRequest {
|
||||
source_uris,
|
||||
source_format,
|
||||
source_storage_options,
|
||||
target_key,
|
||||
source_key,
|
||||
columns,
|
||||
on_missing,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
commit_granularity,
|
||||
priority,
|
||||
};
|
||||
future_into_py(self_.py(), async move {
|
||||
inner.load_columns(request).await.infer_error()
|
||||
})
|
||||
}
|
||||
|
||||
pub fn add_columns(
|
||||
self_: PyRef<'_, Self>,
|
||||
definitions: Vec<(String, String)>,
|
||||
|
||||
@@ -3,14 +3,10 @@
|
||||
|
||||
"""Tests for the type-safe expression builder API."""
|
||||
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from decimal import Decimal
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
from lancedb.expr import Expr, col, func, lit
|
||||
from lancedb.expr import Expr, col, lit, func
|
||||
|
||||
|
||||
# ── unit tests for Expr construction ─────────────────────────────────────────
|
||||
@@ -58,28 +54,6 @@ class TestExprConstruction:
|
||||
with pytest.raises(Exception):
|
||||
func("not_a_real_function", col("x"))
|
||||
|
||||
def test_lit_date(self):
|
||||
e = lit(date(2024, 1, 1))
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_datetime(self):
|
||||
# Naive datetime
|
||||
e = lit(datetime(2024, 1, 1, 10, 0))
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_datetime_tz(self):
|
||||
# Timezone-aware datetime
|
||||
tz = timezone(timedelta(hours=5))
|
||||
dt = datetime(2024, 1, 1, 10, 0, tzinfo=tz)
|
||||
e = lit(dt)
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
def test_lit_decimal_precision(self):
|
||||
# High precision Decimal that would be rounded if converted to float
|
||||
d = Decimal("1.234567890123456789")
|
||||
e = lit(d)
|
||||
assert isinstance(e, Expr)
|
||||
|
||||
|
||||
class TestExprOperators:
|
||||
def test_eq_operator(self):
|
||||
@@ -168,20 +142,6 @@ class TestExprOperators:
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(name = 'alice')"
|
||||
|
||||
def test_reflexive_comparisons(self):
|
||||
# 10 < col("age") swaps to col("age") > 10
|
||||
assert (10 < col("age")).to_sql() == "(age > 10)"
|
||||
assert (10 <= col("age")).to_sql() == "(age >= 10)"
|
||||
assert (10 > col("age")).to_sql() == "(age < 10)"
|
||||
assert (10 >= col("age")).to_sql() == "(age <= 10)"
|
||||
assert (10 == col("age")).to_sql() == "(age = 10)"
|
||||
assert (10 != col("age")).to_sql() == "(age <> 10)"
|
||||
|
||||
def test_reflexive_logical(self):
|
||||
# True & Expr calls Expr.__rand__(True)
|
||||
assert (True & (col("age") > 18)).to_sql() == "(true AND (age > 18))"
|
||||
assert (False | (col("age") > 18)).to_sql() == "(false OR (age > 18))"
|
||||
|
||||
|
||||
class TestExprBytesLiteral:
|
||||
def test_bytes_to_sql(self):
|
||||
@@ -322,40 +282,6 @@ class TestExprRepr:
|
||||
{e: 1}
|
||||
|
||||
|
||||
class TestExprReflexive:
|
||||
def test_reflexive_eq(self):
|
||||
e = 1 == col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x = 1)"
|
||||
|
||||
def test_reflexive_ne(self):
|
||||
e = 1 != col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x <> 1)"
|
||||
|
||||
def test_reflexive_lt(self):
|
||||
# 1 < x => (x > 1)
|
||||
e = 1 < col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x > 1)"
|
||||
|
||||
def test_reflexive_gt(self):
|
||||
# 1 > x => (x < 1)
|
||||
e = 1 > col("x")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(x < 1)"
|
||||
|
||||
def test_reflexive_and(self):
|
||||
e = True & col("active")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(true AND active)"
|
||||
|
||||
def test_reflexive_or(self):
|
||||
e = False | col("inactive")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(false OR inactive)"
|
||||
|
||||
|
||||
# ── integration tests: end-to-end query against a real table ─────────────────
|
||||
|
||||
|
||||
@@ -506,72 +432,6 @@ class TestColNamingIntegration:
|
||||
assert sorted(result["upper_name"].to_pylist()) == ["ALICE", "BOB", "CHARLIE"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def type_check_table(tmp_path):
|
||||
"""Fixture that creates a table with Date32 and Decimal128 columns."""
|
||||
db = lancedb.connect(str(tmp_path))
|
||||
schema = pa.schema(
|
||||
[
|
||||
("date", pa.date32()),
|
||||
("decimal", pa.decimal128(10, 2)),
|
||||
("binary", pa.binary()),
|
||||
]
|
||||
)
|
||||
data = pa.table(
|
||||
{
|
||||
"date": [date(2024, 1, 1), date(2024, 1, 2)],
|
||||
"decimal": [Decimal("10.50"), Decimal("20.75")],
|
||||
"binary": [b"\x01", b"\x02"],
|
||||
},
|
||||
schema=schema,
|
||||
)
|
||||
return db.create_table("extended_types", data)
|
||||
|
||||
|
||||
class TestExtendedTypeIntegration:
|
||||
"""Integration tests verifying that typed literals work correctly in filters."""
|
||||
|
||||
def test_date_integration(self, type_check_table):
|
||||
"""Verify that Date32 literals are correctly parsed and filtered."""
|
||||
result = (
|
||||
type_check_table.search()
|
||||
.where(col("date") == lit(date(2024, 1, 1)))
|
||||
.to_arrow()
|
||||
)
|
||||
assert result.num_rows == 1
|
||||
assert result["date"][0].as_py() == date(2024, 1, 1)
|
||||
|
||||
def test_decimal_integration(self, tmp_path):
|
||||
"""A Decimal literal must retain full 128-bit precision in a filter.
|
||||
|
||||
1.234567890123456789 and 1.234567890123456790 differ only in the last
|
||||
digit and are indistinguishable once truncated to f64. The filter
|
||||
therefore returns the single expected row only if ``lit(Decimal)``
|
||||
produces a true ``Decimal128`` scalar rather than being coerced to f64.
|
||||
"""
|
||||
low = Decimal("1.234567890123456789")
|
||||
high = Decimal("1.234567890123456790")
|
||||
|
||||
db = lancedb.connect(str(tmp_path / "decimal_precision"))
|
||||
schema = pa.schema([("val", pa.decimal128(19, 18))])
|
||||
table = db.create_table(
|
||||
"decimal_precision",
|
||||
pa.table({"val": [low, high]}, schema=schema),
|
||||
)
|
||||
|
||||
result = table.search().where(col("val") < lit(high)).to_arrow()
|
||||
assert result.num_rows == 1
|
||||
assert result["val"][0].as_py() == low
|
||||
|
||||
def test_binary_integration(self, type_check_table):
|
||||
"""Verify that Binary literals are correctly filtered."""
|
||||
result = (
|
||||
type_check_table.search().where(col("binary") == lit(b"\x01")).to_arrow()
|
||||
)
|
||||
assert result.num_rows == 1
|
||||
assert result["binary"][0].as_py() == b"\x01"
|
||||
|
||||
|
||||
# ── bytes / binary column integration tests ───────────────────────────────────
|
||||
|
||||
|
||||
+1
-17
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.31.0-beta.6"
|
||||
version = "0.31.0-beta.4"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
@@ -14,7 +14,6 @@ rust-version.workspace = true
|
||||
ahash = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
arrow-array = { workspace = true }
|
||||
arrow-buffer = { workspace = true }
|
||||
arrow-data = { workspace = true }
|
||||
arrow-schema = { workspace = true }
|
||||
arrow-select = { workspace = true }
|
||||
@@ -108,8 +107,6 @@ http-body = "1" # Matching reqwest
|
||||
rstest = "0.23.0"
|
||||
test-log = "0.2"
|
||||
serial_test = "3"
|
||||
[target.'cfg(unix)'.dev-dependencies]
|
||||
pprof = { version = "0.14", features = ["flamegraph"] }
|
||||
|
||||
|
||||
[features]
|
||||
@@ -128,12 +125,6 @@ azure = [
|
||||
"lance-namespace-impls/dir-azure",
|
||||
"lance-namespace-impls/credential-vendor-azure",
|
||||
]
|
||||
cos = ["lance/tencent", "lance-io/tencent"]
|
||||
goosefs = [
|
||||
"lance/goosefs",
|
||||
"lance-io/goosefs",
|
||||
"lance-namespace-impls/dir-goosefs",
|
||||
]
|
||||
huggingface = [
|
||||
"lance/huggingface",
|
||||
"lance-io/huggingface",
|
||||
@@ -172,16 +163,9 @@ required-features = ["sentence-transformers"]
|
||||
name = "bedrock"
|
||||
required-features = ["bedrock"]
|
||||
|
||||
[[example]]
|
||||
name = "bench_streaming_dataloader"
|
||||
|
||||
[[example]]
|
||||
name = "simple"
|
||||
|
||||
[[example]]
|
||||
name = "polars"
|
||||
required-features = ["polars"]
|
||||
|
||||
[[example]]
|
||||
name = "full_text_search"
|
||||
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -1,47 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
//! This example demonstrates ingesting a Polars DataFrame into LanceDB and
|
||||
//! reading it back out as a Polars DataFrame.
|
||||
|
||||
use lancedb::arrow::IntoPolars;
|
||||
use lancedb::query::ExecutableQuery;
|
||||
use lancedb::{Result, connect};
|
||||
use polars::prelude::{DataFrame, NamedFrom, Series};
|
||||
|
||||
fn make_dataframe() -> DataFrame {
|
||||
let ids = Series::new("id", &[1i32, 2, 3, 4, 5]);
|
||||
let names = Series::new("name", &["Alice", "Bob", "Carol", "Dave", "Eve"]);
|
||||
let scores = Series::new("score", &[9.5f64, 8.1, 7.3, 9.0, 6.5]);
|
||||
DataFrame::new(vec![ids, names, scores]).unwrap()
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let db = connect(tmp.path().to_str().unwrap()).execute().await?;
|
||||
|
||||
// Ingest a Polars DataFrame directly — DataFrame now implements Scannable.
|
||||
let df = make_dataframe();
|
||||
println!("Input DataFrame:\n{df}");
|
||||
|
||||
let table = db.create_table("people", df).execute().await?;
|
||||
|
||||
// Append more rows.
|
||||
let more = DataFrame::new(vec![
|
||||
Series::new("id", &[6i32, 7]),
|
||||
Series::new("name", &["Frank", "Grace"]),
|
||||
Series::new("score", &[7.8f64, 8.9]),
|
||||
])
|
||||
.unwrap();
|
||||
table.add(more).execute().await?;
|
||||
|
||||
// Read back as a Polars DataFrame.
|
||||
let result_df = table.query().execute().await?.into_polars().await?;
|
||||
|
||||
println!(
|
||||
"\nRound-tripped DataFrame ({} rows):\n{result_df}",
|
||||
result_df.height()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
+47
-14
@@ -3,19 +3,7 @@
|
||||
|
||||
use std::{pin::Pin, sync::Arc};
|
||||
|
||||
// Re-export the arrow crates we depend on so downstream consumers can build
|
||||
// `RecordBatch`/arrays/builders against the exact same arrow line lancedb was
|
||||
// compiled against, instead of declaring their own (potentially mismatched)
|
||||
// direct arrow dependencies. See https://github.com/lancedb/lancedb/issues/3575.
|
||||
pub use arrow;
|
||||
pub use arrow_array;
|
||||
pub use arrow_buffer;
|
||||
pub use arrow_cast;
|
||||
pub use arrow_data;
|
||||
pub use arrow_ipc;
|
||||
pub use arrow_ord;
|
||||
pub use arrow_schema;
|
||||
pub use arrow_select;
|
||||
use datafusion_common::DataFusionError;
|
||||
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
|
||||
use futures::{Stream, StreamExt, TryStreamExt};
|
||||
@@ -124,14 +112,54 @@ impl<S: Stream<Item = Result<arrow_array::RecordBatch>>> RecordBatchStream
|
||||
|
||||
/// A trait for converting incoming data to Arrow
|
||||
///
|
||||
/// Integrations should implement this trait to allow data to be
|
||||
/// imported directly from the integration. For example, implementing
|
||||
/// this trait for `Vec<Vec<...>>` would allow the `Vec` to be directly
|
||||
/// used in methods like [`crate::connection::Connection::create_table`]
|
||||
/// or [`crate::table::Table::add`]
|
||||
pub trait IntoArrow {
|
||||
/// Convert the data into an iterator of Arrow batches
|
||||
fn into_arrow(self) -> Result<Box<dyn arrow_array::RecordBatchReader + Send>>;
|
||||
}
|
||||
|
||||
pub type BoxedRecordBatchReader = Box<dyn arrow_array::RecordBatchReader + Send>;
|
||||
|
||||
impl<T: arrow_array::RecordBatchReader + Send + 'static> IntoArrow for T {
|
||||
fn into_arrow(self) -> Result<Box<dyn arrow_array::RecordBatchReader + Send>> {
|
||||
Ok(Box::new(self))
|
||||
}
|
||||
}
|
||||
|
||||
/// A trait for converting incoming data to Arrow asynchronously
|
||||
///
|
||||
/// Serves the same purpose as [`IntoArrow`], but for asynchronous data.
|
||||
///
|
||||
/// Note: Arrow has no async equivalent to RecordBatchReader and so
|
||||
pub trait IntoArrowStream {
|
||||
/// Convert the data into a stream of Arrow batches
|
||||
fn into_arrow(self) -> Result<SendableRecordBatchStream>;
|
||||
}
|
||||
|
||||
impl<S: Stream<Item = Result<arrow_array::RecordBatch>>> SimpleRecordBatchStream<S> {
|
||||
pub fn new(stream: S, schema: Arc<arrow_schema::Schema>) -> Self {
|
||||
Self { schema, stream }
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoArrowStream for SendableRecordBatchStream {
|
||||
fn into_arrow(self) -> Result<SendableRecordBatchStream> {
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoArrowStream for datafusion_physical_plan::SendableRecordBatchStream {
|
||||
fn into_arrow(self) -> Result<SendableRecordBatchStream> {
|
||||
let schema = self.schema();
|
||||
let stream = self.map_err(|df_err| df_err.into());
|
||||
Ok(Box::pin(SimpleRecordBatchStream::new(stream, schema)))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait LanceDbDatagenExt {
|
||||
fn into_ldb_stream(
|
||||
self,
|
||||
@@ -236,7 +264,9 @@ impl IntoPolars for SendableRecordBatchStream {
|
||||
#[cfg(all(test, feature = "polars"))]
|
||||
mod tests {
|
||||
use super::SendableRecordBatchStream;
|
||||
use crate::arrow::{IntoPolars, PolarsDataFrameRecordBatchReader, SimpleRecordBatchStream};
|
||||
use crate::arrow::{
|
||||
IntoArrow, IntoPolars, PolarsDataFrameRecordBatchReader, SimpleRecordBatchStream,
|
||||
};
|
||||
use polars::prelude::{DataFrame, NamedFrom, Series};
|
||||
|
||||
fn get_record_batch_reader_from_polars() -> Box<dyn arrow_array::RecordBatchReader + Send> {
|
||||
@@ -250,7 +280,10 @@ mod tests {
|
||||
float_series = Series::new("float", &[2.0]);
|
||||
let df2 = DataFrame::new(vec![string_series, int_series, float_series]).unwrap();
|
||||
|
||||
Box::new(PolarsDataFrameRecordBatchReader::new(df1.vstack(&df2).unwrap()).unwrap())
|
||||
PolarsDataFrameRecordBatchReader::new(df1.vstack(&df2).unwrap())
|
||||
.unwrap()
|
||||
.into_arrow()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+111
-29
@@ -23,8 +23,10 @@ use crate::connection::create_table::CreateTableBuilder;
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::database::listing::ListingDatabase;
|
||||
use crate::database::{
|
||||
CloneTableRequest, Database, DatabaseOptions, OpenTableRequest, ReadConsistency,
|
||||
TableNamesRequest,
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, Database,
|
||||
DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo, JobInfo, MaterializedViewInfo,
|
||||
MvRefreshPlan, OpenTableRequest, ReadConsistency, RefreshMaterializedViewRequest,
|
||||
TableLineageRequest, TableNamesRequest,
|
||||
};
|
||||
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
|
||||
use crate::error::{Error, Result};
|
||||
@@ -488,6 +490,113 @@ impl Connection {
|
||||
)
|
||||
}
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
// Server-backed features (LanceDB Enterprise / Cloud); local
|
||||
// databases return NotSupported for now.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
pub async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
self.internal.create_function(request).await
|
||||
}
|
||||
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
pub async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
self.internal.list_functions().await
|
||||
}
|
||||
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
pub async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_function(name).await
|
||||
}
|
||||
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
pub async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
self.internal.create_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
pub async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
self.internal.refresh_materialized_view(request).await
|
||||
}
|
||||
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
pub async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
self.internal.table_lineage(request).await
|
||||
}
|
||||
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH).
|
||||
pub async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
self.internal
|
||||
.explain_refresh_materialized_view(name, full, src_version)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
pub async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
self.internal
|
||||
.alter_materialized_view(name, auto_refresh)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
pub async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
self.internal.drop_materialized_view(name).await
|
||||
}
|
||||
|
||||
/// List registered materialized view definitions.
|
||||
pub async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
self.internal.list_materialized_views().await
|
||||
}
|
||||
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
pub async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
self.internal.list_jobs().await
|
||||
}
|
||||
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was flagged for cancellation.
|
||||
pub async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
self.internal.cancel_job(job_id).await
|
||||
}
|
||||
|
||||
/// Look up a single server-side job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table) enables an O(1) server-side lookup; `None`
|
||||
/// scans the database's active jobs. A `None` result means unknown / not
|
||||
/// active.
|
||||
pub async fn get_job(&self, job_id: &str, table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
self.internal.get_job(job_id, table_hint).await
|
||||
}
|
||||
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables.
|
||||
/// Pass `job_id` to narrow to a single job.
|
||||
pub async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
self.internal.job_history(job_id).await
|
||||
}
|
||||
|
||||
/// Per-row UDF errors (SHOW ERRORS) across the database's tables, optionally
|
||||
/// filtered by `job_id` and/or `table`.
|
||||
pub async fn errors(
|
||||
&self,
|
||||
job_id: Option<&str>,
|
||||
table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
self.internal.errors(job_id, table).await
|
||||
}
|
||||
|
||||
/// Rename a table in the database.
|
||||
///
|
||||
/// This is only supported in LanceDB Cloud.
|
||||
@@ -1472,33 +1581,6 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
// A default connection to a commit-engine URI must disable the manifest
|
||||
// namespace: the `__manifest` table under the root has no `ddbTableName`
|
||||
// query, so building its commit handler would fail. See the DynamoDB
|
||||
// regression fixed alongside the Lance v9 bump.
|
||||
#[cfg(feature = "dynamodb")]
|
||||
#[tokio::test]
|
||||
async fn test_default_connect_disables_manifest_for_commit_engine() {
|
||||
for uri in [
|
||||
"s3+ddb://bucket/db?ddbTableName=manifest",
|
||||
"s3://bucket/db?engine=ddb&ddbTableName=manifest",
|
||||
] {
|
||||
let db = connect(uri)
|
||||
.storage_option("region", "us-east-1")
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("expected {uri} to connect, got {e:?}"));
|
||||
|
||||
let (ns_impl, properties) = db.namespace_client_config().await.unwrap();
|
||||
assert_eq!(ns_impl, "dir");
|
||||
assert_eq!(
|
||||
properties.get("manifest_enabled"),
|
||||
Some(&"false".to_string()),
|
||||
"manifest should be disabled for commit-engine URI {uri}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_manifest_enabled_connection_migrates_root_listing_table() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
@@ -185,43 +185,6 @@ impl Scannable for SendableRecordBatchStream {
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "polars")]
|
||||
impl Scannable for polars::frame::DataFrame {
|
||||
fn schema(&self) -> SchemaRef {
|
||||
crate::polars_arrow_convertors::convert_polars_df_schema_to_arrow_rb_schema(
|
||||
self.schema().clone(),
|
||||
)
|
||||
.expect("failed to convert Polars DataFrame schema to Arrow schema")
|
||||
}
|
||||
|
||||
fn scan_as_stream(&mut self) -> SendableRecordBatchStream {
|
||||
let schema = Scannable::schema(self);
|
||||
let batches: crate::Result<Vec<RecordBatch>> =
|
||||
match crate::arrow::PolarsDataFrameRecordBatchReader::new(self.clone()) {
|
||||
Err(e) => Err(e),
|
||||
Ok(reader) => reader.map(|b| b.map_err(Into::into)).collect(),
|
||||
};
|
||||
match batches {
|
||||
Err(e) => Box::pin(SimpleRecordBatchStream {
|
||||
schema,
|
||||
stream: once(async move { Err(e) }),
|
||||
}),
|
||||
Ok(batches) => {
|
||||
let stream = futures::stream::iter(batches.into_iter().map(Ok));
|
||||
Box::pin(SimpleRecordBatchStream { schema, stream })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn num_rows(&self) -> Option<usize> {
|
||||
Some(self.height())
|
||||
}
|
||||
|
||||
fn rescannable(&self) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl StreamingWriteSource for Box<dyn Scannable> {
|
||||
fn arrow_schema(&self) -> SchemaRef {
|
||||
@@ -1126,60 +1089,4 @@ mod tests {
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "polars")]
|
||||
mod polars_tests {
|
||||
use super::*;
|
||||
use crate::arrow::IntoPolars;
|
||||
use crate::query::ExecutableQuery;
|
||||
use polars::prelude::{DataFrame, NamedFrom, Series};
|
||||
|
||||
fn make_df() -> DataFrame {
|
||||
DataFrame::new(vec![
|
||||
Series::new("id", &[1i32, 2, 3]),
|
||||
Series::new("val", &[1.1f64, 2.2, 3.3]),
|
||||
])
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dataframe_scannable_round_trip() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let db = crate::connect(tmp.path().to_str().unwrap())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let df = make_df();
|
||||
let table = db.create_table("t", df.clone()).execute().await.unwrap();
|
||||
|
||||
// Append the same rows again.
|
||||
table.add(df.clone()).execute().await.unwrap();
|
||||
|
||||
let result = table
|
||||
.query()
|
||||
.execute()
|
||||
.await
|
||||
.unwrap()
|
||||
.into_polars()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result.height(), df.height() * 2);
|
||||
assert_eq!(result.schema(), df.schema());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_dataframe_scannable_rescannable() {
|
||||
let mut df = make_df();
|
||||
assert!(df.rescannable());
|
||||
|
||||
let batches1: Vec<RecordBatch> = df.scan_as_stream().try_collect().await.unwrap();
|
||||
assert_eq!(batches1.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
|
||||
|
||||
// Can be scanned again.
|
||||
let batches2: Vec<RecordBatch> = df.scan_as_stream().try_collect().await.unwrap();
|
||||
assert_eq!(batches2.iter().map(|b| b.num_rows()).sum::<usize>(), 3);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -27,7 +27,7 @@ use lance_namespace::models::{
|
||||
};
|
||||
|
||||
use crate::data::scannable::Scannable;
|
||||
use crate::error::Result;
|
||||
use crate::error::{Error, Result};
|
||||
use crate::table::{BaseTable, WriteOptions};
|
||||
|
||||
pub mod listing;
|
||||
@@ -200,6 +200,205 @@ pub enum ReadConsistency {
|
||||
Strong,
|
||||
}
|
||||
|
||||
/// A request to register a UDF (CREATE FUNCTION).
|
||||
///
|
||||
/// Functions are first-class database objects, decoupled from any
|
||||
/// column; computed columns and materialized views reference them by
|
||||
/// name. Server-backed feature (LanceDB Enterprise / Cloud).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateFunctionRequest {
|
||||
/// Function name.
|
||||
pub name: String,
|
||||
/// Implementation language (currently "python").
|
||||
pub language: String,
|
||||
/// SQL return type, e.g. `FLOAT`, `FLOAT[1536]`,
|
||||
/// `STRUCT(a FLOAT, b VARCHAR)`, `TABLE(chunk VARCHAR, idx INT)`.
|
||||
pub return_type: String,
|
||||
/// Function body: source text, or base64 cloudpickle bytes when
|
||||
/// `options["body_format"] = "cloudpickle"`.
|
||||
pub body: String,
|
||||
/// Options: input_columns, pip, num_gpus, batch_size, timeout,
|
||||
/// error_policy, docker_image, body_format, ...
|
||||
pub options: HashMap<String, String>,
|
||||
}
|
||||
|
||||
/// A registered function, as returned by `list_functions`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct FunctionInfo {
|
||||
pub name: String,
|
||||
pub language: String,
|
||||
pub return_type: String,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// A request to create a materialized view (CREATE MATERIALIZED VIEW).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CreateMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// The view's SELECT statement, e.g.
|
||||
/// `SELECT id, embed(body) AS vec FROM articles WHERE id > 1`.
|
||||
/// Bare columns project through; function-call columns compute via
|
||||
/// registered UDFs (a RETURNS TABLE function makes a row-expanding
|
||||
/// chunker view).
|
||||
pub query: String,
|
||||
/// Refresh automatically when the source table changes.
|
||||
pub auto_refresh: bool,
|
||||
/// Register the definition only; skip the initial population.
|
||||
pub with_no_data: bool,
|
||||
/// Optional source column to partition the view's table function on. If the
|
||||
/// column has an IVF vector index the server partitions by its clusters
|
||||
/// (image-dedup style); otherwise it groups by distinct value.
|
||||
pub partition_by: Option<String>,
|
||||
}
|
||||
|
||||
impl CreateMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>, query: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
query: query.into(),
|
||||
auto_refresh: false,
|
||||
with_no_data: false,
|
||||
partition_by: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A request to refresh a materialized view.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RefreshMaterializedViewRequest {
|
||||
/// View name.
|
||||
pub name: String,
|
||||
/// Force a full rebuild (recompute and replace every row) instead of the
|
||||
/// default incremental refresh.
|
||||
pub full: bool,
|
||||
/// Pin the refresh to a source-table version; latest when absent.
|
||||
pub src_version: Option<u64>,
|
||||
/// Initial worker count.
|
||||
pub num_workers: Option<u32>,
|
||||
/// Elastic worker ceiling.
|
||||
pub max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
/// A request for the derived-compute lineage of a table/view (or one of its
|
||||
/// columns). The response is server-defined lineage JSON, returned opaque so
|
||||
/// this client need not model the server's lineage schema.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TableLineageRequest {
|
||||
/// Table or view name.
|
||||
pub name: String,
|
||||
/// Column for column-level lineage; whole table/view when absent.
|
||||
pub column: Option<String>,
|
||||
/// "upstream" | "downstream" | "both" (server default when absent).
|
||||
pub direction: Option<String>,
|
||||
/// Column-hops to walk; transitive when absent.
|
||||
pub depth: Option<u32>,
|
||||
}
|
||||
|
||||
impl RefreshMaterializedViewRequest {
|
||||
pub fn new(name: impl Into<String>) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
full: false,
|
||||
src_version: None,
|
||||
num_workers: None,
|
||||
max_workers: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A registered materialized view definition, as returned by
|
||||
/// `list_materialized_views`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MaterializedViewInfo {
|
||||
pub name: String,
|
||||
pub source_table: String,
|
||||
/// Source columns projected through.
|
||||
pub projection: Vec<String>,
|
||||
/// `alias=expression` per UDF-computed column.
|
||||
pub udf_columns: Vec<String>,
|
||||
pub filter: Option<String>,
|
||||
pub auto_refresh: bool,
|
||||
}
|
||||
|
||||
/// A row from `list_jobs`: one inflight server-side job (index build,
|
||||
/// compaction, column refresh, view refresh, ...).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
/// Lifecycle state: "running", "cancelling", or "stale".
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub age_seconds: Option<i64>,
|
||||
pub command: Option<String>,
|
||||
pub units_done: Option<i64>,
|
||||
pub units_total: Option<i64>,
|
||||
/// Whether the job's final commit has completed (output visible).
|
||||
pub committed: bool,
|
||||
pub rows_skipped: u64,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `job_history`: one durable, completed/terminal server-side job
|
||||
/// record (SHOW JOB HISTORY), read from a table's `_job_history` store. Unlike
|
||||
/// `JobInfo` (live, inflight jobs) this carries created/updated/completed
|
||||
/// timestamps and the lifecycle event log.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobHistoryInfo {
|
||||
pub table: String,
|
||||
pub job_id: String,
|
||||
pub job_type: String,
|
||||
pub state: String,
|
||||
pub column: Option<String>,
|
||||
pub created_ms: i64,
|
||||
pub updated_ms: i64,
|
||||
pub completed_ms: Option<i64>,
|
||||
pub rows_processed: Option<i64>,
|
||||
pub rows_skipped: Option<i64>,
|
||||
pub error: Option<String>,
|
||||
/// Newline-joined lifecycle event log, oldest first.
|
||||
pub events: Option<String>,
|
||||
}
|
||||
|
||||
/// A row from `errors`: one per-row UDF failure recorded by `error_policy=skip`
|
||||
/// (SHOW ERRORS).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JobErrorInfo {
|
||||
pub job_id: String,
|
||||
pub table: String,
|
||||
pub column: String,
|
||||
pub error_type: String,
|
||||
pub error_message: String,
|
||||
pub fragment_id: Option<i64>,
|
||||
pub source_row_id: Option<i64>,
|
||||
pub table_version: Option<i64>,
|
||||
pub age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
/// The plan a `REFRESH MATERIALIZED VIEW` would execute, as returned by
|
||||
/// `explain_refresh_materialized_view` (EXPLAIN REFRESH). No work is run.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MvRefreshPlan {
|
||||
pub table_name: String,
|
||||
/// Whether a refresh would do anything (rebuild or non-empty units).
|
||||
pub has_work: bool,
|
||||
pub source_version: u64,
|
||||
pub last_refreshed_version: Option<u64>,
|
||||
pub full_refresh: bool,
|
||||
/// Source changed non-append-only since the last refresh -> rebuild.
|
||||
pub rebuild: bool,
|
||||
/// Number of row-range work units the refresh would process.
|
||||
pub units_total: u64,
|
||||
}
|
||||
|
||||
fn not_supported<T>(what: &str) -> Result<T> {
|
||||
Err(Error::NotSupported {
|
||||
message: format!("{} is not supported by this database", what),
|
||||
})
|
||||
}
|
||||
|
||||
/// The `Database` trait defines the interface for database implementations.
|
||||
///
|
||||
/// A database is responsible for managing tables and their metadata.
|
||||
@@ -245,6 +444,99 @@ pub trait Database:
|
||||
///
|
||||
/// See [`CloneTableRequest`] for detailed documentation and examples.
|
||||
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
|
||||
// -- Derived compute: functions, materialized views, jobs -------------
|
||||
//
|
||||
// Server-backed features (LanceDB Enterprise / Cloud). The defaults
|
||||
// return NotSupported; the remote database overrides them. Local
|
||||
// single-node implementations are planned.
|
||||
|
||||
/// Register a UDF (CREATE FUNCTION).
|
||||
async fn create_function(&self, _request: CreateFunctionRequest) -> Result<()> {
|
||||
not_supported("create_function")
|
||||
}
|
||||
/// List registered functions (SHOW FUNCTIONS).
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
not_supported("list_functions")
|
||||
}
|
||||
/// Drop a registered function (DROP FUNCTION).
|
||||
async fn drop_function(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_function")
|
||||
}
|
||||
/// Create a materialized view (CREATE MATERIALIZED VIEW). Returns
|
||||
/// the initial-population job id, absent when `with_no_data`.
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
_request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
not_supported("create_materialized_view")
|
||||
}
|
||||
/// Refresh a materialized view; returns the refresh job id.
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
_request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
not_supported("refresh_materialized_view")
|
||||
}
|
||||
/// Derived-compute lineage of a table/view (or column), as server-defined
|
||||
/// JSON. Read-only.
|
||||
async fn table_lineage(&self, _request: TableLineageRequest) -> Result<String> {
|
||||
not_supported("table_lineage")
|
||||
}
|
||||
/// Plan a materialized-view refresh without submitting work
|
||||
/// (EXPLAIN REFRESH). `full` plans a full rebuild (incremental
|
||||
/// planning requires stable row IDs on the source).
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
_name: &str,
|
||||
_full: bool,
|
||||
_src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
not_supported("explain_refresh_materialized_view")
|
||||
}
|
||||
/// Update a materialized view's options (ALTER MATERIALIZED VIEW).
|
||||
async fn alter_materialized_view(&self, _name: &str, _auto_refresh: bool) -> Result<()> {
|
||||
not_supported("alter_materialized_view")
|
||||
}
|
||||
/// Drop a materialized view definition (DROP MATERIALIZED VIEW).
|
||||
async fn drop_materialized_view(&self, _name: &str) -> Result<()> {
|
||||
not_supported("drop_materialized_view")
|
||||
}
|
||||
/// List registered materialized view definitions.
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
not_supported("list_materialized_views")
|
||||
}
|
||||
/// List inflight server-side jobs across the database's tables.
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
not_supported("list_jobs")
|
||||
}
|
||||
/// Cancel an inflight server-side job by id. Returns true if a
|
||||
/// matching inflight job was found and flagged for cancellation,
|
||||
/// false if none was inflight (best-effort, like SQL `CANCEL JOB`).
|
||||
async fn cancel_job(&self, _job_id: &str) -> Result<bool> {
|
||||
not_supported("cancel_job")
|
||||
}
|
||||
/// Point-access for a single job by id -- the `wait()`/status poll path.
|
||||
/// `table_hint` (the job's table, which `wait()` callers know) enables an
|
||||
/// O(1) server-side lookup. `None` if the job is unknown or not active.
|
||||
async fn get_job(&self, _job_id: &str, _table_hint: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
not_supported("get_job")
|
||||
}
|
||||
/// Durable job history (SHOW JOB HISTORY) across the database's tables,
|
||||
/// optionally narrowed to a single `job_id`.
|
||||
async fn job_history(&self, _job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
not_supported("job_history")
|
||||
}
|
||||
/// Per-row UDF errors (SHOW ERRORS) recorded by `error_policy=skip` across
|
||||
/// the database's tables, optionally filtered by `job_id` and/or `table`.
|
||||
async fn errors(
|
||||
&self,
|
||||
_job_id: Option<&str>,
|
||||
_table: Option<&str>,
|
||||
) -> Result<Vec<JobErrorInfo>> {
|
||||
not_supported("errors")
|
||||
}
|
||||
|
||||
/// Open a table in the database
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>>;
|
||||
/// Rename a table in the database
|
||||
|
||||
@@ -323,21 +323,12 @@ impl ListingDatabase {
|
||||
namespace_client_properties: HashMap<String, String>,
|
||||
read_consistency_interval: Option<std::time::Duration>,
|
||||
session: Arc<lance::session::Session>,
|
||||
uses_commit_engine: bool,
|
||||
) -> Result<Arc<LanceNamespaceDatabase>> {
|
||||
let mut ns_properties = Self::build_namespace_client_properties(
|
||||
let ns_properties = Self::build_namespace_client_properties(
|
||||
uri,
|
||||
&storage_options,
|
||||
namespace_client_properties,
|
||||
);
|
||||
// The directory namespace enables manifest-based listing by default, which
|
||||
// opens a `__manifest` table under the root. That is incompatible with an
|
||||
// external commit engine (e.g. `s3+ddb://`): the manifest table URI has no
|
||||
// `ddbTableName` query, so building its commit handler fails. Disable the
|
||||
// manifest so these connections fall back to directory listing.
|
||||
if uses_commit_engine {
|
||||
ns_properties.insert("manifest_enabled".to_string(), "false".to_string());
|
||||
}
|
||||
Ok(Arc::new(
|
||||
LanceNamespaceDatabase::connect(
|
||||
"dir",
|
||||
@@ -512,12 +503,6 @@ impl ListingDatabase {
|
||||
}
|
||||
}
|
||||
|
||||
// A commit engine is in use either via an explicit `engine=` query
|
||||
// param or a `<scheme>+<engine>` URI scheme (e.g. `s3+ddb`). The
|
||||
// manifest namespace is incompatible with these, so track it to
|
||||
// disable the manifest below.
|
||||
let uses_commit_engine = engine.is_some() || url.scheme().contains('+');
|
||||
|
||||
// Filter out the commit store query param -- it's a lancedb param
|
||||
url.query_pairs_mut().clear();
|
||||
url.query_pairs_mut().extend_pairs(filtered_querys);
|
||||
@@ -590,7 +575,6 @@ impl ListingDatabase {
|
||||
request.namespace_client_properties.clone(),
|
||||
request.read_consistency_interval,
|
||||
session.clone(),
|
||||
uses_commit_engine,
|
||||
)
|
||||
.await?;
|
||||
|
||||
@@ -645,7 +629,6 @@ impl ListingDatabase {
|
||||
namespace_client_properties,
|
||||
read_consistency_interval,
|
||||
session.clone(),
|
||||
false,
|
||||
)
|
||||
.await?;
|
||||
|
||||
|
||||
@@ -391,7 +391,6 @@ mod tests {
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Percentages(vec![0.05, 0.30]),
|
||||
clump_size: None,
|
||||
},
|
||||
None,
|
||||
)
|
||||
|
||||
@@ -21,7 +21,6 @@ use arrow::compute::concat_batches;
|
||||
use arrow::datatypes::UInt64Type;
|
||||
use arrow_array::{RecordBatch, UInt64Array};
|
||||
use arrow_schema::SchemaRef;
|
||||
use datafusion_expr::{Expr, col, lit};
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use lance::dataset::scanner::DatasetRecordBatchStream;
|
||||
use lance::io::RecordBatchStream;
|
||||
@@ -197,10 +196,17 @@ impl PermutationReader {
|
||||
.expect_ok()?
|
||||
.values();
|
||||
|
||||
let in_list: Vec<Expr> = row_ids.iter().map(|id| lit(*id)).collect();
|
||||
let filter = format!(
|
||||
"_rowid in ({})",
|
||||
row_ids
|
||||
.iter()
|
||||
.map(|o| o.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
|
||||
let base_query = QueryRequest {
|
||||
filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
|
||||
filter: Some(QueryFilter::Sql(filter)),
|
||||
select: selection,
|
||||
with_row_id: true,
|
||||
..Default::default()
|
||||
|
||||
@@ -35,13 +35,9 @@ pub enum SplitStrategy {
|
||||
/// Rows will be randomly assigned to splits
|
||||
///
|
||||
/// A seed can be provided to make the assignment deterministic.
|
||||
///
|
||||
/// A clump_size can be provided to shuffle contiguous groups of rows together,
|
||||
/// preserving I/O locality while still randomising the split assignment.
|
||||
Random {
|
||||
seed: Option<u64>,
|
||||
sizes: SplitSizes,
|
||||
clump_size: Option<u64>,
|
||||
},
|
||||
/// Rows will be assigned to splits based on the values in the specified columns.
|
||||
///
|
||||
@@ -327,17 +323,13 @@ impl Splitter {
|
||||
self.apply_sequential(source, num_rows, &SplitSizes::Counts(vec![num_rows]))
|
||||
.await
|
||||
}
|
||||
SplitStrategy::Random {
|
||||
seed,
|
||||
sizes,
|
||||
clump_size,
|
||||
} => {
|
||||
SplitStrategy::Random { seed, sizes } => {
|
||||
let shuffler = Shuffler::new(ShufflerConfig {
|
||||
seed: *seed,
|
||||
// In this case we are only shuffling row ids so we can use a large max_rows_per_file
|
||||
max_rows_per_file: 10 * 1024 * 1024,
|
||||
temp_dir: self.temp_dir.clone(),
|
||||
clump_size: *clump_size,
|
||||
clump_size: None,
|
||||
});
|
||||
|
||||
let shuffled = shuffler.shuffle(source, num_rows).await?;
|
||||
@@ -700,7 +692,6 @@ mod tests {
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Fixed(3),
|
||||
clump_size: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -727,7 +718,6 @@ mod tests {
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Counts(vec![5, 15, 10]),
|
||||
clump_size: None,
|
||||
},
|
||||
);
|
||||
|
||||
@@ -754,7 +744,6 @@ mod tests {
|
||||
SplitStrategy::Random {
|
||||
seed: Some(42),
|
||||
sizes: SplitSizes::Percentages(vec![0.217, 0.168, 0.17]),
|
||||
clump_size: None,
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
@@ -342,9 +342,3 @@ pub use connection::connect_namespace;
|
||||
/// Re-export Lance Session and ObjectStoreRegistry for custom session creation
|
||||
pub use lance::session::Session;
|
||||
pub use lance_io::object_store::ObjectStoreRegistry;
|
||||
|
||||
/// Re-export DataFusion so consumers can build the `Expr` values that public
|
||||
/// query/merge APIs (e.g. [`query::QueryBase::only_if_expr`]) accept without
|
||||
/// declaring their own (potentially mismatched) direct `datafusion` dependency.
|
||||
/// See <https://github.com/lancedb/lancedb/issues/3575>.
|
||||
pub use datafusion;
|
||||
|
||||
+21
-156
@@ -7,7 +7,7 @@ use std::{future::Future, time::Duration};
|
||||
use arrow::compute::concat_batches;
|
||||
use arrow_array::{Array, Float16Array, Float32Array, Float64Array, RecordBatch, make_array};
|
||||
use arrow_schema::{DataType, SchemaRef};
|
||||
use datafusion_expr::{Expr, col, lit};
|
||||
use datafusion_expr::Expr;
|
||||
use datafusion_physical_plan::ExecutionPlan;
|
||||
use futures::{FutureExt, TryFutureExt, TryStreamExt, stream, try_join};
|
||||
use half::f16;
|
||||
@@ -401,9 +401,6 @@ pub trait QueryBase {
|
||||
///
|
||||
/// Filtering performance can often be improved by creating a scalar index
|
||||
/// on the filter column(s).
|
||||
///
|
||||
/// Calling this multiple times combines the filters with a logical AND
|
||||
/// (i.e. `(previous) AND (new)`) rather than replacing the previous filter.
|
||||
fn only_if(self, filter: impl AsRef<str>) -> Self;
|
||||
|
||||
/// Only return rows which match the filter, using an expression builder.
|
||||
@@ -426,9 +423,6 @@ pub trait QueryBase {
|
||||
///
|
||||
/// Note: Expression filters are not supported for remote/server-side queries.
|
||||
/// Use [`QueryBase::only_if`] with SQL strings for remote tables.
|
||||
///
|
||||
/// Calling this multiple times combines the expressions with a logical AND
|
||||
/// rather than replacing the previous filter.
|
||||
fn only_if_expr(self, filter: datafusion_expr::Expr) -> Self;
|
||||
|
||||
/// Perform a full text search on the table.
|
||||
@@ -541,13 +535,12 @@ impl<T: HasQuery> QueryBase for T {
|
||||
}
|
||||
|
||||
fn only_if(mut self, filter: impl AsRef<str>) -> Self {
|
||||
self.mut_query()
|
||||
.add_filter(QueryFilter::Sql(filter.as_ref().to_string()));
|
||||
self.mut_query().filter = Some(QueryFilter::Sql(filter.as_ref().to_string()));
|
||||
self
|
||||
}
|
||||
|
||||
fn only_if_expr(mut self, filter: datafusion_expr::Expr) -> Self {
|
||||
self.mut_query().add_filter(QueryFilter::Datafusion(filter));
|
||||
self.mut_query().filter = Some(QueryFilter::Datafusion(filter));
|
||||
self
|
||||
}
|
||||
|
||||
@@ -723,39 +716,6 @@ pub enum QueryFilter {
|
||||
Datafusion(Expr),
|
||||
}
|
||||
|
||||
/// Combine two filters with a logical AND.
|
||||
///
|
||||
/// This is used when a query receives more than one filter (for example when
|
||||
/// `where`/`only_if` is called multiple times) so the filters are composed
|
||||
/// with AND rather than the later filter silently replacing the earlier one.
|
||||
///
|
||||
/// SQL string and expression filters are combined within their own
|
||||
/// representation. When the two representations are mixed, the expression is
|
||||
/// lowered to SQL (via [`crate::expr::expr_to_sql_string`]) and the filters are
|
||||
/// combined as SQL strings. Substrait filters cannot be combined and return an
|
||||
/// error.
|
||||
fn and_filters(existing: QueryFilter, new: QueryFilter) -> Result<QueryFilter> {
|
||||
match (existing, new) {
|
||||
(QueryFilter::Sql(lhs), QueryFilter::Sql(rhs)) => {
|
||||
Ok(QueryFilter::Sql(format!("({lhs}) AND ({rhs})")))
|
||||
}
|
||||
(QueryFilter::Datafusion(lhs), QueryFilter::Datafusion(rhs)) => {
|
||||
Ok(QueryFilter::Datafusion(lhs.and(rhs)))
|
||||
}
|
||||
(QueryFilter::Sql(lhs), QueryFilter::Datafusion(rhs)) => {
|
||||
let rhs = crate::expr::expr_to_sql_string(&rhs)?;
|
||||
Ok(QueryFilter::Sql(format!("({lhs}) AND ({rhs})")))
|
||||
}
|
||||
(QueryFilter::Datafusion(lhs), QueryFilter::Sql(rhs)) => {
|
||||
let lhs = crate::expr::expr_to_sql_string(&lhs)?;
|
||||
Ok(QueryFilter::Sql(format!("({lhs}) AND ({rhs})")))
|
||||
}
|
||||
_ => Err(Error::InvalidInput {
|
||||
message: "cannot combine a Substrait filter with another filter".to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// A basic query into a table without any kind of search
|
||||
///
|
||||
/// This will result in a (potentially filtered) scan if executed
|
||||
@@ -770,13 +730,6 @@ pub struct QueryRequest {
|
||||
/// Apply filter to the returned rows.
|
||||
pub filter: Option<QueryFilter>,
|
||||
|
||||
/// An error recorded while combining repeated filters that could not be
|
||||
/// composed (see [`QueryRequest::add_filter`]). It is surfaced when the
|
||||
/// query is executed via [`QueryRequest::check_filter`]. We defer the error
|
||||
/// because the builder methods that set filters return `Self` rather than a
|
||||
/// `Result`.
|
||||
pub(crate) filter_error: Option<String>,
|
||||
|
||||
/// Perform a full text search on the table.
|
||||
pub full_text_search: Option<FullTextSearchQuery>,
|
||||
|
||||
@@ -822,7 +775,6 @@ impl Default for QueryRequest {
|
||||
limit: None,
|
||||
offset: None,
|
||||
filter: None,
|
||||
filter_error: None,
|
||||
full_text_search: None,
|
||||
select: Select::All,
|
||||
fast_search: false,
|
||||
@@ -836,41 +788,6 @@ impl Default for QueryRequest {
|
||||
}
|
||||
}
|
||||
|
||||
impl QueryRequest {
|
||||
/// Add a filter, combining it with any existing filter using a logical AND.
|
||||
///
|
||||
/// If the new filter cannot be combined with the existing one (because they
|
||||
/// use different representations) the error is recorded and surfaced later
|
||||
/// by [`Self::check_filter`].
|
||||
pub(crate) fn add_filter(&mut self, new: QueryFilter) {
|
||||
self.filter = Some(match self.filter.take() {
|
||||
None => new,
|
||||
Some(existing) => match and_filters(existing, new) {
|
||||
Ok(combined) => combined,
|
||||
Err(err) => {
|
||||
// The filters were consumed while attempting to combine
|
||||
// them; the recorded error is surfaced by `check_filter`
|
||||
// before the query executes.
|
||||
self.filter_error = Some(err.to_string());
|
||||
return;
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
/// Return an error if combining filters failed (see [`Self::add_filter`]).
|
||||
///
|
||||
/// This must be called by every backend before executing a query.
|
||||
pub(crate) fn check_filter(&self) -> Result<()> {
|
||||
if let Some(message) = &self.filter_error {
|
||||
return Err(Error::InvalidInput {
|
||||
message: message.clone(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// A builder for LanceDB queries.
|
||||
///
|
||||
/// See [`crate::Table::query`] for more details on queries
|
||||
@@ -1468,13 +1385,18 @@ impl TakeQuery {
|
||||
///
|
||||
/// See [`crate::Table::take_offsets`] for more details.
|
||||
pub fn from_offsets(parent: Arc<dyn BaseTable>, offsets: Vec<u64>) -> Self {
|
||||
let in_list: Vec<Expr> = offsets.iter().map(|o| lit(*o)).collect();
|
||||
let filter = format!(
|
||||
"_rowoffset in ({})",
|
||||
offsets
|
||||
.iter()
|
||||
.map(|o| o.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
Self {
|
||||
parent,
|
||||
request: QueryRequest {
|
||||
filter: Some(QueryFilter::Datafusion(
|
||||
col("_rowoffset").in_list(in_list, false),
|
||||
)),
|
||||
filter: Some(QueryFilter::Sql(filter)),
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
@@ -1484,11 +1406,18 @@ impl TakeQuery {
|
||||
///
|
||||
/// See [`crate::Table::take_row_ids`] for more details.
|
||||
pub fn from_row_ids(parent: Arc<dyn BaseTable>, row_ids: Vec<u64>) -> Self {
|
||||
let in_list: Vec<Expr> = row_ids.iter().map(|id| lit(*id)).collect();
|
||||
let filter = format!(
|
||||
"_rowid in ({})",
|
||||
row_ids
|
||||
.iter()
|
||||
.map(|o| o.to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join(",")
|
||||
);
|
||||
Self {
|
||||
parent,
|
||||
request: QueryRequest {
|
||||
filter: Some(QueryFilter::Datafusion(col(ROW_ID).in_list(in_list, false))),
|
||||
filter: Some(QueryFilter::Sql(filter)),
|
||||
..Default::default()
|
||||
},
|
||||
}
|
||||
@@ -1753,70 +1682,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_repeated_only_if_combines_with_and() {
|
||||
use crate::expr::{col, lit};
|
||||
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let dataset_path = tmp_dir.path().join("test.lance");
|
||||
let uri = dataset_path.to_str().unwrap();
|
||||
|
||||
let conn = connect(uri).execute().await.unwrap();
|
||||
let table = conn
|
||||
.create_table("my_table", make_non_empty_batches())
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let query = table.query().only_if("id > 0").only_if("id < 100");
|
||||
match &query.request.filter {
|
||||
Some(QueryFilter::Sql(sql)) => assert_eq!(sql, "(id > 0) AND (id < 100)"),
|
||||
other => panic!("expected combined SQL filter, got {other:?}"),
|
||||
}
|
||||
|
||||
// A single filter is left untouched.
|
||||
let query = table.query().only_if("id > 0");
|
||||
match &query.request.filter {
|
||||
Some(QueryFilter::Sql(sql)) => assert_eq!(sql, "id > 0"),
|
||||
other => panic!("expected single SQL filter, got {other:?}"),
|
||||
}
|
||||
|
||||
// Expression filters are combined with a logical AND as well.
|
||||
let query = table
|
||||
.query()
|
||||
.only_if_expr(col("id").gt(lit(0i32)))
|
||||
.only_if_expr(col("id").lt(lit(100i32)));
|
||||
match &query.request.filter {
|
||||
Some(QueryFilter::Datafusion(expr)) => {
|
||||
assert_eq!(
|
||||
expr,
|
||||
&col("id").gt(lit(0i32)).and(col("id").lt(lit(100i32)))
|
||||
);
|
||||
}
|
||||
other => panic!("expected combined Datafusion filter, got {other:?}"),
|
||||
}
|
||||
|
||||
// Mixing an SQL string filter with an expression filter lowers the
|
||||
// expression to SQL and combines them as SQL strings.
|
||||
let query = table
|
||||
.query()
|
||||
.only_if("id > 0")
|
||||
.only_if_expr(col("id").lt(lit(100i32)));
|
||||
match &query.request.filter {
|
||||
Some(QueryFilter::Sql(sql)) => {
|
||||
let expected = format!(
|
||||
"(id > 0) AND ({})",
|
||||
crate::expr::expr_to_sql_string(&col("id").lt(lit(100i32))).unwrap()
|
||||
);
|
||||
assert_eq!(sql, &expected);
|
||||
}
|
||||
other => panic!("expected combined SQL filter, got {other:?}"),
|
||||
}
|
||||
assert!(query.request.check_filter().is_ok());
|
||||
// The combined filter executes without error.
|
||||
query.execute().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_select_with_transform() {
|
||||
// TODO: Switch back to memory://foo after https://github.com/lancedb/lancedb/issues/1051
|
||||
|
||||
@@ -19,8 +19,10 @@ use lance_namespace::models::{
|
||||
|
||||
use crate::Error;
|
||||
use crate::database::{
|
||||
CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
|
||||
OpenTableRequest, ReadConsistency, TableNamesRequest,
|
||||
CloneTableRequest, CreateFunctionRequest, CreateMaterializedViewRequest, CreateTableMode,
|
||||
CreateTableRequest, Database, DatabaseOptions, FunctionInfo, JobErrorInfo, JobHistoryInfo,
|
||||
JobInfo, MaterializedViewInfo, MvRefreshPlan, OpenTableRequest, ReadConsistency,
|
||||
RefreshMaterializedViewRequest, TableLineageRequest, TableNamesRequest,
|
||||
};
|
||||
use crate::error::Result;
|
||||
use crate::remote::util::stream_as_body;
|
||||
@@ -33,6 +35,248 @@ use super::client::{
|
||||
use super::table::RemoteTable;
|
||||
use super::util::parse_server_version;
|
||||
|
||||
// Wire types for the derived-compute routes (functions, materialized
|
||||
// views, jobs). Field shapes mirror the server's REST contract.
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateFunctionRequest {
|
||||
language: String,
|
||||
return_type: String,
|
||||
body: String,
|
||||
options: std::collections::HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteFunctionEntry {
|
||||
name: String,
|
||||
language: String,
|
||||
return_type: String,
|
||||
#[serde(default)]
|
||||
description: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListFunctionsResponse {
|
||||
functions: Vec<RemoteFunctionEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCreateMaterializedViewRequest {
|
||||
query: String,
|
||||
auto_refresh: bool,
|
||||
with_no_data: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
partition_by: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteCreateMaterializedViewResponse {
|
||||
#[serde(default)]
|
||||
job_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteRefreshMaterializedViewRequest {
|
||||
#[serde(skip_serializing_if = "std::ops::Not::not")]
|
||||
full: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
num_workers: Option<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
max_workers: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteRefreshMaterializedViewResponse {
|
||||
job_id: String,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteExplainRefreshRequest {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
full: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
src_version: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteExplainRefreshResponse {
|
||||
table_name: String,
|
||||
has_work: bool,
|
||||
source_version: u64,
|
||||
last_refreshed_version: Option<u64>,
|
||||
full_refresh: bool,
|
||||
rebuild: bool,
|
||||
units_total: u64,
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteAlterMaterializedViewRequest {
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteMaterializedViewEntry {
|
||||
name: String,
|
||||
source_table: String,
|
||||
#[serde(default)]
|
||||
projection: Vec<String>,
|
||||
#[serde(default)]
|
||||
udf_columns: Vec<String>,
|
||||
#[serde(default)]
|
||||
filter: Option<String>,
|
||||
#[serde(default)]
|
||||
auto_refresh: bool,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListMaterializedViewsResponse {
|
||||
views: Vec<RemoteMaterializedViewEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteJobEntry {
|
||||
table: String,
|
||||
job_id: String,
|
||||
job_type: String,
|
||||
state: String,
|
||||
#[serde(default)]
|
||||
column: Option<String>,
|
||||
#[serde(default)]
|
||||
age_seconds: Option<i64>,
|
||||
#[serde(default)]
|
||||
command: Option<String>,
|
||||
#[serde(default)]
|
||||
units_done: Option<i64>,
|
||||
#[serde(default)]
|
||||
units_total: Option<i64>,
|
||||
#[serde(default)]
|
||||
committed: bool,
|
||||
#[serde(default)]
|
||||
rows_skipped: u64,
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteListJobsResponse {
|
||||
jobs: Vec<RemoteJobEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteGetJobResponse {
|
||||
#[serde(default)]
|
||||
job: Option<RemoteJobEntry>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteCancelJobResponse {
|
||||
cancelled: bool,
|
||||
}
|
||||
|
||||
impl From<RemoteJobEntry> for JobInfo {
|
||||
fn from(j: RemoteJobEntry) -> Self {
|
||||
JobInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
age_seconds: j.age_seconds,
|
||||
command: j.command,
|
||||
units_done: j.units_done,
|
||||
units_total: j.units_total,
|
||||
committed: j.committed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteJobHistoryEntry {
|
||||
table: String,
|
||||
job_id: String,
|
||||
job_type: String,
|
||||
state: String,
|
||||
#[serde(default)]
|
||||
column: Option<String>,
|
||||
created_ms: i64,
|
||||
updated_ms: i64,
|
||||
#[serde(default)]
|
||||
completed_ms: Option<i64>,
|
||||
#[serde(default)]
|
||||
rows_processed: Option<i64>,
|
||||
#[serde(default)]
|
||||
rows_skipped: Option<i64>,
|
||||
#[serde(default)]
|
||||
error: Option<String>,
|
||||
#[serde(default)]
|
||||
events: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteJobHistoryResponse {
|
||||
jobs: Vec<RemoteJobHistoryEntry>,
|
||||
}
|
||||
|
||||
impl From<RemoteJobHistoryEntry> for JobHistoryInfo {
|
||||
fn from(j: RemoteJobHistoryEntry) -> Self {
|
||||
JobHistoryInfo {
|
||||
table: j.table,
|
||||
job_id: j.job_id,
|
||||
job_type: j.job_type,
|
||||
state: j.state,
|
||||
column: j.column,
|
||||
created_ms: j.created_ms,
|
||||
updated_ms: j.updated_ms,
|
||||
completed_ms: j.completed_ms,
|
||||
rows_processed: j.rows_processed,
|
||||
rows_skipped: j.rows_skipped,
|
||||
error: j.error,
|
||||
events: j.events,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorEntry {
|
||||
job_id: String,
|
||||
table: String,
|
||||
column: String,
|
||||
error_type: String,
|
||||
error_message: String,
|
||||
#[serde(default)]
|
||||
fragment_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
source_row_id: Option<i64>,
|
||||
#[serde(default)]
|
||||
table_version: Option<i64>,
|
||||
#[serde(default)]
|
||||
age_seconds: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RemoteErrorsResponse {
|
||||
errors: Vec<RemoteErrorEntry>,
|
||||
}
|
||||
|
||||
impl From<RemoteErrorEntry> for JobErrorInfo {
|
||||
fn from(e: RemoteErrorEntry) -> Self {
|
||||
JobErrorInfo {
|
||||
job_id: e.job_id,
|
||||
table: e.table,
|
||||
column: e.column,
|
||||
error_type: e.error_type,
|
||||
error_message: e.error_message,
|
||||
fragment_id: e.fragment_id,
|
||||
source_row_id: e.source_row_id,
|
||||
table_version: e.table_version,
|
||||
age_seconds: e.age_seconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Request structure for the remote clone table API
|
||||
#[derive(serde::Serialize)]
|
||||
struct RemoteCloneTableRequest {
|
||||
@@ -641,6 +885,228 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
|
||||
Ok(table)
|
||||
}
|
||||
|
||||
async fn create_function(&self, request: CreateFunctionRequest) -> Result<()> {
|
||||
let body = RemoteCreateFunctionRequest {
|
||||
language: request.language,
|
||||
return_type: request.return_type,
|
||||
body: request.body,
|
||||
options: request.options,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/function/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_functions(&self) -> Result<Vec<FunctionInfo>> {
|
||||
let req = self.client.get("/v1/function/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListFunctionsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.functions
|
||||
.into_iter()
|
||||
.map(|f| FunctionInfo {
|
||||
name: f.name,
|
||||
language: f.language,
|
||||
return_type: f.return_type,
|
||||
description: f.description,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn drop_function(&self, name: &str) -> Result<()> {
|
||||
let req = self.client.post(&format!("/v1/function/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn create_materialized_view(
|
||||
&self,
|
||||
request: CreateMaterializedViewRequest,
|
||||
) -> Result<Option<String>> {
|
||||
let body = RemoteCreateMaterializedViewRequest {
|
||||
query: request.query,
|
||||
auto_refresh: request.auto_refresh,
|
||||
with_no_data: request.with_no_data,
|
||||
partition_by: request.partition_by,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/create", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteCreateMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn refresh_materialized_view(
|
||||
&self,
|
||||
request: RefreshMaterializedViewRequest,
|
||||
) -> Result<String> {
|
||||
let body = RemoteRefreshMaterializedViewRequest {
|
||||
full: request.full,
|
||||
src_version: request.src_version,
|
||||
num_workers: request.num_workers,
|
||||
max_workers: request.max_workers,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/refresh", request.name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteRefreshMaterializedViewResponse =
|
||||
rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn table_lineage(&self, request: TableLineageRequest) -> Result<String> {
|
||||
let mut req = self
|
||||
.client
|
||||
.get(&format!("/v1/table/{}/lineage", request.name));
|
||||
if let Some(column) = &request.column {
|
||||
req = req.query(&[("column", column)]);
|
||||
}
|
||||
if let Some(direction) = &request.direction {
|
||||
req = req.query(&[("direction", direction)]);
|
||||
}
|
||||
if let Some(depth) = request.depth {
|
||||
req = req.query(&[("depth", depth.to_string())]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
// Server-defined lineage JSON, returned opaque (the client does not
|
||||
// model the lineage schema; the Python layer deserializes it).
|
||||
rsp.text().await.err_to_http(request_id)
|
||||
}
|
||||
|
||||
async fn explain_refresh_materialized_view(
|
||||
&self,
|
||||
name: &str,
|
||||
full: bool,
|
||||
src_version: Option<u64>,
|
||||
) -> Result<MvRefreshPlan> {
|
||||
let body = RemoteExplainRefreshRequest {
|
||||
full: Some(full),
|
||||
src_version,
|
||||
};
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/explain_refresh", name))
|
||||
.json(&body);
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteExplainRefreshResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(MvRefreshPlan {
|
||||
table_name: body.table_name,
|
||||
has_work: body.has_work,
|
||||
source_version: body.source_version,
|
||||
last_refreshed_version: body.last_refreshed_version,
|
||||
full_refresh: body.full_refresh,
|
||||
rebuild: body.rebuild,
|
||||
units_total: body.units_total,
|
||||
})
|
||||
}
|
||||
|
||||
async fn alter_materialized_view(&self, name: &str, auto_refresh: bool) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/alter", name))
|
||||
.json(&RemoteAlterMaterializedViewRequest { auto_refresh });
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_materialized_view(&self, name: &str) -> Result<()> {
|
||||
let req = self
|
||||
.client
|
||||
.post(&format!("/v1/materialized_view/{}/drop", name));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
self.client.check_response(&request_id, rsp).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn list_materialized_views(&self) -> Result<Vec<MaterializedViewInfo>> {
|
||||
let req = self.client.get("/v1/materialized_view/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListMaterializedViewsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body
|
||||
.views
|
||||
.into_iter()
|
||||
.map(|v| MaterializedViewInfo {
|
||||
name: v.name,
|
||||
source_table: v.source_table,
|
||||
projection: v.projection,
|
||||
udf_columns: v.udf_columns,
|
||||
filter: v.filter,
|
||||
auto_refresh: v.auto_refresh,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn list_jobs(&self) -> Result<Vec<JobInfo>> {
|
||||
let req = self.client.get("/v1/job/list");
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteListJobsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.jobs.into_iter().map(JobInfo::from).collect())
|
||||
}
|
||||
|
||||
async fn get_job(&self, job_id: &str, table: Option<&str>) -> Result<Option<JobInfo>> {
|
||||
// Point-access poll path: GET /v1/job/{id}, with the table as the O(1)
|
||||
// hint when known. `query` handles URL-encoding the table name.
|
||||
let mut req = self.client.get(&format!("/v1/job/{job_id}"));
|
||||
if let Some(t) = table {
|
||||
req = req.query(&[("table", t)]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteGetJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job.map(JobInfo::from))
|
||||
}
|
||||
|
||||
async fn cancel_job(&self, job_id: &str) -> Result<bool> {
|
||||
let req = self.client.post(&format!("/v1/job/{}/cancel", job_id));
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteCancelJobResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.cancelled)
|
||||
}
|
||||
|
||||
async fn job_history(&self, job_id: Option<&str>) -> Result<Vec<JobHistoryInfo>> {
|
||||
let mut req = self.client.get("/v1/job/history");
|
||||
if let Some(j) = job_id {
|
||||
req = req.query(&[("job", j)]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteJobHistoryResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.jobs.into_iter().map(JobHistoryInfo::from).collect())
|
||||
}
|
||||
|
||||
async fn errors(&self, job_id: Option<&str>, table: Option<&str>) -> Result<Vec<JobErrorInfo>> {
|
||||
let mut req = self.client.get("/v1/job/errors");
|
||||
if let Some(j) = job_id {
|
||||
req = req.query(&[("job", j)]);
|
||||
}
|
||||
if let Some(t) = table {
|
||||
req = req.query(&[("table", t)]);
|
||||
}
|
||||
let (request_id, rsp) = self.client.send(req).await?;
|
||||
let rsp = self.client.check_response(&request_id, rsp).await?;
|
||||
let body: RemoteErrorsResponse = rsp.json().await.err_to_http(request_id)?;
|
||||
Ok(body.errors.into_iter().map(JobErrorInfo::from).collect())
|
||||
}
|
||||
|
||||
async fn open_table(&self, request: OpenTableRequest) -> Result<Arc<dyn BaseTable>> {
|
||||
let identifier = build_table_identifier(
|
||||
&request.name,
|
||||
@@ -1580,6 +2046,223 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_derived_compute_routes() {
|
||||
// create_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["language"], "python");
|
||||
assert_eq!(body["return_type"], "FLOAT[4]");
|
||||
assert_eq!(body["body"], "def embed(x): ...");
|
||||
assert_eq!(body["options"]["pip"], "torch");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.create_function(crate::database::CreateFunctionRequest {
|
||||
name: "embed".into(),
|
||||
language: "python".into(),
|
||||
return_type: "FLOAT[4]".into(),
|
||||
body: "def embed(x): ...".into(),
|
||||
options: [("pip".to_string(), "torch".to_string())].into(),
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// list_functions
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/function/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"functions":[{"name":"embed","language":"python","return_type":"Float32","description":""}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let functions = conn.list_functions().await.unwrap();
|
||||
assert_eq!(functions.len(), 1);
|
||||
assert_eq!(functions[0].name, "embed");
|
||||
|
||||
// drop_function
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/function/embed/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"embed","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_function("embed").await.unwrap();
|
||||
|
||||
// create_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/create");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["query"], "SELECT id, embed(body) AS vec FROM docs");
|
||||
assert_eq!(body["auto_refresh"], true);
|
||||
assert_eq!(body["with_no_data"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","job_id":"j-1"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::CreateMaterializedViewRequest::new(
|
||||
"mv1",
|
||||
"SELECT id, embed(body) AS vec FROM docs",
|
||||
);
|
||||
request.auto_refresh = true;
|
||||
let job_id = conn.create_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id.as_deref(), Some("j-1"));
|
||||
|
||||
// refresh_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/refresh");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("src_version").is_none());
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-2"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
let mut request = crate::database::RefreshMaterializedViewRequest::new("mv1");
|
||||
request.num_workers = Some(2);
|
||||
let job_id = conn.refresh_materialized_view(request).await.unwrap();
|
||||
assert_eq!(job_id, "j-2");
|
||||
|
||||
// alter_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/alter");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["auto_refresh"], false);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.alter_materialized_view("mv1", false).await.unwrap();
|
||||
|
||||
// drop_materialized_view
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/mv1/drop");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"name":"mv1","status":"OK"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
conn.drop_materialized_view("mv1").await.unwrap();
|
||||
|
||||
// list_materialized_views
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/materialized_view/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"views":[{"name":"mv1","source_table":"docs","projection":["id"],"udf_columns":["vec=embed(body)"],"filter":null,"auto_refresh":true}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let views = conn.list_materialized_views().await.unwrap();
|
||||
assert_eq!(views.len(), 1);
|
||||
assert_eq!(views[0].source_table, "docs");
|
||||
assert!(views[0].auto_refresh);
|
||||
|
||||
// list_jobs
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/job/list");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"table":"docs","job_id":"j-3","job_type":"udf_virtual_column_backfill","state":"running","column":"vec","age_seconds":4,"command":null,"units_done":1,"units_total":2,"committed":false,"rows_skipped":0,"error":null}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let jobs = conn.list_jobs().await.unwrap();
|
||||
assert_eq!(jobs.len(), 1);
|
||||
assert_eq!(jobs[0].state, "running");
|
||||
assert_eq!(jobs[0].units_total, Some(2));
|
||||
|
||||
// cancel_job
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::POST);
|
||||
assert_eq!(request.url().path(), "/v1/job/j-3/cancel");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"cancelled":true}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(conn.cancel_job("j-3").await.unwrap());
|
||||
|
||||
// cancel_job: no such inflight job -> false, not an error
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/job/gone/cancel");
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"cancelled":false}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(!conn.cancel_job("gone").await.unwrap());
|
||||
|
||||
// job_history: GET /v1/job/history, no filter
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/job/history");
|
||||
assert!(request.url().query().is_none());
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"jobs":[{"table":"docs","job_id":"j-1","job_type":"udf_virtual_column_backfill","state":"done","column":"vec","created_ms":1000,"updated_ms":2000,"completed_ms":2000,"rows_processed":42,"rows_skipped":3,"error":null,"events":"created\ndone"}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let hist = conn.job_history(None).await.unwrap();
|
||||
assert_eq!(hist.len(), 1);
|
||||
assert_eq!(hist[0].state, "done");
|
||||
assert_eq!(hist[0].rows_processed, Some(42));
|
||||
assert_eq!(hist[0].events.as_deref(), Some("created\ndone"));
|
||||
|
||||
// job_history: ?job= narrows to one job
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.url().path(), "/v1/job/history");
|
||||
assert_eq!(request.url().query(), Some("job=j-1"));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(r#"{"jobs":[]}"#)
|
||||
.unwrap()
|
||||
});
|
||||
assert!(conn.job_history(Some("j-1")).await.unwrap().is_empty());
|
||||
|
||||
// errors: GET /v1/job/errors with job + table filters
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
assert_eq!(request.method(), &reqwest::Method::GET);
|
||||
assert_eq!(request.url().path(), "/v1/job/errors");
|
||||
assert_eq!(request.url().query(), Some("job=j-1&table=docs"));
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(
|
||||
r#"{"errors":[{"job_id":"j-1","table":"docs","column":"vec","error_type":"ValueError","error_message":"boom","fragment_id":0,"source_row_id":42,"table_version":7,"age_seconds":5}]}"#,
|
||||
)
|
||||
.unwrap()
|
||||
});
|
||||
let errs = conn.errors(Some("j-1"), Some("docs")).await.unwrap();
|
||||
assert_eq!(errs.len(), 1);
|
||||
assert_eq!(errs[0].error_type, "ValueError");
|
||||
assert_eq!(errs[0].source_row_id, Some(42));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_clone_table() {
|
||||
let conn = Connection::new_with_handler(|request| {
|
||||
|
||||
+194
-283
@@ -20,7 +20,6 @@ use crate::table::AddColumnsResult;
|
||||
use crate::table::AddResult;
|
||||
use crate::table::DeleteResult;
|
||||
use crate::table::DropColumnsResult;
|
||||
use crate::table::LsmWriteSpec;
|
||||
use crate::table::MergeResult;
|
||||
use crate::table::Tags;
|
||||
use crate::table::UpdateResult;
|
||||
@@ -71,29 +70,18 @@ use tokio::sync::RwLock;
|
||||
const REQUEST_TIMEOUT_HEADER: HeaderName = HeaderName::from_static("x-request-timeout-ms");
|
||||
const MIN_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-version");
|
||||
const MIN_TIMESTAMP_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-timestamp");
|
||||
const MIN_READ_VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-min-read-version");
|
||||
const VERSION_HEADER: HeaderName = HeaderName::from_static("x-lancedb-version");
|
||||
const METRIC_TYPE_KEY: &str = "metric_type";
|
||||
const INDEX_TYPE_KEY: &str = "index_type";
|
||||
const SCHEMA_CACHE_TTL: Duration = Duration::from_secs(30);
|
||||
const SCHEMA_CACHE_REFRESH_WINDOW: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Per-table state driving the freshness headers (`x-lancedb-min-version`,
|
||||
/// `x-lancedb-min-timestamp`, and `x-lancedb-min-read-version`) sent on read
|
||||
/// requests.
|
||||
/// Per-table state driving the freshness headers (`x-lancedb-min-version` and
|
||||
/// `x-lancedb-min-timestamp`) sent on read requests.
|
||||
#[derive(Debug, Default, Clone, Copy)]
|
||||
struct FreshnessState {
|
||||
/// Provides read-your-write within a single handle: writes that return a
|
||||
/// version update this, and reads send it as `x-lancedb-min-version`.
|
||||
min_version: Option<u64>,
|
||||
/// Highest dataset version observed in a *read* response on this handle.
|
||||
/// Reads send it as `x-lancedb-min-read-version` so a load-balanced query
|
||||
/// node whose cache is behind this version must refresh before serving,
|
||||
/// giving monotonic reads across nodes regardless of which one the load
|
||||
/// balancer routes to. Sourced only from reads (always committed dataset
|
||||
/// versions), never from writes (which may return WAL entry ids), so it is
|
||||
/// unaffected by the WAL/version mismatch that retired `min_version`.
|
||||
min_read_version: Option<u64>,
|
||||
/// Wall-clock time captured at the last [`BaseTable::checkout_latest`]
|
||||
/// call. Subsequent reads send
|
||||
/// `max(baseline, now - read_consistency_interval)` as
|
||||
@@ -114,7 +102,6 @@ struct FreshnessState {
|
||||
struct FreshnessHeaders {
|
||||
min_version: Option<u64>,
|
||||
min_timestamp: Option<SystemTime>,
|
||||
min_read_version: Option<u64>,
|
||||
}
|
||||
|
||||
impl FreshnessHeaders {
|
||||
@@ -126,9 +113,6 @@ impl FreshnessHeaders {
|
||||
let dt: chrono::DateTime<chrono::Utc> = ts.into();
|
||||
request = request.header(MIN_TIMESTAMP_HEADER, dt.to_rfc3339());
|
||||
}
|
||||
if let Some(v) = self.min_read_version {
|
||||
request = request.header(MIN_READ_VERSION_HEADER, v.to_string());
|
||||
}
|
||||
request
|
||||
}
|
||||
}
|
||||
@@ -613,7 +597,6 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
body: &mut serde_json::Value,
|
||||
params: &QueryRequest,
|
||||
) -> Result<()> {
|
||||
params.check_filter()?;
|
||||
body["prefilter"] = params.prefilter.into();
|
||||
if let Some(offset) = params.offset {
|
||||
body["offset"] = serde_json::Value::Number(serde_json::Number::from(offset));
|
||||
@@ -901,7 +884,6 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
self.client.read_consistency_interval,
|
||||
SystemTime::now(),
|
||||
),
|
||||
min_read_version: state.min_read_version,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -923,30 +905,6 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
state.min_version = Some(state.min_version.map_or(version, |v| v.max(version)));
|
||||
}
|
||||
|
||||
/// Record a dataset version observed in a *read* response so subsequent
|
||||
/// reads request at least this version via `x-lancedb-min-read-version`,
|
||||
/// giving monotonic reads across load-balanced query nodes. A returned `0`
|
||||
/// (or absent header from an old server) is ignored.
|
||||
fn track_read_version(&self, version: u64) {
|
||||
if version == 0 {
|
||||
return;
|
||||
}
|
||||
let mut state = self.freshness.lock().unwrap();
|
||||
state.min_read_version = Some(state.min_read_version.map_or(version, |v| v.max(version)));
|
||||
}
|
||||
|
||||
/// Parse the `x-lancedb-version` response header (the dataset version a read
|
||||
/// reflects) and fold it into the read-version watermark.
|
||||
fn track_read_version_from_headers(&self, headers: &reqwest::header::HeaderMap) {
|
||||
if let Some(version) = headers
|
||||
.get(&VERSION_HEADER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
{
|
||||
self.track_read_version(version);
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_query(
|
||||
&self,
|
||||
query: &AnyQuery,
|
||||
@@ -970,7 +928,6 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
|
||||
let futures = requests.into_iter().map(|req| async move {
|
||||
let (request_id, response) = self.send(req, true).await?;
|
||||
self.track_read_version_from_headers(response.headers());
|
||||
self.read_arrow_stream(&request_id, response).await
|
||||
});
|
||||
let streams = futures::future::try_join_all(futures);
|
||||
@@ -1588,12 +1545,11 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
*write_guard = None;
|
||||
drop(write_guard);
|
||||
|
||||
// Drop any per-handle read/write tracking; subsequent reads use the
|
||||
// Drop any per-handle write tracking; subsequent reads use the
|
||||
// baseline timestamp captured now to guarantee freshness.
|
||||
*self.freshness.lock().unwrap() = FreshnessState {
|
||||
min_version: None,
|
||||
checkout_baseline: Some(SystemTime::now()),
|
||||
min_read_version: None,
|
||||
};
|
||||
|
||||
// Invalidate schema cache since we're switching versions
|
||||
@@ -1849,7 +1805,6 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
}
|
||||
};
|
||||
|
||||
self.track_read_version_from_headers(response.headers());
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
|
||||
serde_json::from_str(&body).map_err(|e| Error::Http {
|
||||
@@ -2270,7 +2225,8 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
})
|
||||
}
|
||||
|
||||
async fn set_lsm_write_spec(&self, spec: LsmWriteSpec) -> Result<()> {
|
||||
async fn set_lsm_write_spec(&self, spec: crate::table::LsmWriteSpec) -> Result<()> {
|
||||
use crate::table::LsmWriteSpec;
|
||||
self.check_mutable().await?;
|
||||
|
||||
// Map the spec onto the server's request DTO. `sharding` is internally
|
||||
@@ -2322,69 +2278,6 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
|
||||
// Read counterpart to set/unset, resolved server-side against HEAD. The
|
||||
// server reads the spec from the `__lance_mem_wal` system index (shard
|
||||
// column mapped from its Lance field id against the current schema) and
|
||||
// re-encodes it into the same sophon-owned shape the set endpoint
|
||||
// accepts — no lance/lancedb types cross the wire. `lsm_write_spec` is
|
||||
// null when the LSM write path is not enabled for the table.
|
||||
let request = self.post_read(&format!(
|
||||
"/v1/table/{}/get_lsm_write_spec/",
|
||||
self.identifier
|
||||
));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
let body = response.text().await.err_to_http(request_id.clone())?;
|
||||
|
||||
// Mirror of sophon's `Sharding` (internally tagged on `mode`) and
|
||||
// `LsmWriteSpecBody` / `GetLsmWriteSpecResponse`.
|
||||
#[derive(Deserialize)]
|
||||
#[serde(tag = "mode", rename_all = "snake_case")]
|
||||
enum Sharding {
|
||||
Unsharded,
|
||||
Bucket { column: String, num_buckets: u32 },
|
||||
Identity { column: String },
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct LsmWriteSpecBody {
|
||||
sharding: Sharding,
|
||||
#[serde(default)]
|
||||
maintained_indexes: Vec<String>,
|
||||
#[serde(default)]
|
||||
writer_config_defaults: std::collections::HashMap<String, String>,
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct GetLsmWriteSpecResponse {
|
||||
lsm_write_spec: Option<LsmWriteSpecBody>,
|
||||
}
|
||||
|
||||
let parsed: GetLsmWriteSpecResponse =
|
||||
serde_json::from_str(&body).map_err(|e| Error::Http {
|
||||
source: format!("Failed to parse get_lsm_write_spec response: {}", e).into(),
|
||||
request_id,
|
||||
status_code: None,
|
||||
})?;
|
||||
|
||||
let Some(body) = parsed.lsm_write_spec else {
|
||||
// The LSM write path is not enabled for this table.
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let spec = match body.sharding {
|
||||
Sharding::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
} => LsmWriteSpec::bucket(column, num_buckets),
|
||||
Sharding::Identity { column } => LsmWriteSpec::identity(column),
|
||||
Sharding::Unsharded => LsmWriteSpec::unsharded(),
|
||||
}
|
||||
.with_maintained_indexes(body.maintained_indexes)
|
||||
.with_writer_config_defaults(body.writer_config_defaults);
|
||||
|
||||
Ok(Some(spec))
|
||||
}
|
||||
|
||||
async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
|
||||
Ok(Box::new(RemoteTags { inner: self }))
|
||||
}
|
||||
@@ -2416,6 +2309,126 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
message: "optimize is not supported on LanceDB cloud.".into(),
|
||||
})
|
||||
}
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[(String, String)],
|
||||
expression: &str,
|
||||
) -> Result<()> {
|
||||
let new_columns: Vec<serde_json::Value> = columns
|
||||
.iter()
|
||||
.map(|(name, data_type)| {
|
||||
serde_json::json!({
|
||||
"name": name,
|
||||
"computed": { "data_type": data_type, "expression": expression },
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/add_columns/", self.identifier))
|
||||
.json(&serde_json::json!({ "new_columns": new_columns }));
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
self.check_table_response(&request_id, response).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn refresh_column(
|
||||
&self,
|
||||
columns: &[String],
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
let mut body = serde_json::json!({ "columns": columns });
|
||||
if let Some(w) = where_clause {
|
||||
body["where_clause"] = serde_json::Value::String(w);
|
||||
}
|
||||
if let Some(n) = num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(p) = priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/refresh_column", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct RefreshColumnResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: RefreshColumnResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn load_columns(&self, request: crate::table::LoadColumnsRequest) -> Result<String> {
|
||||
let columns: Vec<serde_json::Value> = request
|
||||
.columns
|
||||
.iter()
|
||||
.map(|(target, source)| {
|
||||
serde_json::json!({
|
||||
"target": target,
|
||||
"source": source.clone().unwrap_or_else(|| target.clone()),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
let mut source = serde_json::json!({
|
||||
"uris": request.source_uris,
|
||||
"format": request.source_format,
|
||||
});
|
||||
if let Some(opts) = request.source_storage_options {
|
||||
source["storage_options"] = serde_json::to_value(opts).unwrap_or_default();
|
||||
}
|
||||
let mut body = serde_json::json!({
|
||||
"columns": columns,
|
||||
"source": source,
|
||||
"target_key": request.target_key,
|
||||
});
|
||||
if let Some(k) = request.source_key {
|
||||
body["source_key"] = serde_json::Value::String(k);
|
||||
}
|
||||
if let Some(m) = request.on_missing {
|
||||
body["on_missing"] = serde_json::Value::String(m);
|
||||
}
|
||||
if let Some(n) = request.num_workers {
|
||||
body["num_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.max_workers {
|
||||
body["max_workers"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.batch_size {
|
||||
body["batch_size"] = n.into();
|
||||
}
|
||||
if let Some(n) = request.commit_granularity {
|
||||
body["commit_granularity"] = n.into();
|
||||
}
|
||||
if let Some(p) = request.priority {
|
||||
body["priority"] = serde_json::Value::String(p);
|
||||
}
|
||||
let http_request = self
|
||||
.client
|
||||
.post(&format!("/v1/table/{}/load_columns", self.identifier))
|
||||
.json(&body);
|
||||
let (request_id, response) = self.send(http_request, true).await?;
|
||||
let response = self.check_table_response(&request_id, response).await?;
|
||||
#[derive(serde::Deserialize)]
|
||||
struct LoadColumnsResponse {
|
||||
job_id: String,
|
||||
}
|
||||
let body: LoadColumnsResponse = response.json().await.err_to_http(request_id)?;
|
||||
Ok(body.job_id)
|
||||
}
|
||||
|
||||
async fn add_columns(
|
||||
&self,
|
||||
transforms: NewColumnTransform,
|
||||
@@ -2908,6 +2921,75 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_refresh_column() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/refresh_column");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(body["columns"], serde_json::json!(["vec"]));
|
||||
assert_eq!(body["num_workers"], 2);
|
||||
assert!(body.get("where_clause").is_none());
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"j-9"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let job_id = table
|
||||
.refresh_column(&["vec".to_string()], None, Some(2), None, None, None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(job_id, "j-9");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_load_columns() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(request.url().path(), "/v1/table/my_table/load_columns");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap();
|
||||
assert_eq!(
|
||||
body["columns"],
|
||||
serde_json::json!([{"target": "embedding", "source": "emb"}])
|
||||
);
|
||||
assert_eq!(body["source"]["format"], "parquet");
|
||||
assert_eq!(
|
||||
body["source"]["uris"],
|
||||
serde_json::json!(["s3://b/x.parquet"])
|
||||
);
|
||||
assert_eq!(body["target_key"], "document_id");
|
||||
assert_eq!(body["source_key"], "doc_id");
|
||||
assert_eq!(body["on_missing"], "null");
|
||||
assert_eq!(body["num_workers"], 4);
|
||||
|
||||
http::Response::builder()
|
||||
.status(202)
|
||||
.body(r#"{"job_id":"lc-7"}"#)
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let request = crate::table::LoadColumnsRequest {
|
||||
source_uris: vec!["s3://b/x.parquet".to_string()],
|
||||
source_format: "parquet".to_string(),
|
||||
source_storage_options: None,
|
||||
target_key: "document_id".to_string(),
|
||||
source_key: Some("doc_id".to_string()),
|
||||
columns: vec![("embedding".to_string(), Some("emb".to_string()))],
|
||||
on_missing: Some("null".to_string()),
|
||||
num_workers: Some(4),
|
||||
max_workers: None,
|
||||
batch_size: None,
|
||||
commit_granularity: None,
|
||||
priority: None,
|
||||
};
|
||||
let job_id = table.load_columns(request).await.unwrap();
|
||||
assert_eq!(job_id, "lc-7");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_version() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
@@ -5424,74 +5506,6 @@ mod tests {
|
||||
table.unset_lsm_write_spec().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_lsm_write_spec() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(request.method(), "POST");
|
||||
assert_eq!(
|
||||
request.url().path(),
|
||||
"/v1/table/my_table/get_lsm_write_spec/"
|
||||
);
|
||||
|
||||
// The server resolves the spec and re-encodes it into the same
|
||||
// sophon-owned shape the set endpoint accepts (`Sharding` internally
|
||||
// tagged on `mode`, wrapped in `lsm_write_spec`).
|
||||
let response = serde_json::json!({
|
||||
"lsm_write_spec": {
|
||||
"sharding": { "mode": "bucket", "column": "id", "num_buckets": 4 },
|
||||
"maintained_indexes": ["id_idx"],
|
||||
"writer_config_defaults": { "durable_write": "false" },
|
||||
}
|
||||
});
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(response.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
|
||||
let spec = table
|
||||
.get_lsm_write_spec()
|
||||
.await
|
||||
.unwrap()
|
||||
.expect("a spec should be reported");
|
||||
match spec {
|
||||
crate::table::LsmWriteSpec::Bucket {
|
||||
column,
|
||||
num_buckets,
|
||||
maintained_indexes,
|
||||
writer_config_defaults,
|
||||
} => {
|
||||
assert_eq!(column, "id");
|
||||
assert_eq!(num_buckets, 4);
|
||||
assert_eq!(maintained_indexes, vec!["id_idx".to_string()]);
|
||||
assert_eq!(
|
||||
writer_config_defaults
|
||||
.get("durable_write")
|
||||
.map(String::as_str),
|
||||
Some("false")
|
||||
);
|
||||
}
|
||||
other => panic!("expected a bucket spec, got {:?}", other),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_lsm_write_spec_absent() {
|
||||
let table = Table::new_with_handler("my_table", |request| {
|
||||
assert_eq!(
|
||||
request.url().path(),
|
||||
"/v1/table/my_table/get_lsm_write_spec/"
|
||||
);
|
||||
// Null spec → the LSM write path is not enabled.
|
||||
let response = serde_json::json!({ "lsm_write_spec": null });
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(response.to_string())
|
||||
.unwrap()
|
||||
});
|
||||
assert!(table.get_lsm_write_spec().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_wait_for_index() {
|
||||
let table = _make_table_with_indices(0);
|
||||
@@ -7299,7 +7313,6 @@ mod tests {
|
||||
let state = FreshnessState {
|
||||
min_version: None,
|
||||
checkout_baseline: Some(baseline),
|
||||
min_read_version: None,
|
||||
};
|
||||
assert_eq!(compute_min_timestamp(&state, None, now), Some(baseline));
|
||||
|
||||
@@ -7324,7 +7337,6 @@ mod tests {
|
||||
let state = FreshnessState {
|
||||
min_version: None,
|
||||
checkout_baseline: Some(baseline),
|
||||
min_read_version: None,
|
||||
};
|
||||
assert_eq!(
|
||||
compute_min_timestamp(&state, Some(Duration::from_secs(10)), now),
|
||||
@@ -7336,7 +7348,6 @@ mod tests {
|
||||
let state = FreshnessState {
|
||||
min_version: None,
|
||||
checkout_baseline: Some(recent_baseline),
|
||||
min_read_version: None,
|
||||
};
|
||||
assert_eq!(
|
||||
compute_min_timestamp(&state, Some(Duration::from_secs(60)), now),
|
||||
@@ -7481,106 +7492,6 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
/// A handler that records every request's headers and answers each read with
|
||||
/// an `x-lancedb-version` response header taken from `versions` (by call
|
||||
/// index, saturating at the last entry). An empty string means "no header".
|
||||
fn read_version_handler(
|
||||
versions: &'static [&'static str],
|
||||
) -> (
|
||||
impl Fn(reqwest::Request) -> http::Response<String> + Clone + Send + Sync + 'static,
|
||||
Arc<std::sync::Mutex<Vec<http::HeaderMap>>>,
|
||||
) {
|
||||
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
|
||||
let requests_c = requests.clone();
|
||||
let call = Arc::new(AtomicUsize::new(0));
|
||||
let handler = move |request: reqwest::Request| {
|
||||
requests_c.lock().unwrap().push(request.headers().clone());
|
||||
let i = call.fetch_add(1, Ordering::SeqCst).min(versions.len() - 1);
|
||||
let mut builder = http::Response::builder().status(200);
|
||||
if !versions[i].is_empty() {
|
||||
builder = builder.header("x-lancedb-version", versions[i]);
|
||||
}
|
||||
builder.body("42".to_string()).unwrap()
|
||||
};
|
||||
(handler, requests)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_version_watermark_tracked_and_sent() {
|
||||
let (handler, requests) = read_version_handler(&["100", "100"]);
|
||||
let table = Table::new_with_handler("my_table", handler);
|
||||
|
||||
// First read has no watermark yet; the response advertises version 100,
|
||||
// so the second read must floor the server at 100.
|
||||
table.count_rows(None).await.unwrap();
|
||||
table.count_rows(None).await.unwrap();
|
||||
|
||||
let reqs = requests.lock().unwrap();
|
||||
assert!(!reqs[0].contains_key("x-lancedb-min-read-version"));
|
||||
assert_eq!(
|
||||
reqs[1]
|
||||
.get("x-lancedb-min-read-version")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"100"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_version_watermark_keeps_max() {
|
||||
// Server reports 100 then a stale 50; the watermark must not regress.
|
||||
let (handler, requests) = read_version_handler(&["100", "50", "50"]);
|
||||
let table = Table::new_with_handler("my_table", handler);
|
||||
|
||||
table.count_rows(None).await.unwrap();
|
||||
table.count_rows(None).await.unwrap();
|
||||
table.count_rows(None).await.unwrap();
|
||||
|
||||
let reqs = requests.lock().unwrap();
|
||||
assert_eq!(
|
||||
reqs[2]
|
||||
.get("x-lancedb-min-read-version")
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap(),
|
||||
"100"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_version_absent_header_no_watermark() {
|
||||
// An old server that doesn't return the version header leaves the
|
||||
// watermark unset, preserving backward compatibility.
|
||||
let (handler, requests) = read_version_handler(&[""]);
|
||||
let table = Table::new_with_handler("my_table", handler);
|
||||
|
||||
table.count_rows(None).await.unwrap();
|
||||
table.count_rows(None).await.unwrap();
|
||||
|
||||
let reqs = requests.lock().unwrap();
|
||||
assert!(!reqs[1].contains_key("x-lancedb-min-read-version"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_version_watermark_reset_on_checkout_latest() {
|
||||
let (handler, requests) = read_version_handler(&["100", "100"]);
|
||||
let table = Table::new_with_handler("my_table", handler);
|
||||
|
||||
table.count_rows(None).await.unwrap();
|
||||
table.checkout_latest().await.unwrap();
|
||||
table.count_rows(None).await.unwrap();
|
||||
|
||||
// The read after checkout_latest starts from a clean slate.
|
||||
let reqs = requests.lock().unwrap();
|
||||
assert!(
|
||||
!reqs
|
||||
.last()
|
||||
.unwrap()
|
||||
.contains_key("x-lancedb-min-read-version")
|
||||
);
|
||||
}
|
||||
|
||||
/// Like `capturing_handler`, but keeps a per-path snapshot of the headers
|
||||
/// from every request so tests can assert on a specific endpoint.
|
||||
#[allow(clippy::type_complexity)]
|
||||
|
||||
+110
-98
@@ -471,6 +471,33 @@ impl LsmWriteSpec {
|
||||
}
|
||||
}
|
||||
|
||||
/// Request to fill existing table columns from an external source by
|
||||
/// primary-key join (Geneva `Table.load_columns()` parity). Server-backed
|
||||
/// feature (LanceDB Enterprise / Cloud).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoadColumnsRequest {
|
||||
/// External source URIs.
|
||||
pub source_uris: Vec<String>,
|
||||
/// Source format: "parquet" | "lance" | "ipc".
|
||||
pub source_format: String,
|
||||
/// Source-only storage options (e.g. cloud credentials).
|
||||
pub source_storage_options: Option<HashMap<String, String>>,
|
||||
/// Destination primary-key column.
|
||||
pub target_key: String,
|
||||
/// Source primary-key column. Defaults to `target_key` when None.
|
||||
pub source_key: Option<String>,
|
||||
/// Value column mappings as `(target, source)`; a None source defaults to
|
||||
/// the target name.
|
||||
pub columns: Vec<(String, Option<String>)>,
|
||||
/// Missing-row policy: "carry" (default) | "null" | "error".
|
||||
pub on_missing: Option<String>,
|
||||
pub num_workers: Option<u32>,
|
||||
pub max_workers: Option<u32>,
|
||||
pub batch_size: Option<u32>,
|
||||
pub commit_granularity: Option<u32>,
|
||||
pub priority: Option<String>,
|
||||
}
|
||||
|
||||
/// A trait for anything "table-like". This is used for both native tables (which target
|
||||
/// Lance datasets) and remote tables (which target LanceDB cloud)
|
||||
///
|
||||
@@ -581,16 +608,6 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
message: "unset_lsm_write_spec is not supported on this table type".into(),
|
||||
})
|
||||
}
|
||||
/// Read the [`LsmWriteSpec`] currently installed on this table, returning
|
||||
/// `None` when the MemWAL LSM write path is not enabled.
|
||||
///
|
||||
/// The default implementation returns `NotSupported`. Implementations that
|
||||
/// support the MemWAL LSM write path must override this.
|
||||
async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
|
||||
Err(Error::NotSupported {
|
||||
message: "get_lsm_write_spec is not supported on this table type".into(),
|
||||
})
|
||||
}
|
||||
/// Drain and close any cached MemWAL shard writers for this table.
|
||||
///
|
||||
/// The default implementation is a no-op; table types that maintain
|
||||
@@ -630,6 +647,47 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
transforms: NewColumnTransform,
|
||||
read_columns: Option<Vec<String>>,
|
||||
) -> Result<AddColumnsResult>;
|
||||
/// Declare computed columns bound to a registered function: each
|
||||
/// `(name, sql_type)` is added all-null with the expression stored
|
||||
/// as its binding; no compute happens here (the server's lazy
|
||||
/// detector or refresh_column fills them). Several columns map a
|
||||
/// struct-returning function's fields positionally. Server-backed
|
||||
/// feature; the default returns NotSupported.
|
||||
async fn add_computed_columns(
|
||||
&self,
|
||||
_columns: &[(String, String)],
|
||||
_expression: &str,
|
||||
) -> Result<()> {
|
||||
Err(Error::NotSupported {
|
||||
message: "computed columns are not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Trigger recompute of computed columns. The expression is
|
||||
/// resolved server-side from each column's stored binding; columns
|
||||
/// bound to the same struct-returning function refresh together.
|
||||
/// Returns the refresh job id. Server-backed feature (LanceDB
|
||||
/// Enterprise / Cloud); the default returns NotSupported.
|
||||
async fn refresh_column(
|
||||
&self,
|
||||
_columns: &[String],
|
||||
_where_clause: Option<String>,
|
||||
_num_workers: Option<u32>,
|
||||
_max_workers: Option<u32>,
|
||||
_batch_size: Option<u32>,
|
||||
_priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
Err(Error::NotSupported {
|
||||
message: "refresh_column is not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Fill existing columns from an external source by primary-key join
|
||||
/// (Geneva `load_columns`). Returns the load job id. Server-backed feature;
|
||||
/// the default returns NotSupported.
|
||||
async fn load_columns(&self, _request: LoadColumnsRequest) -> Result<String> {
|
||||
Err(Error::NotSupported {
|
||||
message: "load_columns is not supported by this table".into(),
|
||||
})
|
||||
}
|
||||
/// Alter columns in the table.
|
||||
async fn alter_columns(&self, alterations: &[ColumnAlteration]) -> Result<AlterColumnsResult>;
|
||||
/// Drop columns from the table.
|
||||
@@ -1471,6 +1529,48 @@ impl Table {
|
||||
self.inner.add_columns(transforms, read_columns).await
|
||||
}
|
||||
|
||||
/// Declare computed columns bound to a registered function
|
||||
/// (`(name, sql_type)` pairs + a `f(args)` expression). No compute
|
||||
/// happens here. Server-backed feature.
|
||||
pub async fn add_computed_columns(
|
||||
&self,
|
||||
columns: &[(String, String)],
|
||||
expression: &str,
|
||||
) -> Result<()> {
|
||||
self.inner.add_computed_columns(columns, expression).await
|
||||
}
|
||||
|
||||
/// Trigger recompute of computed columns (REFRESH COLUMN). The
|
||||
/// expression comes from each column's stored binding; columns
|
||||
/// bound to the same struct-returning function refresh together.
|
||||
/// Returns the refresh job id. Server-backed feature.
|
||||
pub async fn refresh_column(
|
||||
&self,
|
||||
columns: &[String],
|
||||
where_clause: Option<String>,
|
||||
num_workers: Option<u32>,
|
||||
max_workers: Option<u32>,
|
||||
batch_size: Option<u32>,
|
||||
priority: Option<String>,
|
||||
) -> Result<String> {
|
||||
self.inner
|
||||
.refresh_column(
|
||||
columns,
|
||||
where_clause,
|
||||
num_workers,
|
||||
max_workers,
|
||||
batch_size,
|
||||
priority,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Fill existing columns from an external Parquet/Lance/IPC source by
|
||||
/// primary-key join (Geneva `Table.load_columns()`). Returns the job id.
|
||||
pub async fn load_columns(&self, request: LoadColumnsRequest) -> Result<String> {
|
||||
self.inner.load_columns(request).await
|
||||
}
|
||||
|
||||
/// Change a column's name or nullability.
|
||||
pub async fn alter_columns(
|
||||
&self,
|
||||
@@ -1548,29 +1648,6 @@ impl Table {
|
||||
self.inner.unset_lsm_write_spec().await
|
||||
}
|
||||
|
||||
/// Read the [`LsmWriteSpec`] currently installed on this table.
|
||||
///
|
||||
/// Returns `Ok(None)` when the MemWAL LSM write path is not enabled (no
|
||||
/// spec has been set, or it was removed with [`Table::unset_lsm_write_spec`]).
|
||||
/// The returned spec — including its [`LsmWriteSpec::maintained_indexes`] and
|
||||
/// [`LsmWriteSpec::writer_config_defaults`] — mirrors what was passed to
|
||||
/// [`Table::set_lsm_write_spec`].
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```
|
||||
/// # use lancedb::table::Table;
|
||||
/// # async fn example(table: &Table) -> Result<(), Box<dyn std::error::Error>> {
|
||||
/// if let Some(spec) = table.get_lsm_write_spec().await? {
|
||||
/// println!("LSM write path enabled: {:?}", spec);
|
||||
/// }
|
||||
/// # Ok(())
|
||||
/// # }
|
||||
/// ```
|
||||
pub async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
|
||||
self.inner.get_lsm_write_spec().await
|
||||
}
|
||||
|
||||
/// Drain and close any cached MemWAL shard writers held for this table.
|
||||
///
|
||||
/// When an [`LsmWriteSpec`] is installed, `merge_insert` opens MemWAL shard
|
||||
@@ -2883,10 +2960,6 @@ impl BaseTable for NativeTable {
|
||||
merge::lsm::unset_lsm_write_spec(self).await
|
||||
}
|
||||
|
||||
async fn get_lsm_write_spec(&self) -> Result<Option<LsmWriteSpec>> {
|
||||
merge::lsm::get_lsm_write_spec(self).await
|
||||
}
|
||||
|
||||
async fn close_lsm_writers(&self) -> Result<()> {
|
||||
merge::lsm::close_lsm_writers(self).await
|
||||
}
|
||||
@@ -4448,67 +4521,6 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_lsm_write_spec() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
let uri = tmp_dir.path().to_str().unwrap();
|
||||
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("id", DataType::Int64, false),
|
||||
Field::new("region", DataType::Utf8, true),
|
||||
]));
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(arrow_array::Int64Array::from(vec![1, 2, 3])),
|
||||
Arc::new(StringArray::from(vec!["a", "b", "c"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let reader: Box<dyn arrow_array::RecordBatchReader + Send> =
|
||||
Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema.clone()));
|
||||
let conn = ConnectBuilder::new(uri)
|
||||
.read_consistency_interval(Duration::from_secs(0))
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let table = conn.create_table("t", reader).execute().await.unwrap();
|
||||
|
||||
// No spec installed yet.
|
||||
assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
|
||||
|
||||
// A real scalar index is needed to name it as a maintained index.
|
||||
table
|
||||
.create_index(&["id"], Index::Auto)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap();
|
||||
let idx_name = table.list_indices().await.unwrap()[0].name.clone();
|
||||
|
||||
// Bucket spec round-trips exactly, including the routing column (recovered
|
||||
// from its field id), maintained indexes, and writer config defaults.
|
||||
let spec = LsmWriteSpec::bucket("id", 4)
|
||||
.with_maintained_indexes([idx_name])
|
||||
.with_writer_config_defaults([("durable_write", "false")]);
|
||||
table.set_lsm_write_spec(spec.clone()).await.unwrap();
|
||||
assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
|
||||
|
||||
// After unset, no spec is reported.
|
||||
table.unset_lsm_write_spec().await.unwrap();
|
||||
assert_eq!(table.get_lsm_write_spec().await.unwrap(), None);
|
||||
|
||||
// Identity sharding round-trips (column recovered from the schema).
|
||||
let spec = LsmWriteSpec::identity("region");
|
||||
table.set_lsm_write_spec(spec.clone()).await.unwrap();
|
||||
assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
|
||||
table.unset_lsm_write_spec().await.unwrap();
|
||||
|
||||
// Unsharded round-trips (no routing column).
|
||||
let spec = LsmWriteSpec::unsharded();
|
||||
table.set_lsm_write_spec(spec.clone()).await.unwrap();
|
||||
assert_eq!(table.get_lsm_write_spec().await.unwrap(), Some(spec));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
pub async fn test_stats() {
|
||||
let tmp_dir = tempdir().unwrap();
|
||||
|
||||
@@ -32,7 +32,7 @@ use lance::dataset::mem_wal::{
|
||||
};
|
||||
use lance::index::DatasetIndexExt;
|
||||
use lance_core::datatypes::Schema as LanceSchema;
|
||||
use lance_index::mem_wal::{MemWalIndexDetails, ShardingField, ShardingSpec};
|
||||
use lance_index::mem_wal::{MemWalIndexDetails, ShardingSpec};
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -148,97 +148,6 @@ pub(crate) async fn unset_lsm_write_spec(table: &NativeTable) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// get_lsm_write_spec
|
||||
// =============================================================================
|
||||
|
||||
/// Read the [`LsmWriteSpec`] currently installed on the table, if any.
|
||||
///
|
||||
/// Reconstructs the spec from the MemWAL index's stored details — the same
|
||||
/// metadata [`set_lsm_write_spec`] writes — read directly via
|
||||
/// [`mem_wal_index_details`](DatasetMemWalExt::mem_wal_index_details) (the raw
|
||||
/// manifest record, not the `describe_indices` enrichment path, so it is
|
||||
/// unaffected by system-index filtering there). Returns `Ok(None)` when no
|
||||
/// MemWAL index is installed.
|
||||
#[allow(clippy::redundant_pub_crate)]
|
||||
pub(crate) async fn get_lsm_write_spec(table: &NativeTable) -> Result<Option<LsmWriteSpec>> {
|
||||
let dataset = table.dataset.get().await?;
|
||||
let Some(details) = dataset.mem_wal_index_details().await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let spec = lsm_write_spec_from_details(&details, dataset.schema())?;
|
||||
Ok(Some(spec))
|
||||
}
|
||||
|
||||
/// Reconstruct the public [`LsmWriteSpec`] from stored MemWAL index details.
|
||||
///
|
||||
/// The sharding transform and its parameters recover the mode; the routing
|
||||
/// column is resolved from the sharding field's source id back to its schema
|
||||
/// name (only bucket / identity carry a column). `maintained_indexes` and
|
||||
/// `writer_config_defaults` are copied verbatim, so the result round-trips
|
||||
/// what `set_lsm_write_spec` installed.
|
||||
fn lsm_write_spec_from_details(
|
||||
details: &MemWalIndexDetails,
|
||||
schema: &LanceSchema,
|
||||
) -> Result<LsmWriteSpec> {
|
||||
let spec = details
|
||||
.sharding_specs
|
||||
.first()
|
||||
.ok_or_else(|| Error::Runtime {
|
||||
message: "get_lsm_write_spec: MemWAL index has no sharding spec".to_string(),
|
||||
})?;
|
||||
let field = spec.fields.first().ok_or_else(|| Error::Runtime {
|
||||
message: "get_lsm_write_spec: MemWAL index has an empty sharding spec".to_string(),
|
||||
})?;
|
||||
|
||||
let base = match field.transform.as_deref() {
|
||||
Some(BUCKET_TRANSFORM) => {
|
||||
let num_buckets = field
|
||||
.parameters
|
||||
.get(NUM_BUCKETS_PARAM)
|
||||
.and_then(|s| s.parse::<u32>().ok())
|
||||
.filter(|n| *n > 0)
|
||||
.ok_or_else(|| Error::Runtime {
|
||||
message: "get_lsm_write_spec: MemWAL bucket spec has a missing or invalid num_buckets parameter".to_string(),
|
||||
})?;
|
||||
LsmWriteSpec::bucket(sharding_column(field, schema)?, num_buckets)
|
||||
}
|
||||
Some(IDENTITY_TRANSFORM) => LsmWriteSpec::identity(sharding_column(field, schema)?),
|
||||
Some(UNSHARDED_TRANSFORM) => LsmWriteSpec::unsharded(),
|
||||
other => {
|
||||
return Err(Error::Runtime {
|
||||
message: format!(
|
||||
"get_lsm_write_spec: MemWAL index has an unsupported sharding transform {:?}",
|
||||
other
|
||||
),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(base
|
||||
.with_maintained_indexes(details.maintained_indexes.clone())
|
||||
.with_writer_config_defaults(details.writer_config_defaults.clone()))
|
||||
}
|
||||
|
||||
/// Resolve the single routing column name from a sharding field's source id.
|
||||
///
|
||||
/// `set_lsm_write_spec` records the shard column by its Lance field id, so the
|
||||
/// name is recovered by looking that id up in the current schema.
|
||||
fn sharding_column(field: &ShardingField, schema: &LanceSchema) -> Result<String> {
|
||||
let source_id = *field.source_ids.first().ok_or_else(|| Error::Runtime {
|
||||
message: "get_lsm_write_spec: sharding field has no source column".to_string(),
|
||||
})?;
|
||||
schema
|
||||
.field_by_id(source_id)
|
||||
.map(|f| f.name.clone())
|
||||
.ok_or_else(|| Error::Runtime {
|
||||
message: format!(
|
||||
"get_lsm_write_spec: sharding source field id {} not found in schema",
|
||||
source_id
|
||||
),
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// close_lsm_writers
|
||||
// =============================================================================
|
||||
|
||||
@@ -35,15 +35,6 @@ pub enum AnyQuery {
|
||||
VectorQuery(VectorQueryRequest),
|
||||
}
|
||||
|
||||
impl AnyQuery {
|
||||
pub(crate) fn base(&self) -> &QueryRequest {
|
||||
match self {
|
||||
Self::Query(query) => query,
|
||||
Self::VectorQuery(query) => &query.base,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//Decide between namespace or local
|
||||
pub async fn execute_query(
|
||||
table: &NativeTable,
|
||||
@@ -117,7 +108,6 @@ pub async fn create_plan(
|
||||
AnyQuery::VectorQuery(query) => query.clone(),
|
||||
AnyQuery::Query(query) => VectorQueryRequest::from_plain_query(query.clone()),
|
||||
};
|
||||
query.base.check_filter()?;
|
||||
|
||||
let ds_ref = table.dataset.get().await?;
|
||||
let schema = ds_ref.schema();
|
||||
@@ -367,7 +357,6 @@ async fn execute_namespace_query(
|
||||
|
||||
/// Convert an AnyQuery to the namespace QueryTableRequest format.
|
||||
fn convert_to_namespace_query(query: &AnyQuery) -> Result<NsQueryTableRequest> {
|
||||
query.base().check_filter()?;
|
||||
match query {
|
||||
AnyQuery::VectorQuery(vq) => {
|
||||
// Extract the query vector(s)
|
||||
|
||||
Reference in New Issue
Block a user