Compare commits

..

3 Commits

Author SHA1 Message Date
Will Jones 5406e3c7c2 feat(listing): soft-delete root tables via an embedded V1 namespace
Root-table lifecycle in a ListingDatabase now flows through an embedded V1
(manifest-disabled) DirectoryNamespace, so a root drop is a soft-delete rather than
an immediate remove_dir_all:

- drop_table writes a delete marker and leaves the data for a later purge (TTL).
- create_table on a soft-deleted name revives it (clears the marker under the
  namespace's lifecycle lock, then overwrites via the native create path).
- open_table / table_names / list_tables treat soft-deleted tables as absent.
- table listing is now a single O(1) read_dir in the namespace instead of a
  per-table probe here.

A ListingDatabase now holds two embedded namespaces: the existing manifest-backed
`namespace_database` for child namespaces (multi-level table ids), and a new
`root_namespace_database` (V1, manifest-off) that owns root soft-delete/purge/
table_status. `namespace_client()` still returns the manifest namespace so child
namespace ops are unaffected.

Also preserves TableNotFound through LanceNamespaceDatabase::drop_table instead of
flattening it to a generic Runtime error, so dropping a missing table reports the
right error.

Removes the now-dead native drop_tables / commit-handler path and the
object_store/base_path fields it needed.

Depends on lance-format/lance#7541.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-01 14:48:34 -07:00
Will Jones 99a1c04d2a chore: temporarily pin lance to soft-delete PR rev
Points the lance workspace deps at lance-format/lance rev 4acefffd5 (PR
lance-format/lance#7541), which adds the soft-delete lifecycle trait methods this
change depends on. Revert to a release tag once a new lance beta is published.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-30 17:08:07 -07:00
lancedb automation 027f60a8b2 chore: update lance dependency to v9.0.0-beta.8 2026-06-24 23:29:23 +00:00
97 changed files with 751 additions and 6516 deletions
-137
View File
@@ -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
View File
@@ -1,5 +1,5 @@
[tool.bumpversion]
current_version = "0.31.0-beta.6"
current_version = "0.31.0-beta.2"
parse = """(?x)
(?P<major>0|[1-9]\\d*)\\.
(?P<minor>0|[1-9]\\d*)\\.
+2 -2
View File
@@ -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) }}
+3 -3
View File
@@ -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
+1 -1
View File
@@ -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
+2 -2
View File
@@ -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'
+2 -2
View File
@@ -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) }}
+1 -1
View File
@@ -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:
+1 -1
View File
@@ -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: |
+1 -1
View File
@@ -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
+10 -10
View File
@@ -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.
+19 -19
View File
@@ -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) }}
+7 -7
View File
@@ -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
+7 -7
View File
@@ -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
+71 -88
View File
@@ -157,9 +157,9 @@ dependencies = [
[[package]]
name = "anyhow"
version = "1.0.103"
version = "1.0.102"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3"
checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
[[package]]
name = "approx"
@@ -1297,6 +1297,15 @@ version = "2.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3"
[[package]]
name = "bitpacking"
version = "0.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "96a7139abd3d9cebf8cd6f920a389cf3dc9576172e32f4563f188cae3c3eb019"
dependencies = [
"crunchy",
]
[[package]]
name = "bitvec"
version = "1.0.1"
@@ -3177,9 +3186,9 @@ dependencies = [
[[package]]
name = "env_filter"
version = "2.0.0"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217"
checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef"
dependencies = [
"log",
"regex",
@@ -3187,9 +3196,9 @@ dependencies = [
[[package]]
name = "env_logger"
version = "0.11.11"
version = "0.11.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6"
checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a"
dependencies = [
"anstream",
"anstyle",
@@ -3423,8 +3432,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"rand 0.9.4",
@@ -4726,8 +4735,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arc-swap",
"arrow",
@@ -4745,6 +4754,7 @@ dependencies = [
"async_cell",
"aws-credential-types",
"aws-sdk-dynamodb",
"bitpacking",
"byteorder",
"bytes",
"chrono",
@@ -4763,7 +4773,6 @@ dependencies = [
"humantime",
"itertools 0.14.0",
"lance-arrow",
"lance-bitpacking",
"lance-core",
"lance-datafusion",
"lance-encoding",
@@ -4801,8 +4810,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4823,7 +4832,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.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4837,7 +4846,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.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -4846,19 +4855,18 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrayref",
"crunchy",
"paste",
"seq-macro",
]
[[package]]
name = "lance-core"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -4896,8 +4904,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"arrow-array",
@@ -4927,8 +4935,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"arrow-array",
@@ -4945,8 +4953,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"proc-macro2",
"quote",
@@ -4955,8 +4963,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -4991,8 +4999,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5022,8 +5030,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arc-swap",
"arrow",
@@ -5035,6 +5043,7 @@ dependencies = [
"async-channel",
"async-recursion",
"async-trait",
"bitpacking",
"bitvec",
"bytes",
"chrono",
@@ -5052,7 +5061,6 @@ dependencies = [
"jsonb",
"lance-arrow",
"lance-arrow-stats",
"lance-bitpacking",
"lance-core",
"lance-datafusion",
"lance-datagen",
@@ -5088,8 +5096,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"arrow-arith",
@@ -5130,8 +5138,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5147,8 +5155,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"async-trait",
@@ -5160,8 +5168,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"arrow-ipc",
@@ -5215,8 +5223,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5231,8 +5239,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"arrow",
"arrow-array",
@@ -5271,22 +5279,22 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "9.0.0-beta.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
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.16"
source = "git+https://github.com/lance-format/lance.git?tag=v9.0.0-beta.16#78a814b6b448b433f5ebfdb9a8dbdec7b904c59f"
version = "9.0.0-beta.8"
source = "git+https://github.com/lance-format/lance.git?rev=4acefffd5d38f88003fce681ae1d0871077ce5e7#4acefffd5d38f88003fce681ae1d0871077ce5e7"
dependencies = [
"icu_segmenter",
"jieba-rs",
@@ -5299,13 +5307,12 @@ dependencies = [
[[package]]
name = "lancedb"
version = "0.31.0-beta.6"
version = "0.31.0-beta.2"
dependencies = [
"ahash",
"anyhow",
"arrow",
"arrow-array",
"arrow-buffer",
"arrow-cast",
"arrow-data",
"arrow-ipc",
@@ -5361,7 +5368,6 @@ dependencies = [
"pin-project",
"polars",
"polars-arrow",
"pprof 0.14.1",
"rand 0.9.4",
"random_word",
"regex",
@@ -5378,14 +5384,13 @@ dependencies = [
"tokenizers",
"tokio",
"url",
"urlencoding",
"uuid",
"walkdir",
]
[[package]]
name = "lancedb-nodejs"
version = "0.31.0-beta.6"
version = "0.31.0-beta.2"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5410,7 +5415,7 @@ dependencies = [
[[package]]
name = "lancedb-python"
version = "0.34.0-beta.6"
version = "0.34.0-beta.2"
dependencies = [
"arrow",
"async-trait",
@@ -5643,9 +5648,9 @@ dependencies = [
[[package]]
name = "log"
version = "0.4.33"
version = "0.4.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad"
checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a"
[[package]]
name = "loom"
@@ -5953,9 +5958,9 @@ dependencies = [
[[package]]
name = "napi"
version = "3.10.3"
version = "3.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c71997d6f7ad4a756966e452426848ac27d3b37a295302d63afbbcce0270f93"
checksum = "fbd9f9295f3ff5921e78a71222c3361a8216f7760b1a99a6ad4e8441de18bbb9"
dependencies = [
"bitflags 2.11.1",
"chrono",
@@ -5978,9 +5983,9 @@ checksum = "c9c366d2c8c60b86fa632df75f745509b52f9128f91a6bad4c796e44abb505e1"
[[package]]
name = "napi-derive"
version = "3.5.9"
version = "3.5.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d4ba572deef53e2c386759a8c2014175a62679d74ff83adc205c8bc0e0285727"
checksum = "89b3f766e04667e6da0e181e2da4f85475d5a6513b7cf6a80bea184e224a5b42"
dependencies = [
"convert_case",
"ctor 1.0.5",
@@ -5992,9 +5997,9 @@ dependencies = [
[[package]]
name = "napi-derive-backend"
version = "5.1.1"
version = "5.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ddd961eb2aa8965e3f29722d754f3a86907eb1984e2fbcbe3fe87b9a02d6bfba"
checksum = "0d5af30503edf933ce7377cf6d4c877a62b0f1107ea05585f1b5e430e88d5baf"
dependencies = [
"convert_case",
"proc-macro2",
@@ -7327,28 +7332,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"
@@ -10144,9 +10127,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "uuid"
version = "1.23.4"
version = "1.23.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53"
checksum = "144d6b123cef80b301b8f72a9e2ca4370ddec21950d0a103dd22c437006d2db7"
dependencies = [
"getrandom 0.4.2",
"js-sys",
+14 -15
View File
@@ -13,25 +13,24 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=9.0.0-beta.16", default-features = false, "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=9.0.0-beta.16", default-features = false, "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=9.0.0-beta.16", default-features = false, "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=9.0.0-beta.16", "tag" = "v9.0.0-beta.16", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=9.0.0-beta.8", default-features = false, "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=9.0.0-beta.8", default-features = false, "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=9.0.0-beta.8", default-features = false, "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=9.0.0-beta.8", "rev" = "4acefffd5d38f88003fce681ae1d0871077ce5e7", "git" = "https://github.com/lance-format/lance.git" }
ahash = "0.8"
# Note that this one does not include pyarrow
arrow = { version = "58.0.0", optional = false }
arrow-array = "58.0.0"
arrow-buffer = "58.0.0"
arrow-data = "58.0.0"
arrow-ipc = "58.0.0"
arrow-ord = "58.0.0"
-12
View File
@@ -119,18 +119,6 @@ ignore = [
# 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" },
]
# ---------------------------------------------------------------------------
+1 -1
View File
@@ -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.2</version>
</dependency>
```
-3
View File
@@ -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
-3
View File
@@ -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
-29
View File
@@ -1,29 +0,0 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / OAuthFlowType
# Enumeration: OAuthFlowType
OAuth authentication flow types.
## Enumeration Members
### AzureManagedIdentity
```ts
AzureManagedIdentity: "azure_managed_identity";
```
Azure Managed Identity via IMDS.
***
### ClientCredentials
```ts
ClientCredentials: "client_credentials";
```
Client Credentials grant (service-to-service / M2M).
-3
View File
@@ -12,7 +12,6 @@
## Enumerations
- [FullTextQueryType](enumerations/FullTextQueryType.md)
- [OAuthFlowType](enumerations/OAuthFlowType.md)
- [Occur](enumerations/Occur.md)
- [Operator](enumerations/Operator.md)
@@ -86,8 +85,6 @@
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
- [MergeResult](interfaces/MergeResult.md)
- [NativeOAuthConfig](interfaces/NativeOAuthConfig.md)
- [OAuthConfig](interfaces/OAuthConfig.md)
- [OpenTableOptions](interfaces/OpenTableOptions.md)
- [OptimizeOptions](interfaces/OptimizeOptions.md)
- [OptimizeStats](interfaces/OptimizeStats.md)
@@ -64,19 +64,6 @@ client used by manifest-enabled native connections.
***
### oauthConfig?
```ts
optional oauthConfig: NativeOAuthConfig;
```
(For LanceDB cloud only): OAuth configuration for IdP-based
authentication (e.g., Azure Entra ID). When set, token acquisition
and refresh are handled entirely in Rust. TypeScript users should pass
the public `OAuthConfig` type exported from `@lancedb/lancedb`.
***
### readConsistencyInterval?
```ts
@@ -1,88 +0,0 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / NativeOAuthConfig
# Interface: NativeOAuthConfig
OAuth configuration for LanceDB authentication.
This is the generated napi-rs binding shape. TypeScript users should prefer
the public `OAuthConfig` type exported from `@lancedb/lancedb`.
All token acquisition and refresh is handled in the Rust layer.
## Properties
### clientId
```ts
clientId: string;
```
Application / Client ID.
***
### clientSecret?
```ts
optional clientSecret: string;
```
Client secret (required for client_credentials).
***
### flow?
```ts
optional flow: string;
```
Authentication flow: "client_credentials" or "azure_managed_identity"
***
### issuerUrl
```ts
issuerUrl: string;
```
OIDC issuer URL or OAuth authority URL.
For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
***
### managedIdentityClientId?
```ts
optional managedIdentityClientId: string;
```
Client ID for user-assigned managed identity (azure_managed_identity).
***
### refreshBufferSecs?
```ts
optional refreshBufferSecs: number;
```
Seconds before expiry to trigger proactive refresh (default: 300).
Keep this well below the token TTL; if it is greater than or equal to
the TTL, each request refreshes the token.
***
### scopes
```ts
scopes: string[];
```
OAuth scopes to request. For Azure managed identity, exactly one scope
or resource is required. For example: `["api://{app_id}/.default"]`
-111
View File
@@ -1,111 +0,0 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / OAuthConfig
# Interface: OAuthConfig
OAuth configuration for LanceDB authentication.
This is the public TypeScript OAuth configuration type. The generated
`NativeOAuthConfig` type has the same runtime shape but is an implementation
detail of the napi-rs binding.
All token acquisition and refresh is handled in the Rust layer.
This config is passed through to Rust via napi-rs.
## Examples
```typescript
const config: OAuthConfig = {
issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
clientId: "app-id",
clientSecret: "secret",
scopes: ["api://lancedb-api/.default"],
};
```
```typescript
const config: OAuthConfig = {
issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
clientId: "app-id",
scopes: ["api://lancedb-api/.default"],
flow: OAuthFlowType.AzureManagedIdentity,
};
```
## Properties
### clientId
```ts
clientId: string;
```
Application / Client ID.
***
### clientSecret?
```ts
optional clientSecret: string;
```
Client secret (required for ClientCredentials).
***
### flow?
```ts
optional flow: OAuthFlowType;
```
Authentication flow (default: ClientCredentials).
***
### issuerUrl
```ts
issuerUrl: string;
```
OIDC issuer URL or OAuth authority URL.
For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
***
### managedIdentityClientId?
```ts
optional managedIdentityClientId: string;
```
Client ID for user-assigned managed identity (AzureManagedIdentity).
***
### refreshBufferSecs?
```ts
optional refreshBufferSecs: number;
```
Seconds before expiry to trigger proactive refresh (default: 300).
Keep this well below the token TTL; if it is greater than or equal to
the TTL, each request refreshes the token.
***
### scopes
```ts
scopes: string[];
```
OAuth scopes to request.
For Azure managed identity, exactly one scope or resource is required.
For example: `["api://{app_id}/.default"]`
@@ -8,14 +8,6 @@
## Properties
### clumpSize?
```ts
optional clumpSize: number;
```
***
### counts?
```ts
+1 -1
View File
@@ -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.2</version>
<relativePath>../pom.xml</relativePath>
</parent>
+2 -2
View File
@@ -6,7 +6,7 @@
<groupId>com.lancedb</groupId>
<artifactId>lancedb-parent</artifactId>
<version>0.31.0-beta.6</version>
<version>0.31.0-beta.2</version>
<packaging>pom</packaging>
<name>${project.artifactId}</name>
<description>LanceDB Java SDK Parent POM</description>
@@ -28,7 +28,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
<arrow.version>15.0.0</arrow.version>
<lance-core.version>9.0.0-beta.16</lance-core.version>
<lance-core.version>9.0.0-beta.8</lance-core.version>
<spotless.skip>false</spotless.skip>
<spotless.version>2.30.0</spotless.version>
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
+1 -1
View File
@@ -1,7 +1,7 @@
[package]
name = "lancedb-nodejs"
edition.workspace = true
version = "0.31.0-beta.6"
version = "0.31.0-beta.2"
publish = false
license.workspace = true
description.workspace = true
-14
View File
@@ -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()
-3
View File
@@ -52,7 +52,6 @@ export {
SplitHashOptions,
SplitSequentialOptions,
ShuffleOptions,
OAuthConfig as NativeOAuthConfig,
} from "./native.js";
export {
@@ -131,8 +130,6 @@ export {
TokenResponse,
} from "./header";
export { OAuthConfig, OAuthFlowType } from "./oauth";
export { MergeInsertBuilder, WriteExecutionOptions } from "./merge";
export * as embedding from "./embedding";
-76
View File
@@ -1,76 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
/**
* OAuth authentication flow types.
*/
export enum OAuthFlowType {
/** Client Credentials grant (service-to-service / M2M). */
ClientCredentials = "client_credentials",
/** Azure Managed Identity via IMDS. */
AzureManagedIdentity = "azure_managed_identity",
}
/**
* OAuth configuration for LanceDB authentication.
*
* This is the public TypeScript OAuth configuration type. The generated
* `NativeOAuthConfig` type has the same runtime shape but is an implementation
* detail of the napi-rs binding.
*
* All token acquisition and refresh is handled in the Rust layer.
* This config is passed through to Rust via napi-rs.
*
* @example Client Credentials (service-to-service):
* ```typescript
* const config: OAuthConfig = {
* issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
* clientId: "app-id",
* clientSecret: "secret",
* scopes: ["api://lancedb-api/.default"],
* };
* ```
*
* @example Azure Managed Identity:
* ```typescript
* const config: OAuthConfig = {
* issuerUrl: "https://login.microsoftonline.com/{tenant}/v2.0",
* clientId: "app-id",
* scopes: ["api://lancedb-api/.default"],
* flow: OAuthFlowType.AzureManagedIdentity,
* };
* ```
*/
export interface OAuthConfig {
/**
* OIDC issuer URL or OAuth authority URL.
* For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
*/
issuerUrl: string;
/** Application / Client ID. */
clientId: string;
/**
* OAuth scopes to request.
* For Azure managed identity, exactly one scope or resource is required.
* For example: `["api://{app_id}/.default"]`
*/
scopes: string[];
/** Authentication flow (default: ClientCredentials). */
flow?: OAuthFlowType;
/** Client secret (required for ClientCredentials). */
clientSecret?: string;
/** Client ID for user-assigned managed identity (AzureManagedIdentity). */
managedIdentityClientId?: string;
/**
* Seconds before expiry to trigger proactive refresh (default: 300).
* Keep this well below the token TTL; if it is greater than or equal to
* the TTL, each request refreshes the token.
*/
refreshBufferSecs?: number;
}
-3
View File
@@ -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));
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-darwin-arm64",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["darwin"],
"cpu": ["arm64"],
"main": "lancedb.darwin-arm64.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-gnu",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-musl",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-gnu",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-musl",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-arm64-msvc",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": [
"win32"
],
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-x64-msvc",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"os": ["win32"],
"cpu": ["x64"],
"main": "lancedb.win32-x64-msvc.node",
+2 -2
View File
@@ -1,12 +1,12 @@
{
"name": "@lancedb/lancedb",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "@lancedb/lancedb",
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"cpu": [
"x64",
"arm64"
+1 -1
View File
@@ -11,7 +11,7 @@
"ann"
],
"private": false,
"version": "0.31.0-beta.6",
"version": "0.31.0-beta.2",
"main": "dist/index.js",
"exports": {
".": "./dist/index.js",
-6
View File
@@ -112,12 +112,6 @@ impl Connection {
builder = builder.client_config(rust_config);
if let Some(oauth_config) = options.oauth_config {
let config: lancedb::remote::oauth::OAuthConfig =
oauth_config.try_into().default_error()?;
builder = builder.oauth_config(config);
}
if let Some(api_key) = options.api_key {
builder = builder.api_key(&api_key);
}
-5
View File
@@ -65,11 +65,6 @@ pub struct ConnectionOptions {
/// (For LanceDB cloud only): the host to use for LanceDB cloud. Used
/// for testing purposes.
pub host_override: Option<String>,
/// (For LanceDB cloud only): OAuth configuration for IdP-based
/// authentication (e.g., Azure Entra ID). When set, token acquisition
/// and refresh are handled entirely in Rust. TypeScript users should pass
/// the public `OAuthConfig` type exported from `@lancedb/lancedb`.
pub oauth_config: Option<remote::OAuthConfig>,
}
#[napi(object)]
+6 -4
View File
@@ -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();
+1 -7
View File
@@ -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(),
)
})
-121
View File
@@ -3,7 +3,6 @@
use std::collections::HashMap;
use lancedb::error::Error;
use napi_derive::*;
/// Timeout configuration for remote HTTP client.
@@ -141,84 +140,6 @@ impl From<TlsConfig> for lancedb::remote::TlsConfig {
}
}
/// OAuth configuration for LanceDB authentication.
///
/// This is the generated napi-rs binding shape. TypeScript users should prefer
/// the public `OAuthConfig` type exported from `@lancedb/lancedb`.
///
/// All token acquisition and refresh is handled in the Rust layer.
#[napi(object)]
#[derive(Clone)]
pub struct OAuthConfig {
/// OIDC issuer URL or OAuth authority URL.
/// For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
pub issuer_url: String,
/// Application / Client ID.
pub client_id: String,
/// OAuth scopes to request. For Azure managed identity, exactly one scope
/// or resource is required. For example: `["api://{app_id}/.default"]`
pub scopes: Vec<String>,
/// Authentication flow: "client_credentials" or "azure_managed_identity"
pub flow: Option<String>,
/// Client secret (required for client_credentials).
pub client_secret: Option<String>,
/// Client ID for user-assigned managed identity (azure_managed_identity).
pub managed_identity_client_id: Option<String>,
/// Seconds before expiry to trigger proactive refresh (default: 300).
/// Keep this well below the token TTL; if it is greater than or equal to
/// the TTL, each request refreshes the token.
pub refresh_buffer_secs: Option<u32>,
}
impl std::fmt::Debug for OAuthConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OAuthConfig")
.field("issuer_url", &self.issuer_url)
.field("client_id", &self.client_id)
.field("scopes", &self.scopes)
.field("flow", &self.flow)
.field(
"client_secret",
&self.client_secret.as_deref().map(|_| "<redacted>"),
)
.field(
"managed_identity_client_id",
&self.managed_identity_client_id,
)
.field("refresh_buffer_secs", &self.refresh_buffer_secs)
.finish()
}
}
impl TryFrom<OAuthConfig> for lancedb::remote::oauth::OAuthConfig {
type Error = Error;
fn try_from(config: OAuthConfig) -> Result<Self, Self::Error> {
use lancedb::remote::oauth::OAuthFlow;
let flow = match config.flow.as_deref().unwrap_or("client_credentials") {
"client_credentials" => OAuthFlow::ClientCredentials,
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
client_id: config.managed_identity_client_id,
},
other => {
return Err(Error::InvalidInput {
message: format!("Unknown OAuth flow type: {other}"),
});
}
};
Ok(Self {
issuer_url: config.issuer_url,
client_id: config.client_id,
client_secret: config.client_secret,
scopes: config.scopes,
flow,
refresh_buffer_secs: config.refresh_buffer_secs.map(|v| v as u64),
})
}
}
impl From<ClientConfig> for lancedb::remote::ClientConfig {
fn from(config: ClientConfig) -> Self {
Self {
@@ -235,45 +156,3 @@ impl From<ClientConfig> for lancedb::remote::ClientConfig {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_unknown_oauth_flow_returns_invalid_input() {
let config = OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["scope".to_string()],
flow: Some("typo".to_string()),
client_secret: None,
managed_identity_client_id: None,
refresh_buffer_secs: None,
};
let err = lancedb::remote::oauth::OAuthConfig::try_from(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "Unknown OAuth flow type: typo"
));
}
#[test]
fn test_oauth_config_debug_redacts_client_secret() {
let config = OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["scope".to_string()],
flow: Some("client_credentials".to_string()),
client_secret: Some("super-secret".to_string()),
managed_identity_client_id: None,
refresh_buffer_secs: None,
};
let debug = format!("{config:?}");
assert!(!debug.contains("super-secret"));
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
}
}
+1 -1
View File
@@ -1,5 +1,5 @@
[tool.bumpversion]
current_version = "0.34.0-beta.6"
current_version = "0.34.0-beta.2"
parse = """(?x)
(?P<major>0|[1-9]\\d*)\\.
(?P<minor>0|[1-9]\\d*)\\.
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.34.0-beta.6"
version = "0.34.0-beta.2"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -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()
-8
View File
@@ -89,8 +89,6 @@ def connect(
If presented, connect to LanceDB cloud.
Otherwise, connect to a database on file system or cloud storage.
Can be set via environment variable `LANCEDB_API_KEY`.
OAuth configuration is currently supported only by ``connect_async``;
synchronous LanceDB Cloud connections require an API key.
region: str, default "us-east-1"
The region to use for LanceDB Cloud.
host_override: str, optional
@@ -342,7 +340,6 @@ async def connect_async(
session: Optional[Session] = None,
manifest_enabled: bool = False,
namespace_client_properties: Optional[Dict[str, str]] = None,
oauth_config=None,
) -> AsyncConnection:
"""Connect to a LanceDB database.
@@ -392,10 +389,6 @@ async def connect_async(
namespace_client_properties : dict, optional
Additional directory namespace client properties to use with
``manifest_enabled=True``.
oauth_config : OAuthConfig, optional
OAuth configuration for LanceDB Cloud/Enterprise. This is supported by
``connect_async`` only; synchronous ``connect`` uses API key
authentication for ``db://`` URIs.
Examples
--------
@@ -442,7 +435,6 @@ async def connect_async(
session,
manifest_enabled,
namespace_client_properties,
oauth_config,
)
)
+2 -23
View File
@@ -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:
@@ -283,24 +280,6 @@ async def connect(
session: Optional[Session],
manifest_enabled: bool = False,
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:
+3 -17
View File
@@ -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))
-9
View File
@@ -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
+169 -258
View File
@@ -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
@@ -364,23 +373,6 @@ def _convert_pyarrow_schema_to_json(schema: pa.Schema) -> JsonArrowSchema:
return JsonArrowSchema(fields=fields, metadata=meta)
def _builds_namespace_natively(
namespace_client_impl: Optional[str],
namespace_client_properties: Optional[Dict[str, str]],
) -> bool:
"""Whether ``connect_namespace_client`` builds the namespace client natively
in Rust (installing the read-freshness context provider) rather than wrapping
the pre-built Python client.
Must mirror Rust ``build_namespace_natively`` in ``python/src/connection.rs``.
"""
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 +383,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 +391,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 +432,23 @@ 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(
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 +494,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(
@@ -565,7 +543,6 @@ class LanceNamespaceDBConnection(DBConnection):
namespace_path=namespace_path,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
_async=async_table,
)
@@ -591,8 +568,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
@@ -603,7 +580,6 @@ class LanceNamespaceDBConnection(DBConnection):
namespace_path=namespace_path,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
_async=async_table,
)
if branch is not None:
@@ -614,9 +590,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 +609,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 +628,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 +659,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 +693,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 +728,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 +760,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 +794,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 +856,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,15 +869,12 @@ 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,
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,
_inner: Optional[AsyncConnection] = None,
):
"""
Initialize an async namespace-based LanceDB connection.
@@ -926,12 +900,6 @@ class AsyncLanceNamespaceDBConnection:
namespace.create_table() instead of using declare_table + local write.
Default is None (no pushdown, all operations run locally).
namespace_client_impl : Optional[str]
The namespace implementation name used to create this connection.
Required (with ``namespace_client_properties``) for the Rust client to
be built natively and install the read-freshness provider.
namespace_client_properties : Optional[Dict[str, str]]
The namespace properties used to create this connection.
"""
self._namespace_client = namespace_client
self.read_consistency_interval = read_consistency_interval
@@ -940,37 +908,23 @@ class AsyncLanceNamespaceDBConnection:
self._namespace_client_pushdown_operations = set(
namespace_client_pushdown_operations or []
)
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(
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=None,
namespace_client_properties=None,
)
)
async def table_names(
self,
@@ -994,9 +948,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,
@@ -1036,7 +992,6 @@ class AsyncLanceNamespaceDBConnection:
namespace_path=namespace_path,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
)
async def open_table(
@@ -1059,8 +1014,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
@@ -1074,14 +1029,15 @@ class AsyncLanceNamespaceDBConnection:
namespace_path=namespace_path,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
)
async def drop_table(self, name: str, namespace_path: Optional[List[str]] = None):
"""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 +1051,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 +1070,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 +1101,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 +1134,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 +1168,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 +1199,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 +1232,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 +1252,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 +1302,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 +1377,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
)
@@ -1474,6 +1387,4 @@ def connect_namespace_async(
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,
)
-13
View File
@@ -48,14 +48,6 @@ class PermutationBuilder:
By default, the permutation builder will create a single split that contains all
rows in the same order as the base table.
"""
if not hasattr(table, "_inner"):
raise TypeError(
f"PermutationBuilder requires a local LanceTable, "
f"got {type(table).__name__}. "
"The permutation API is not supported on remote tables. "
"Remote tables connect to LanceDB Cloud or Enterprise and do not have "
"direct access to the underlying Lance dataset needed for permutations."
)
self._async = async_permutation_builder(table)
def split_random(
@@ -65,7 +57,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 +79,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 +87,6 @@ class PermutationBuilder:
counts=counts,
fixed=fixed,
seed=seed,
clump_size=clump_size,
split_names=split_names,
)
return self
+2 -36
View File
@@ -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)
-3
View File
@@ -9,7 +9,6 @@ from typing import List, Optional
from lancedb import __version__
from .header import HeaderProvider
from .oauth import OAuthConfig, OAuthFlowType
__all__ = [
"TimeoutConfig",
@@ -17,8 +16,6 @@ __all__ = [
"TlsConfig",
"ClientConfig",
"HeaderProvider",
"OAuthConfig",
"OAuthFlowType",
]
-75
View File
@@ -1,75 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional
class OAuthFlowType(str, Enum):
"""OAuth authentication flow types."""
CLIENT_CREDENTIALS = "client_credentials"
"""Client Credentials grant (service-to-service / M2M)."""
AZURE_MANAGED_IDENTITY = "azure_managed_identity"
"""Azure Managed Identity via IMDS."""
@dataclass
class OAuthConfig:
"""OAuth configuration for LanceDB authentication.
All token acquisition and refresh is handled in the Rust layer.
This config is passed through to Rust via PyO3.
Parameters
----------
issuer_url : str
OIDC issuer URL or OAuth authority URL.
For Azure: ``https://login.microsoftonline.com/{tenant_id}/v2.0``
client_id : str
Application / Client ID.
scopes : List[str]
OAuth scopes to request.
For Azure managed identity, exactly one scope or resource is required.
For example: ``["api://{app_id}/.default"]``
flow : OAuthFlowType
Authentication flow to use. Default: CLIENT_CREDENTIALS.
client_secret : Optional[str]
Client secret (required for CLIENT_CREDENTIALS).
managed_identity_client_id : Optional[str]
Client ID for user-assigned managed identity (AZURE_MANAGED_IDENTITY).
refresh_buffer_secs : Optional[int]
Seconds before expiry to trigger proactive refresh (default: 300).
Keep this well below the token TTL; if it is greater than or equal to
the TTL, each request refreshes the token.
Examples
--------
Client Credentials (service-to-service):
>>> config = OAuthConfig(
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
... client_id="app-id",
... client_secret="secret",
... scopes=["api://lancedb-api/.default"],
... )
Azure Managed Identity:
>>> config = OAuthConfig(
... issuer_url="https://login.microsoftonline.com/{tenant}/v2.0",
... client_id="app-id",
... scopes=["api://lancedb-api/.default"],
... flow=OAuthFlowType.AZURE_MANAGED_IDENTITY,
... )
"""
issuer_url: str
client_id: str
scopes: List[str]
flow: OAuthFlowType = OAuthFlowType.CLIENT_CREDENTIALS
client_secret: Optional[str] = field(default=None, repr=False)
managed_identity_client_id: Optional[str] = None
refresh_buffer_secs: Optional[int] = None
+1 -8
View File
@@ -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)
-607
View File
@@ -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)
+9 -63
View File
@@ -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),
@@ -2032,7 +2022,6 @@ class LanceTable(Table):
namespace_client: Optional[Any] = None,
managed_versioning: Optional[bool] = None,
pushdown_operations: Optional[set] = None,
route_pushdown_to_rust: bool = False,
_async: AsyncTable = None,
):
if namespace_path is None:
@@ -2042,14 +2031,6 @@ class LanceTable(Table):
self._location = location # Store location for use in _dataset_path
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
# When the connection built the namespace client natively (e.g. an
# enterprise "rest" connection), the underlying Rust table already
# executes QueryTable pushdown itself -- and, unlike this Python urllib3
# path, it routes through the read-freshness context provider that emits
# the ``x-lancedb-min-timestamp`` header. So we must defer pushdown to
# Rust instead of calling the Python ``namespace_client.query_table``
# directly, or reads silently bypass read-freshness (stale results).
self._route_pushdown_to_rust = route_pushdown_to_rust
if _async is not None:
self._table = _async
else:
@@ -2152,19 +2133,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,
)
@@ -2267,7 +2241,6 @@ class LanceTable(Table):
namespace_path=self._namespace_path,
namespace_client=self._namespace_client,
pushdown_operations=self._pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
location=self._location,
_async=async_table,
)
@@ -2418,11 +2391,8 @@ class LanceTable(Table):
Returns
-------
pa.Table"""
if (
_should_push_down_query_table(
self._namespace_client, self._pushdown_operations
)
and not self._route_pushdown_to_rust
if _should_push_down_query_table(
self._namespace_client, self._pushdown_operations
):
return self._execute_query(Query()).read_all()
@@ -3374,7 +3344,6 @@ class LanceTable(Table):
location: Optional[str] = None,
namespace_client: Optional[Any] = None,
pushdown_operations: Optional[set] = None,
route_pushdown_to_rust: bool = False,
):
"""
Create a new table.
@@ -3437,7 +3406,6 @@ class LanceTable(Table):
self._location = location
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
self._route_pushdown_to_rust = route_pushdown_to_rust
if data_storage_version is not None:
warnings.warn(
@@ -3551,7 +3519,6 @@ class LanceTable(Table):
_should_push_down_query_table(
self._namespace_client, self._pushdown_operations
)
and not self._route_pushdown_to_rust
and self.current_branch() is None
):
from lancedb.namespace import _execute_server_side_query
@@ -4030,16 +3997,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()
@@ -4302,7 +4260,6 @@ class AsyncTable:
namespace_path: Optional[List[str]] = None,
namespace_client: Optional[Any] = None,
pushdown_operations: Optional[set] = None,
route_pushdown_to_rust: bool = False,
):
"""Create a new AsyncTable object.
@@ -4315,9 +4272,6 @@ class AsyncTable:
self._namespace_path = namespace_path or []
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
# See LanceTable.__init__: defer QueryTable pushdown to Rust (which emits
# the read-freshness header) for natively-built namespace clients.
self._route_pushdown_to_rust = route_pushdown_to_rust
def _set_namespace_context(
self,
@@ -4325,12 +4279,10 @@ class AsyncTable:
namespace_path: Optional[List[str]] = None,
namespace_client: Optional[Any] = None,
pushdown_operations: Optional[set] = None,
route_pushdown_to_rust: bool = False,
) -> "AsyncTable":
self._namespace_path = namespace_path or []
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
self._route_pushdown_to_rust = route_pushdown_to_rust
return self
def __repr__(self):
@@ -4540,11 +4492,8 @@ class AsyncTable:
-------
pa.Table
"""
if (
_should_push_down_query_table(
self._namespace_client, self._pushdown_operations
)
and not self._route_pushdown_to_rust
if _should_push_down_query_table(
self._namespace_client, self._pushdown_operations
):
return (await self._execute_query(Query())).read_all()
@@ -5228,11 +5177,8 @@ class AsyncTable:
batch_size: Optional[int] = None,
timeout: Optional[timedelta] = None,
) -> pa.RecordBatchReader:
if (
_should_push_down_query_table(
self._namespace_client, self._pushdown_operations
)
and not self._route_pushdown_to_rust
if _should_push_down_query_table(
self._namespace_client, self._pushdown_operations
):
from lancedb.namespace import _execute_server_side_query
File diff suppressed because it is too large Load Diff
-43
View File
@@ -242,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):
-126
View File
@@ -5,7 +5,6 @@
import tempfile
import shutil
import importlib
import pytest
import pyarrow as pa
import lancedb
@@ -66,9 +65,6 @@ def _namespace_lance_table(namespace_client: _NamespaceClient) -> LanceTable:
table._namespace_path = ["geneva"]
table._namespace_client = namespace_client
table._pushdown_operations = {"QueryTable"}
# This test exercises the Python-side pushdown path (non-native client), so
# pushdown is not routed to Rust.
table._route_pushdown_to_rust = False
return table
@@ -104,40 +100,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 +561,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):
@@ -898,39 +805,6 @@ class TestPushdownOperations:
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
assert len(db._namespace_client_pushdown_operations) == 0
def test_route_pushdown_to_rust_for_native_rest(self):
"""A natively-built rest connection must defer QueryTable pushdown to
Rust so reads carry the x-lancedb-min-timestamp read-freshness header."""
db = lancedb.connect_namespace(
"rest",
{"uri": "http://localhost:12345"},
namespace_client_pushdown_operations=["QueryTable"],
)
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."""
db = lancedb.connect_namespace("dir", {"root": self.temp_dir})
assert db._route_pushdown_to_rust is True
def test_async_route_pushdown_to_rust_for_native_rest(self):
"""The async connection must not silently bypass the read-freshness fix:
a natively-built rest connection defers pushdown to Rust (regression test
for the async path omitting the freshness header)."""
db = lancedb.connect_namespace_async(
"rest",
{"uri": "http://localhost:12345"},
namespace_client_pushdown_operations=["QueryTable"],
)
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."""
db = lancedb.connect_namespace_async("dir", {"root": self.temp_dir})
assert db._route_pushdown_to_rust is True
def test_lance_table_to_arrow_uses_query_pushdown(self):
namespace_client = _NamespaceClient()
table = _namespace_lance_table(namespace_client)
-55
View File
@@ -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()
+3 -5
View File
@@ -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)
-32
View File
@@ -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(
{
-10
View File
@@ -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
+12 -115
View File
@@ -539,7 +539,7 @@ impl Connection {
}
#[pyfunction]
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None, oauth_config=None))]
#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None, manifest_enabled=false, namespace_client_properties=None))]
#[allow(clippy::too_many_arguments)]
pub fn connect(
py: Python<'_>,
@@ -553,7 +553,6 @@ pub fn connect(
session: Option<crate::session::Session>,
manifest_enabled: bool,
namespace_client_properties: Option<HashMap<String, String>>,
oauth_config: Option<crate::oauth::PyOAuthConfig>,
) -> PyResult<Bound<'_, PyAny>> {
future_into_py(py, async move {
let mut builder = lancedb::connect(&uri);
@@ -583,11 +582,6 @@ pub fn connect(
if let Some(client_config) = client_config {
builder = builder.client_config(client_config.into());
}
if let Some(oauth_config) = oauth_config {
let config: lancedb::remote::oauth::OAuthConfig =
oauth_config.try_into().infer_error()?;
builder = builder.oauth_config(config);
}
if let Some(session) = session {
builder = builder.session(session.inner.clone());
}
@@ -616,38 +610,24 @@ pub fn connect_namespace_client(
namespace_client_impl: Option<String>,
namespace_client_properties: Option<HashMap<String, String>>,
) -> PyResult<Connection> {
let namespace_client = extract_namespace_arc(py, namespace_client)?;
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 ns_impl = namespace_client_impl.unwrap_or_else(|| "python".to_string());
let ns_properties = namespace_client_properties.unwrap_or_default();
let storage_options = storage_options.unwrap_or_default();
let session = session.map(|s| s.inner.clone());
// Prefer building the namespace natively from (impl, properties) so the
// read-freshness provider installed
let database = if build_namespace_natively(namespace_client_impl.as_deref(), &ns_properties) {
let ns_impl = namespace_client_impl.expect("impl present per build_namespace_natively");
crate::runtime::block_on(LanceNamespaceDatabase::connect(
&ns_impl,
ns_properties,
storage_options,
read_consistency_interval,
session,
namespace_client_pushdown_operations,
))
.infer_error()?
} else {
let namespace_client = extract_namespace_arc(py, namespace_client)?;
LanceNamespaceDatabase::from_namespace_client(
namespace_client,
namespace_client_impl.unwrap_or_else(|| "python".to_string()),
ns_properties,
storage_options,
read_consistency_interval,
session,
namespace_client_pushdown_operations,
)
};
let database = LanceNamespaceDatabase::from_namespace_client(
namespace_client,
ns_impl,
ns_properties,
storage_options,
read_consistency_interval,
session,
namespace_client_pushdown_operations,
);
Ok(Connection::new(LanceConnection::new(
Arc::new(database),
@@ -655,56 +635,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
fn build_namespace_natively(
namespace_client_impl: Option<&str>,
namespace_client_properties: &HashMap<String, String>,
) -> bool {
matches!(namespace_client_impl, Some("rest")) && !namespace_client_properties.is_empty()
}
#[derive(FromPyObject)]
pub struct PyClientConfig {
user_agent: String,
@@ -803,36 +733,3 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn props(pairs: &[(&str, &str)]) -> HashMap<String, String> {
pairs
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect()
}
#[test]
fn native_build_only_for_rest_with_properties() {
let rest = props(&[("uri", "http://localhost:10024")]);
// rest + non-empty properties -> build natively (installs the
// read-freshness provider so checkout_latest() busts the server cache).
assert!(build_namespace_natively(Some("rest"), &rest));
// dir is local (no server cache) -> wrap the pre-built client unchanged.
assert!(!build_namespace_natively(
Some("dir"),
&props(&[("root", "/tmp")])
));
// No impl: only a pre-built client was handed in -> wrap it as-is.
assert!(!build_namespace_natively(None, &rest));
// rest but no properties: nothing to build a connection from -> wrap.
assert!(!build_namespace_natively(Some("rest"), &HashMap::new()));
}
}
+8 -60
View File
@@ -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.
+1 -3
View File
@@ -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;
@@ -26,7 +26,6 @@ pub mod expr;
pub mod header;
pub mod index;
pub mod namespace;
pub mod oauth;
pub mod permutation;
pub mod query;
pub mod runtime;
@@ -62,7 +61,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)?)?;
-72
View File
@@ -1,72 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use pyo3::FromPyObject;
use lancedb::error::Error;
use lancedb::remote::oauth::{OAuthConfig, OAuthFlow};
/// Python-side OAuth configuration, extracted via FromPyObject.
/// Maps to `lancedb.remote.oauth.OAuthConfig` Python dataclass.
#[derive(FromPyObject)]
pub struct PyOAuthConfig {
pub issuer_url: String,
pub client_id: String,
pub scopes: Vec<String>,
pub flow: String,
pub client_secret: Option<String>,
pub managed_identity_client_id: Option<String>,
pub refresh_buffer_secs: Option<u64>,
}
impl TryFrom<PyOAuthConfig> for OAuthConfig {
type Error = Error;
fn try_from(py: PyOAuthConfig) -> Result<Self, Self::Error> {
let flow = match py.flow.as_str() {
"client_credentials" => OAuthFlow::ClientCredentials,
"azure_managed_identity" => OAuthFlow::AzureManagedIdentity {
client_id: py.managed_identity_client_id,
},
other => {
return Err(Error::InvalidInput {
message: format!("Unknown OAuth flow type: {other}"),
});
}
};
Ok(Self {
issuer_url: py.issuer_url,
client_id: py.client_id,
client_secret: py.client_secret,
scopes: py.scopes,
flow,
refresh_buffer_secs: py.refresh_buffer_secs,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_unknown_oauth_flow_returns_invalid_input() {
let config = PyOAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
scopes: vec!["scope".to_string()],
flow: "typo".to_string(),
client_secret: None,
managed_identity_client_id: None,
refresh_buffer_secs: None,
};
let err = OAuthConfig::try_from(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "Unknown OAuth flow type: typo"
));
}
}
+2 -10
View File
@@ -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)
})
}
-9
View File
@@ -56,15 +56,6 @@ fn get_runtime() -> &'static runtime::Runtime {
unsafe { &*new_ptr }
}
/// Block the current thread on a future using the shared runtime.
///
/// For sync `#[pyfunction]`s that need to drive an async operation (e.g.
/// building a namespace client). Must not be called from within the runtime's
/// own worker threads.
pub fn block_on<F: std::future::Future>(fut: F) -> F::Output {
get_runtime().block_on(fut)
}
/// Runs in async-signal context after `fork()` in the child. We can only
/// touch atomics here; we deliberately leak the previous runtime because
/// dropping a tokio `Runtime` would try to join its (now-dead) worker
@@ -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 ───────────────────────────────────
-33
View File
@@ -1,33 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import importlib.util
import sys
from pathlib import Path
def _load_oauth_module():
oauth_path = (
Path(__file__).parents[1] / "python" / "lancedb" / "remote" / "oauth.py"
)
spec = importlib.util.spec_from_file_location("lancedb_remote_oauth", oauth_path)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
def test_oauth_config_repr_redacts_client_secret():
oauth = _load_oauth_module()
config = oauth.OAuthConfig(
issuer_url="https://issuer.example.com",
client_id="client-id",
scopes=["scope"],
client_secret="super-secret",
)
rendered = repr(config)
assert "super-secret" not in rendered
assert "client_secret" not in rendered
+3 -21
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb"
version = "0.31.0-beta.6"
version = "0.31.0-beta.2"
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 }
@@ -51,7 +50,7 @@ lance-namespace = { workspace = true }
lance-namespace-impls = { workspace = true }
moka = { workspace = true }
pin-project = { workspace = true }
tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] }
tokio = { version = "1.23", features = ["rt-multi-thread"] }
log.workspace = true
async-trait = "0"
bytes = "1"
@@ -76,7 +75,6 @@ reqwest = { version = "0.12.0", default-features = false, features = [
"stream",
], optional = true }
http = { version = "1", optional = true } # Matching what is in reqwest
urlencoding = { version = "2", optional = true }
uuid = { version = "1.7.0", features = ["v4", "v5"] }
polars-arrow = { version = ">=0.37,<0.40.0", optional = true }
polars = { version = ">=0.37,<0.40.0", optional = true }
@@ -95,7 +93,6 @@ semver = { workspace = true }
anyhow = "1"
tempfile = "3.5.0"
random_word = { version = "0.4.3", features = ["en"] }
tokio = { version = "1.23", features = ["io-util", "macros", "net", "rt-multi-thread", "sync"] }
uuid = { version = "1.7.0", features = ["v4"] }
walkdir = "2"
aws-sdk-dynamodb = { version = "1.55.0" }
@@ -108,8 +105,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]
@@ -134,13 +129,7 @@ huggingface = [
"lance-namespace-impls/dir-huggingface",
]
dynamodb = ["lance/dynamodb", "aws"]
remote = [
"dep:reqwest",
"dep:http",
"dep:urlencoding",
"lance-namespace-impls/rest",
"lance-namespace-impls/rest-adapter",
]
remote = ["dep:reqwest", "dep:http", "lance-namespace-impls/rest", "lance-namespace-impls/rest-adapter"]
fp16kernels = ["lance-linalg/fp16kernels"]
s3-test = []
bedrock = ["dep:aws-sdk-bedrockruntime"]
@@ -166,16 +155,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(())
}
-47
View File
@@ -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
View File
@@ -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]
+4 -129
View File
@@ -667,8 +667,6 @@ pub struct ConnectRequest {
pub struct ConnectBuilder {
request: ConnectRequest,
embedding_registry: Option<Arc<dyn EmbeddingRegistry>>,
#[cfg(feature = "remote")]
oauth_config: Option<crate::remote::OAuthConfig>,
}
#[cfg(feature = "remote")]
@@ -690,8 +688,6 @@ impl ConnectBuilder {
session: None,
},
embedding_registry: None,
#[cfg(feature = "remote")]
oauth_config: None,
}
}
@@ -780,19 +776,6 @@ impl ConnectBuilder {
self
}
/// Configure OAuth authentication for LanceDB Cloud/Enterprise.
///
/// This creates an [`OAuthHeaderProvider`](crate::remote::OAuthHeaderProvider)
/// from the given config and sets it as the header provider. OAuth cannot
/// be combined with an API key or another header provider.
///
/// Token acquisition and refresh are handled in Rust.
#[cfg(feature = "remote")]
pub fn oauth_config(mut self, config: crate::remote::OAuthConfig) -> Self {
self.oauth_config = Some(config);
self
}
/// Provide a custom [`EmbeddingRegistry`] to use for this connection.
pub fn embedding_registry(mut self, registry: Arc<dyn EmbeddingRegistry>) -> Self {
self.embedding_registry = Some(registry);
@@ -938,40 +921,9 @@ impl ConnectBuilder {
let region = options.region.ok_or_else(|| Error::InvalidInput {
message: "A region is required when connecting to LanceDb Cloud".to_string(),
})?;
let api_key = match (&self.oauth_config, &options.api_key) {
(Some(_), None) => String::new(),
(Some(_), Some(_)) => {
return Err(Error::InvalidInput {
message:
"api_key and oauth_config cannot both be set when connecting to LanceDb Cloud"
.to_string(),
});
}
(None, Some(key)) => key.clone(),
(None, None) => {
return Err(Error::InvalidInput {
message:
"An api_key or oauth_config is required when connecting to LanceDb Cloud"
.to_string(),
});
}
};
if self.oauth_config.is_some() && self.request.client_config.header_provider.is_some() {
return Err(Error::InvalidInput {
message:
"oauth_config and client_config.header_provider cannot both be set when connecting to LanceDb Cloud"
.to_string(),
});
}
let mut client_config = self.request.client_config;
if let Some(oauth_config) = self.oauth_config {
let provider = crate::remote::OAuthHeaderProvider::new(oauth_config)?;
client_config.header_provider =
Some(Arc::new(provider) as Arc<dyn crate::remote::HeaderProvider>);
}
let api_key = options.api_key.ok_or_else(|| Error::InvalidInput {
message: "An api_key is required when connecting to LanceDb Cloud".to_string(),
})?;
let storage_options = StorageOptions(options.storage_options.clone());
let internal = Arc::new(crate::remote::db::RemoteDatabase::try_new(
@@ -979,7 +931,7 @@ impl ConnectBuilder {
&api_key,
&region,
options.host_override,
client_config,
self.request.client_config,
storage_options.into(),
self.request.read_consistency_interval,
)?);
@@ -1288,83 +1240,6 @@ mod tests {
assert_eq!(Some(&"EXPLICIT-VALUE".to_string()), options.get(opts_key));
}
#[cfg(feature = "remote")]
#[tokio::test]
async fn test_connect_rejects_api_key_with_oauth_config() {
let oauth_config = crate::remote::OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: crate::remote::OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
let result = ConnectBuilder::new("db://my-container/my-prefix")
.region("us-east-1")
.api_key("my-api-key")
.oauth_config(oauth_config)
.execute()
.await;
match result {
Err(Error::InvalidInput { message })
if message
== "api_key and oauth_config cannot both be set when connecting to LanceDb Cloud" =>
{}
Err(err) => panic!("expected InvalidInput, got {err:?}"),
Ok(_) => panic!("expected api_key and oauth_config to be rejected"),
}
}
#[cfg(feature = "remote")]
#[tokio::test]
async fn test_connect_rejects_header_provider_with_oauth_config() {
#[derive(Debug)]
struct TestHeaderProvider;
#[async_trait::async_trait]
impl crate::remote::HeaderProvider for TestHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
Ok(HashMap::from([(
"authorization".to_string(),
"Bearer token".to_string(),
)]))
}
}
let oauth_config = crate::remote::OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: crate::remote::OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
let client_config = crate::remote::ClientConfig {
header_provider: Some(
Arc::new(TestHeaderProvider) as Arc<dyn crate::remote::HeaderProvider>
),
..Default::default()
};
let result = ConnectBuilder::new("db://my-container/my-prefix")
.region("us-east-1")
.client_config(client_config)
.oauth_config(oauth_config)
.execute()
.await;
match result {
Err(Error::InvalidInput { message })
if message
== "oauth_config and client_config.header_provider cannot both be set when connecting to LanceDb Cloud" =>
{}
Err(err) => panic!("expected InvalidInput, got {err:?}"),
Ok(_) => panic!("expected header_provider and oauth_config to be rejected"),
}
}
#[cfg(not(windows))]
#[tokio::test]
async fn test_connect_relative() {
-93
View File
@@ -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);
}
}
}
+262 -122
View File
@@ -14,7 +14,6 @@ use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore};
use lance_datafusion::utils::StreamingWriteSource;
use lance_encoding::version::LanceFileVersion;
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
use lance_table::io::commit::commit_handler_from_url;
use object_store::local::LocalFileSystem;
use snafu::ResultExt;
@@ -235,11 +234,9 @@ impl ListingDatabaseOptionsBuilder {
/// We will have two tables named `table1` and `table2`.
#[derive(Debug)]
pub struct ListingDatabase {
object_store: Arc<ObjectStore>,
query_string: Option<String>,
pub(crate) uri: String,
pub(crate) base_path: object_store::path::Path,
// the object store wrapper to use on write path
pub(crate) store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
@@ -258,8 +255,13 @@ pub struct ListingDatabase {
// Session for object stores and caching
session: Arc<lance::session::Session>,
// Namespace-backed database for child namespace operations
// Namespace-backed database for child namespace operations (manifest mode).
namespace_database: Arc<LanceNamespaceDatabase>,
// V1 (manifest-disabled) directory namespace for root table lifecycle, so root
// drops are soft-deletes and purge/table_status are available. Shares the same root
// as `namespace_database` but in directory mode.
root_namespace_database: Arc<LanceNamespaceDatabase>,
}
impl std::fmt::Display for ListingDatabase {
@@ -280,7 +282,6 @@ impl std::fmt::Display for ListingDatabase {
}
}
const LANCE_EXTENSION: &str = "lance";
const ENGINE: &str = "engine";
const MIRRORED_STORE: &str = "mirroredStore";
@@ -342,6 +343,39 @@ impl ListingDatabase {
))
}
/// Build the V1 (manifest-disabled) directory namespace used for *root* table
/// lifecycle ops.
///
/// Root tables in a listing database are flat `<name>.lance` directories; soft-delete
/// (drop/purge/TTL) is a V1-only mechanism, so root ops go through this namespace.
/// Child namespaces are manifest-backed and handled by the separate
/// (manifest-enabled) `namespace_database`.
async fn connect_root_namespace_database(
uri: &str,
storage_options: HashMap<String, String>,
namespace_client_properties: HashMap<String, String>,
read_consistency_interval: Option<std::time::Duration>,
session: Arc<lance::session::Session>,
) -> Result<Arc<LanceNamespaceDatabase>> {
let mut ns_properties = Self::build_namespace_client_properties(
uri,
&storage_options,
namespace_client_properties,
);
ns_properties.insert("manifest_enabled".to_string(), "false".to_string());
Ok(Arc::new(
LanceNamespaceDatabase::connect(
"dir",
ns_properties,
storage_options,
read_consistency_interval,
Some(session),
HashSet::new(),
)
.await?,
))
}
async fn prepare_namespace_root(
uri: &str,
storage_options: &HashMap<String, String>,
@@ -548,7 +582,7 @@ impl ListingDatabase {
},
..Default::default()
};
let (object_store, base_path) = ObjectStore::from_uri_and_params(
let (object_store, _base_path) = ObjectStore::from_uri_and_params(
session.store_registry(),
&plain_uri,
&os_params,
@@ -577,12 +611,18 @@ impl ListingDatabase {
session.clone(),
)
.await?;
let root_namespace_database = Self::connect_root_namespace_database(
&table_base_uri,
options.storage_options.clone(),
request.namespace_client_properties.clone(),
request.read_consistency_interval,
session.clone(),
)
.await?;
Ok(Self {
uri: table_base_uri,
query_string,
base_path,
object_store,
store_wrapper: write_store_wrapper,
read_consistency_interval: request.read_consistency_interval,
storage_options: options.storage_options,
@@ -590,6 +630,7 @@ impl ListingDatabase {
new_table_config: options.new_table_config,
session,
namespace_database,
root_namespace_database,
})
}
Err(_) => {
@@ -613,7 +654,7 @@ impl ListingDatabase {
session: Option<Arc<lance::session::Session>>,
) -> Result<Self> {
let session = session.unwrap_or_else(|| Arc::new(lance::session::Session::default()));
let (object_store, base_path) = ObjectStore::from_uri_and_params(
let (object_store, _base_path) = ObjectStore::from_uri_and_params(
session.store_registry(),
path,
&ObjectStoreParams::default(),
@@ -624,6 +665,14 @@ impl ListingDatabase {
}
let namespace_database = Self::connect_namespace_database(
path,
HashMap::new(),
namespace_client_properties.clone(),
read_consistency_interval,
session.clone(),
)
.await?;
let root_namespace_database = Self::connect_root_namespace_database(
path,
HashMap::new(),
namespace_client_properties,
@@ -635,8 +684,6 @@ impl ListingDatabase {
Ok(Self {
uri: path.to_string(),
query_string: None,
base_path,
object_store,
store_wrapper: None,
read_consistency_interval,
storage_options: HashMap::new(),
@@ -644,6 +691,7 @@ impl ListingDatabase {
new_table_config,
session,
namespace_database,
root_namespace_database,
})
}
@@ -705,42 +753,10 @@ impl ListingDatabase {
self.namespace_database.clone()
}
async fn drop_tables(&self, names: Vec<String>) -> Result<()> {
let object_store_params = ObjectStoreParams {
storage_options_accessor: if self.storage_options.is_empty() {
None
} else {
Some(Arc::new(StorageOptionsAccessor::with_static_options(
self.storage_options.clone(),
)))
},
..Default::default()
};
let mut uri = self.uri.clone();
if let Some(query_string) = &self.query_string {
uri.push_str(&format!("?{}", query_string));
}
let commit_handler = commit_handler_from_url(&uri, &Some(object_store_params)).await?;
for name in names {
let dir_name = format!("{}.{}", name, LANCE_EXTENSION);
let full_path = self.base_path.clone().join(dir_name.clone());
commit_handler.delete(&full_path).await?;
self.object_store
.remove_dir_all(full_path.clone())
.await
.map_err(|err| match err {
// this error is not lance::Error::DatasetNotFound, as the method
// `remove_dir_all` may be used to remove something not be a dataset
lance::Error::NotFound { .. } => Error::TableNotFound {
name: name.clone(),
source: Box::new(err),
},
_ => Error::from(err),
})?;
}
Ok(())
/// The V1 directory namespace used for root table lifecycle (soft-delete drop, purge,
/// table_status, O(1) listing).
fn root_namespace_database(&self) -> Arc<LanceNamespaceDatabase> {
self.root_namespace_database.clone()
}
/// Inherit storage options from the connection into the target map
@@ -946,88 +962,43 @@ impl Database for ListingDatabase {
if !request.namespace_path.is_empty() {
return self.namespace_database().table_names(request).await;
}
let mut f = self
.object_store
.read_dir(self.base_path.clone())
.await?
.iter()
.map(Path::new)
.filter(|path| {
let is_lance = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e == LANCE_EXTENSION);
is_lance.unwrap_or(false)
})
.filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from)))
.collect::<Vec<String>>();
f.sort();
if let Some(start_after) = request.start_after {
let index = f
.iter()
.position(|name| name.as_str() > start_after.as_str())
.unwrap_or(f.len());
f.drain(0..index);
}
if let Some(limit) = request.limit {
f.truncate(limit as usize);
}
Ok(f)
// Root tables: the V1 namespace lists them in a single read_dir (O(1) requests)
// and excludes soft-deleted tables, instead of a per-table probe here.
self.root_namespace_database().table_names(request).await
}
async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
if request.id.as_ref().map(|v| !v.is_empty()).unwrap_or(false) {
return self.namespace_database().list_tables(request).await;
}
let mut f = self
.object_store
.read_dir(self.base_path.clone())
.await?
.iter()
.map(Path::new)
.filter(|path| {
let is_lance = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e == LANCE_EXTENSION);
is_lance.unwrap_or(false)
})
.filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from)))
.collect::<Vec<String>>();
f.sort();
// Handle pagination with page_token
if let Some(ref page_token) = request.page_token {
let index = f
.iter()
.position(|name| name.as_str() > page_token.as_str())
.unwrap_or(f.len());
f.drain(0..index);
}
// Determine if there's a next page
let next_page_token = if let Some(limit) = request.limit {
if f.len() > limit as usize {
let token = f[limit as usize].clone();
f.truncate(limit as usize);
Some(token)
} else {
None
}
} else {
None
};
Ok(ListTablesResponse {
tables: f,
page_token: next_page_token,
})
self.root_namespace_database().list_tables(request).await
}
async fn create_table(&self, request: CreateTableRequest) -> Result<Arc<dyn BaseTable>> {
if !request.namespace_path.is_empty() {
return self.namespace_database().create_table(request).await;
}
let mut request = request;
// Re-creating a soft-deleted table is a revive: clear the delete marker (via the
// V1 root namespace, under its lifecycle lock so a concurrent purge can't race),
// making the table live again, then overwrite its data through the native create
// path below (preserving lineage as a new version). A plain native create would
// leave the marker in place, keeping the table hidden.
if matches!(
self.root_namespace_database()
.namespace_client()
.await?
.table_status(Some(vec![request.name.clone()]))
.await?,
lance_namespace::TableLifecycle::SoftDeleted { .. }
) {
self.root_namespace_database()
.namespace_client()
.await?
.undelete_table(Some(vec![request.name.clone()]))
.await?;
request.mode = CreateTableMode::Overwrite;
}
// Use provided location if available, otherwise derive from table name
let table_uri = request
.location
@@ -1146,6 +1117,19 @@ impl Database for ListingDatabase {
if !request.namespace_path.is_empty() {
return self.namespace_database().open_table(request).await;
}
// A soft-deleted (dropped-but-not-purged) table must read as absent even though
// its data still exists on disk. Consult the V1 root namespace (which owns the
// marker); if soft-deleted, route to it so the open surfaces TableNotFound.
if matches!(
self.root_namespace_database()
.namespace_client()
.await?
.table_status(Some(vec![request.name.clone()]))
.await?,
lance_namespace::TableLifecycle::SoftDeleted { .. }
) {
return self.root_namespace_database().open_table(request).await;
}
// Use provided location if available, otherwise derive from table name
let table_uri = request
.location
@@ -1245,20 +1229,23 @@ impl Database for ListingDatabase {
.drop_table(name, namespace_path)
.await;
}
self.drop_tables(vec![name.to_string()]).await
// Root table: route through the V1 namespace so the drop is a soft-delete (writes
// a marker, leaves data for later purge) rather than an immediate remove_dir_all.
self.root_namespace_database()
.drop_table(name, namespace_path)
.await
}
#[allow(deprecated)]
async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()> {
// Check if namespace parameter is provided
if !namespace_path.is_empty() {
return self
.namespace_database()
.drop_all_tables(namespace_path)
.await;
}
let tables = self.table_names(TableNamesRequest::default()).await?;
self.drop_tables(tables).await
self.root_namespace_database()
.drop_all_tables(namespace_path)
.await
}
fn as_any(&self) -> &dyn std::any::Any {
@@ -1266,6 +1253,9 @@ impl Database for ListingDatabase {
}
async fn namespace_client(&self) -> Result<Arc<dyn lance_namespace::LanceNamespace>> {
// Returns the manifest-backed namespace so callers can operate on child
// namespaces (multi-level table ids) through the client. Root-table soft-delete
// lifecycle (table_status/purge) is reached via the V1 root namespace internally.
self.namespace_database.namespace_client().await
}
@@ -2615,4 +2605,154 @@ mod tests {
.unwrap();
assert!(post_drop.tables.is_empty());
}
/// Root-table drop is a soft-delete routed through the V1 namespace: the table is
/// hidden from listing/open but its data survives until purged, and re-creating it
/// revives it. Verifies the consolidation end-to-end at the ListingDatabase level.
#[tokio::test]
async fn test_root_table_soft_delete_lifecycle() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let create = |name: &str| CreateTableRequest {
name: name.to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema.clone())) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
};
let open = |name: &str| OpenTableRequest {
name: name.to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
};
db.create_table(create("t")).await.unwrap();
db.drop_table("t", &[]).await.unwrap();
// Hidden from listing and not openable...
#[allow(deprecated)]
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert!(!names.contains(&"t".to_string()));
assert!(matches!(
db.open_table(open("t")).await,
Err(Error::TableNotFound { .. })
));
// ...but data survives: it shows up as purgable via the V1 root namespace.
let root_ns = db
.root_namespace_database()
.namespace_client()
.await
.unwrap();
let purgable = root_ns.list_purgable_tables(None).await.unwrap();
assert_eq!(purgable.len(), 1);
assert_eq!(purgable[0].id, vec!["t".to_string()]);
// Re-creating revives it.
db.create_table(create("t")).await.unwrap();
db.open_table(open("t")).await.unwrap();
#[allow(deprecated)]
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert!(names.contains(&"t".to_string()));
assert!(root_ns.list_purgable_tables(None).await.unwrap().is_empty());
// Drop then purge reclaims it for good.
db.drop_table("t", &[]).await.unwrap();
let purged = root_ns.purge_tables(None).await.unwrap();
assert_eq!(purged, vec![vec!["t".to_string()]]);
assert!(root_ns.list_purgable_tables(None).await.unwrap().is_empty());
}
/// Migration: a table written by the *old* path (a plain `<root>/<name>.lance`
/// dataset directory with no namespace markers — exactly what a pre-existing listing
/// database looks like on disk) is fully usable by the new embedded-namespace routing.
/// It lists, opens, and reads normally, and dropping it is now a soft-delete (data
/// survives on disk) rather than an immediate removal.
#[tokio::test]
async fn test_legacy_directory_layout_migrates() {
use arrow_array::RecordBatchIterator;
use lance::dataset::{Dataset, WriteParams};
let tempdir = tempdir().unwrap();
let root = tempdir.path().to_str().unwrap().to_string();
// Write a dataset the way the pre-namespace path did: raw lance, no declare/reserve
// marker, no manifest — just `<root>/legacy_table.lance`.
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let legacy_uri = format!("{}/legacy_table.lance", root);
Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
&legacy_uri,
Some(WriteParams::default()),
)
.await
.unwrap();
// Connect the new ListingDatabase over that pre-existing directory.
let request = ConnectRequest {
uri: root.clone(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let open = |name: &str| OpenTableRequest {
name: name.to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
};
// The legacy table is listed, opens, and reads its rows.
#[allow(deprecated)]
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert!(names.contains(&"legacy_table".to_string()));
let table = db.open_table(open("legacy_table")).await.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 3);
// Dropping it is now a soft-delete: hidden from listing / not openable, but the
// dataset directory still exists on disk (reclaimed only by a later purge).
db.drop_table("legacy_table", &[]).await.unwrap();
#[allow(deprecated)]
let names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert!(!names.contains(&"legacy_table".to_string()));
assert!(db.open_table(open("legacy_table")).await.is_err());
assert!(
PathBuf::from(&legacy_uri).exists(),
"soft-delete must leave the legacy dataset directory intact"
);
// And it can be revived by re-creating the name.
let root_ns = db
.root_namespace_database()
.namespace_client()
.await
.unwrap();
assert_eq!(
root_ns.list_purgable_tables(None).await.unwrap().len(),
1,
"the dropped legacy table is purgable"
);
}
}
+3 -3
View File
@@ -583,9 +583,9 @@ impl Database for LanceNamespaceDatabase {
self.namespace
.drop_table(drop_request)
.await
.map_err(|e| Error::Runtime {
message: format!("Failed to drop table: {}", e),
})?;
// Preserve TableNotFound (e.g. dropping a non-existent table) rather than
// flattening every failure to a generic Runtime error.
.map_err(|e| map_namespace_lance_error(e, name))?;
Ok(())
}
@@ -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,
},
);
-6
View File
@@ -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
View File
@@ -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
-2
View File
@@ -8,7 +8,6 @@
pub(crate) mod client;
pub(crate) mod db;
pub mod oauth;
mod retry;
pub(crate) mod table;
pub(crate) mod util;
@@ -21,4 +20,3 @@ const JSON_CONTENT_TYPE: &str = "application/json";
pub use client::{ClientConfig, HeaderProvider, RetryConfig, TimeoutConfig, TlsConfig};
pub use db::{RemoteDatabaseOptions, RemoteDatabaseOptionsBuilder};
pub use oauth::{OAuthConfig, OAuthFlow, OAuthHeaderProvider};
-907
View File
@@ -1,907 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use log::debug;
use reqwest::Client;
use serde::Deserialize;
use tokio::sync::RwLock;
use crate::error::{Error, Result};
use crate::remote::client::HeaderProvider;
const DEFAULT_REFRESH_BUFFER_SECS: u64 = 300;
const DEFAULT_TOKEN_TTL_SECS: u64 = 3600;
const AZURE_IMDS_ENDPOINT: &str = "http://169.254.169.254/metadata/identity/oauth2/token";
const AZURE_IMDS_API_VERSION: &str = "2018-02-01";
/// OAuth authentication flow configuration.
#[derive(Debug, Clone)]
pub enum OAuthFlow {
/// Client Credentials grant (service-to-service / M2M).
/// Requires `client_secret` in [`OAuthConfig`].
ClientCredentials,
/// Azure Managed Identity via IMDS.
/// Works on Azure VMs, AKS, App Service, and Azure Functions.
/// IMDS requests bypass proxy settings because the endpoint is link-local.
AzureManagedIdentity {
/// Client ID for user-assigned managed identity.
/// Omit for system-assigned managed identity.
client_id: Option<String>,
},
}
/// OAuth configuration for LanceDB authentication.
///
/// All token acquisition and refresh is handled in the Rust layer.
/// Python and TypeScript bindings expose this as a plain config object.
#[derive(Clone)]
pub struct OAuthConfig {
/// OIDC issuer URL or OAuth authority URL.
/// For Azure: `https://login.microsoftonline.com/{tenant_id}/v2.0`
pub issuer_url: String,
/// Application / Client ID.
pub client_id: String,
/// Client secret (required for `ClientCredentials`, optional for others).
pub client_secret: Option<String>,
/// OAuth scopes to request.
/// For Azure managed identity, exactly one scope or resource is required.
/// For example: `["api://{app_id}/.default"]`
pub scopes: Vec<String>,
/// Authentication flow to use.
pub flow: OAuthFlow,
/// Seconds before token expiry to trigger proactive refresh (default: 300).
/// Keep this well below the token TTL; if it is greater than or equal to
/// the TTL, each request refreshes the token.
pub refresh_buffer_secs: Option<u64>,
}
impl std::fmt::Debug for OAuthConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OAuthConfig")
.field("issuer_url", &self.issuer_url)
.field("client_id", &self.client_id)
.field(
"client_secret",
&self.client_secret.as_deref().map(|_| "<redacted>"),
)
.field("scopes", &self.scopes)
.field("flow", &self.flow)
.field("refresh_buffer_secs", &self.refresh_buffer_secs)
.finish()
}
}
// -- OIDC Discovery --
#[derive(Clone, Debug, Deserialize)]
struct OidcDiscovery {
token_endpoint: String,
}
// -- Token Response --
#[derive(Deserialize)]
struct TokenResponse {
access_token: String,
/// Token lifetime in seconds.
/// Some providers (Azure IMDS) return this as a string, so we accept both.
#[serde(default, deserialize_with = "deserialize_optional_u64_or_string")]
expires_in: Option<u64>,
#[serde(default)]
#[allow(dead_code)]
token_type: Option<String>,
}
impl std::fmt::Debug for TokenResponse {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TokenResponse")
.field("access_token", &"<redacted>")
.field("expires_in", &self.expires_in)
.field("token_type", &self.token_type)
.finish()
}
}
fn deserialize_optional_u64_or_string<'de, D>(
deserializer: D,
) -> std::result::Result<Option<u64>, D::Error>
where
D: serde::Deserializer<'de>,
{
use serde::de;
struct U64OrString;
impl<'de> de::Visitor<'de> for U64OrString {
type Value = Option<u64>;
fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
formatter.write_str("an integer, an integer-valued float, a numeric string, or null")
}
fn visit_u64<E: de::Error>(self, v: u64) -> std::result::Result<Self::Value, E> {
Ok(Some(v))
}
fn visit_i64<E: de::Error>(self, v: i64) -> std::result::Result<Self::Value, E> {
if v < 0 {
return Err(E::custom(format!("invalid expires_in value: {v}")));
}
Ok(Some(v as u64))
}
fn visit_f64<E: de::Error>(self, v: f64) -> std::result::Result<Self::Value, E> {
if !v.is_finite() || v < 0.0 || v.fract() != 0.0 || v > u64::MAX as f64 {
return Err(E::custom(format!("invalid expires_in value: {v}")));
}
Ok(Some(v as u64))
}
fn visit_str<E: de::Error>(self, v: &str) -> std::result::Result<Self::Value, E> {
v.parse::<u64>().map(Some).map_err(de::Error::custom)
}
fn visit_none<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
Ok(None)
}
fn visit_unit<E: de::Error>(self) -> std::result::Result<Self::Value, E> {
Ok(None)
}
}
deserializer.deserialize_any(U64OrString)
}
// -- Internal Token State --
struct TokenState {
access_token: Option<String>,
expires_at: Option<Instant>,
}
impl TokenState {
fn new() -> Self {
Self {
access_token: None,
expires_at: None,
}
}
fn is_expired(&self, buffer: Duration) -> bool {
match (self.access_token.as_ref(), self.expires_at) {
(Some(_), Some(expires_at)) => Instant::now() + buffer >= expires_at,
(None, _) => true,
(Some(_), None) => true,
}
}
fn update(&mut self, resp: &TokenResponse) {
self.access_token = Some(resp.access_token.clone());
let expires_in = resp.expires_in.unwrap_or(DEFAULT_TOKEN_TTL_SECS);
self.expires_at = Some(Instant::now() + Duration::from_secs(expires_in));
}
}
#[async_trait]
trait TokenSource: Send + Sync + std::fmt::Debug {
async fn fetch_token(&self) -> Result<TokenResponse>;
}
struct ClientCredentialsSource {
issuer_url: String,
client_id: String,
client_secret: String,
scopes: Vec<String>,
http_client: Client,
discovery: RwLock<Option<OidcDiscovery>>,
}
impl std::fmt::Debug for ClientCredentialsSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ClientCredentialsSource")
.field("issuer_url", &self.issuer_url)
.field("client_id", &self.client_id)
.field("client_secret", &"<redacted>")
.field("scopes", &self.scopes)
.finish()
}
}
impl ClientCredentialsSource {
fn new(
issuer_url: String,
client_id: String,
client_secret: Option<String>,
scopes: Vec<String>,
) -> Result<Self> {
let client_secret = client_secret.ok_or(Error::InvalidInput {
message: "client_secret is required for ClientCredentials flow".to_string(),
})?;
Self::validate_issuer_transport(&issuer_url)?;
let http_client = Client::builder()
.timeout(Duration::from_secs(30))
.build()
.map_err(|e| Error::Runtime {
message: format!("Failed to create HTTP client for OAuth: {e}"),
})?;
Ok(Self {
issuer_url,
client_id,
client_secret,
scopes,
http_client,
discovery: RwLock::new(None),
})
}
fn validate_issuer_transport(issuer_url: &str) -> Result<()> {
let issuer = url::Url::parse(issuer_url).map_err(|e| Error::InvalidInput {
message: format!("Invalid OAuth issuer_url: {e}"),
})?;
match issuer.scheme() {
"https" => Ok(()),
"http" if Self::is_loopback_issuer(&issuer) => Ok(()),
_ => Err(Error::InvalidInput {
message:
"ClientCredentials OAuth issuer_url must use https, except for loopback hosts"
.to_string(),
}),
}
}
fn is_loopback_issuer(issuer: &url::Url) -> bool {
let Some(host) = issuer.host_str() else {
return false;
};
host.eq_ignore_ascii_case("localhost")
|| host
.parse::<IpAddr>()
.map(|addr| addr.is_loopback())
.unwrap_or(false)
}
async fn get_discovery(&self) -> Result<OidcDiscovery> {
{
let cached = self.discovery.read().await;
if let Some(ref disc) = *cached {
return Ok(disc.clone());
}
}
let mut cache = self.discovery.write().await;
// Double-check
if let Some(ref disc) = *cache {
return Ok(disc.clone());
}
let discovery_url = format!(
"{}/.well-known/openid-configuration",
self.issuer_url.trim_end_matches('/')
);
debug!("Fetching OIDC discovery from {}", discovery_url);
let resp = self
.http_client
.get(&discovery_url)
.send()
.await
.map_err(|e| Error::Runtime {
message: format!("Failed to fetch OIDC discovery document: {e}"),
})?;
if !resp.status().is_success() {
return Err(Error::Runtime {
message: format!(
"OIDC discovery failed with status {}: {}",
resp.status(),
resp.text().await.unwrap_or_default()
),
});
}
let disc: OidcDiscovery = resp.json().await.map_err(|e| Error::Runtime {
message: format!("Failed to parse OIDC discovery document: {e}"),
})?;
let result = disc.clone();
*cache = Some(disc);
Ok(result)
}
async fn get_token_endpoint(&self) -> Result<String> {
self.get_discovery().await.map(|disc| disc.token_endpoint)
}
fn scopes_string(&self) -> String {
self.scopes.join(" ")
}
async fn post_token_request(
&self,
endpoint: &str,
params: &[(&str, &str)],
) -> Result<TokenResponse> {
let resp = self
.http_client
.post(endpoint)
.form(params)
.send()
.await
.map_err(|e| Error::Runtime {
message: format!("Token request to {endpoint} failed: {e}"),
})?;
if !resp.status().is_success() {
return Err(Error::Runtime {
message: format!(
"Token request failed with status {}: {}",
resp.status(),
resp.text().await.unwrap_or_default()
),
});
}
resp.json().await.map_err(|e| Error::Runtime {
message: format!("Failed to parse token response: {e}"),
})
}
}
#[async_trait]
impl TokenSource for ClientCredentialsSource {
async fn fetch_token(&self) -> Result<TokenResponse> {
let token_endpoint = self.get_token_endpoint().await?;
let scope = self.scopes_string();
let params = [
("grant_type", "client_credentials"),
("client_id", self.client_id.as_str()),
("client_secret", self.client_secret.as_str()),
("scope", scope.as_str()),
];
self.post_token_request(&token_endpoint, &params).await
}
}
struct AzureImdsSource {
client_id: Option<String>,
resource: String,
http_client: Client,
}
impl std::fmt::Debug for AzureImdsSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AzureImdsSource")
.field("client_id", &self.client_id)
.field("resource", &self.resource)
.finish()
}
}
impl AzureImdsSource {
fn new(scopes: Vec<String>, client_id: Option<String>) -> Result<Self> {
let resource = Self::resource_from_scopes(&scopes)?;
let http_client = Client::builder()
.timeout(Duration::from_secs(30))
.no_proxy()
.build()
.map_err(|e| Error::Runtime {
message: format!("Failed to create HTTP client for Azure IMDS OAuth: {e}"),
})?;
Ok(Self {
client_id,
resource,
http_client,
})
}
fn resource_from_scopes(scopes: &[String]) -> Result<String> {
let [scope] = scopes else {
return Err(Error::InvalidInput {
message: "AzureManagedIdentity flow requires exactly one OAuth scope or resource"
.to_string(),
});
};
Ok(scope.strip_suffix("/.default").unwrap_or(scope).to_string())
}
}
#[async_trait]
impl TokenSource for AzureImdsSource {
async fn fetch_token(&self) -> Result<TokenResponse> {
let mut url = format!(
"{AZURE_IMDS_ENDPOINT}?api-version={AZURE_IMDS_API_VERSION}&resource={}",
urlencoding::encode(&self.resource),
);
if let Some(cid) = self.client_id.as_deref() {
url.push_str(&format!("&client_id={}", urlencoding::encode(cid)));
}
let resp = self
.http_client
.get(&url)
.header("Metadata", "true")
.send()
.await
.map_err(|e| Error::Runtime {
message: format!("Azure IMDS request failed: {e}"),
})?;
if !resp.status().is_success() {
return Err(Error::Runtime {
message: format!(
"Azure IMDS returned status {}: {}",
resp.status(),
resp.text().await.unwrap_or_default()
),
});
}
resp.json().await.map_err(|e| Error::Runtime {
message: format!("Failed to parse IMDS token response: {e}"),
})
}
}
/// OAuth header provider that manages the full token lifecycle.
///
/// Implements [`HeaderProvider`] to inject `Authorization: Bearer <token>`
/// headers into every LanceDB request, with automatic token refresh.
pub struct OAuthHeaderProvider {
token_source: Box<dyn TokenSource>,
token_state: Arc<RwLock<TokenState>>,
refresh_buffer: Duration,
}
impl std::fmt::Debug for OAuthHeaderProvider {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("OAuthHeaderProvider")
.field("token_source", &self.token_source)
.finish()
}
}
impl OAuthHeaderProvider {
/// Create a new OAuth header provider from configuration.
pub fn new(config: OAuthConfig) -> Result<Self> {
let OAuthConfig {
issuer_url,
client_id,
client_secret,
scopes,
flow,
refresh_buffer_secs,
} = config;
if scopes.is_empty() {
return Err(Error::InvalidInput {
message: "At least one OAuth scope is required".to_string(),
});
}
let refresh_buffer =
Duration::from_secs(refresh_buffer_secs.unwrap_or(DEFAULT_REFRESH_BUFFER_SECS));
let token_source: Box<dyn TokenSource> = match flow {
OAuthFlow::ClientCredentials => Box::new(ClientCredentialsSource::new(
issuer_url,
client_id,
client_secret,
scopes,
)?),
OAuthFlow::AzureManagedIdentity { client_id } => {
Box::new(AzureImdsSource::new(scopes, client_id)?)
}
};
Ok(Self {
token_source,
token_state: Arc::new(RwLock::new(TokenState::new())),
refresh_buffer,
})
}
/// Get a valid access token, refreshing if necessary.
async fn get_valid_token(&self) -> Result<String> {
// Fast path: check if current token is still valid
{
let state = self.token_state.read().await;
if !state.is_expired(self.refresh_buffer)
&& let Some(ref token) = state.access_token
{
return Ok(token.clone());
}
}
// Slow path: acquire or refresh token
let mut state = self.token_state.write().await;
// Double-check after acquiring write lock
if !state.is_expired(self.refresh_buffer)
&& let Some(ref token) = state.access_token
{
return Ok(token.clone());
}
debug!("Acquiring new OAuth token via {:?}", self.token_source);
let resp = self.token_source.fetch_token().await?;
state.update(&resp);
Ok(resp.access_token)
}
}
#[async_trait]
impl HeaderProvider for OAuthHeaderProvider {
async fn get_headers(&self) -> Result<HashMap<String, String>> {
let token = self.get_valid_token().await?;
Ok(HashMap::from([(
"authorization".to_string(),
format!("Bearer {token}"),
)]))
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::task::JoinHandle;
#[test]
fn test_token_state_expiry() {
let mut state = TokenState::new();
assert!(state.is_expired(Duration::from_secs(0)));
state.access_token = Some("tok".to_string());
state.expires_at = Some(Instant::now() + Duration::from_secs(600));
assert!(!state.is_expired(Duration::from_secs(300)));
assert!(state.is_expired(Duration::from_secs(601)));
state.expires_at = None;
assert!(state.is_expired(Duration::from_secs(0)));
}
#[test]
fn test_token_state_uses_default_expiry() {
let mut state = TokenState::new();
let response = TokenResponse {
access_token: "tok".to_string(),
expires_in: None,
token_type: None,
};
state.update(&response);
assert!(!state.is_expired(Duration::from_secs(DEFAULT_TOKEN_TTL_SECS - 1)));
assert!(state.is_expired(Duration::from_secs(DEFAULT_TOKEN_TTL_SECS + 1)));
}
#[test]
fn test_token_response_accepts_float_expires_in() {
let response: TokenResponse =
serde_json::from_str(r#"{"access_token":"tok","expires_in":3600.0}"#).unwrap();
assert_eq!(response.expires_in, Some(3600));
}
#[test]
fn test_token_response_rejects_negative_expires_in() {
let err =
serde_json::from_str::<TokenResponse>(r#"{"access_token":"tok","expires_in":-1}"#)
.unwrap_err();
assert!(err.to_string().contains("invalid expires_in value: -1"));
}
#[test]
fn test_token_response_debug_redacts_access_token() {
let response = TokenResponse {
access_token: "secret-token".to_string(),
expires_in: Some(3600),
token_type: Some("Bearer".to_string()),
};
let debug = format!("{response:?}");
assert!(!debug.contains("secret-token"));
assert!(debug.contains("access_token: \"<redacted>\""));
}
#[test]
fn test_scopes_string() {
let source = ClientCredentialsSource::new(
"https://login.microsoftonline.com/tenant/v2.0".to_string(),
"app-id".to_string(),
Some("secret".to_string()),
vec!["scope1".to_string(), "scope2".to_string()],
)
.unwrap();
assert_eq!(source.scopes_string(), "scope1 scope2");
}
#[test]
fn test_oauth_config_debug_redacts_client_secret() {
let config = OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: Some("super-secret".to_string()),
scopes: vec!["scope".to_string()],
flow: OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
let debug = format!("{config:?}");
assert!(!debug.contains("super-secret"));
assert!(debug.contains("client_secret: Some(\"<redacted>\")"));
}
#[test]
fn test_oauth_header_provider_debug_redacts_client_secret() {
let config = OAuthConfig {
issuer_url: "https://issuer.example.com".to_string(),
client_id: "client-id".to_string(),
client_secret: Some("super-secret".to_string()),
scopes: vec!["scope".to_string()],
flow: OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
let provider = OAuthHeaderProvider::new(config).unwrap();
let debug = format!("{provider:?}");
assert!(!debug.contains("super-secret"));
assert!(debug.contains("client_secret: \"<redacted>\""));
}
#[test]
fn test_managed_identity_resource_from_default_scope() {
assert_eq!(
AzureImdsSource::resource_from_scopes(&["api://test/.default".to_string()]).unwrap(),
"api://test"
);
}
#[test]
fn test_managed_identity_resource_without_default_suffix() {
assert_eq!(
AzureImdsSource::resource_from_scopes(&["api://test".to_string()]).unwrap(),
"api://test"
);
}
#[test]
fn test_managed_identity_rejects_multiple_scopes() {
let config = OAuthConfig {
issuer_url: "https://login.microsoftonline.com/tenant/v2.0".to_string(),
client_id: "app-id".to_string(),
client_secret: None,
scopes: vec![
"api://test-a/.default".to_string(),
"api://test-b/.default".to_string(),
],
flow: OAuthFlow::AzureManagedIdentity { client_id: None },
refresh_buffer_secs: None,
};
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[tokio::test]
async fn test_token_endpoint_requires_discovery_success() {
let (issuer_url, server) = spawn_discovery_error_server().await;
let source = ClientCredentialsSource::new(
issuer_url,
"client-id".to_string(),
Some("secret".to_string()),
vec!["scope".to_string()],
)
.unwrap();
let err = source.get_token_endpoint().await.unwrap_err();
assert!(matches!(
err,
Error::Runtime { message }
if message.contains("OIDC discovery failed with status 503")
));
server.await.unwrap();
}
#[test]
fn test_client_credentials_requires_secret() {
let config = OAuthConfig {
issuer_url: "https://login.microsoftonline.com/tenant/v2.0".to_string(),
client_id: "app-id".to_string(),
client_secret: None,
scopes: vec!["scope".to_string()],
flow: OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[test]
fn test_client_credentials_rejects_insecure_non_loopback_issuer() {
let config = OAuthConfig {
issuer_url: "http://issuer.example.com".to_string(),
client_id: "app-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: OAuthFlow::ClientCredentials,
refresh_buffer_secs: None,
};
let err = OAuthHeaderProvider::new(config).unwrap_err();
assert!(matches!(
err,
Error::InvalidInput { message }
if message == "ClientCredentials OAuth issuer_url must use https, except for loopback hosts"
));
}
#[test]
fn test_empty_scopes_rejected() {
let config = OAuthConfig {
issuer_url: "https://login.microsoftonline.com/tenant/v2.0".to_string(),
client_id: "app-id".to_string(),
client_secret: None,
scopes: vec![],
flow: OAuthFlow::AzureManagedIdentity { client_id: None },
refresh_buffer_secs: None,
};
assert!(OAuthHeaderProvider::new(config).is_err());
}
#[tokio::test]
async fn test_client_credentials_token_lifecycle() {
let (issuer_url, token_requests, server) = spawn_oauth_server().await;
let config = OAuthConfig {
issuer_url,
client_id: "client-id".to_string(),
client_secret: Some("secret".to_string()),
scopes: vec!["scope".to_string()],
flow: OAuthFlow::ClientCredentials,
refresh_buffer_secs: Some(0),
};
let provider = OAuthHeaderProvider::new(config).unwrap();
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-1");
assert_eq!(token_requests.load(Ordering::SeqCst), 1);
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-1");
assert_eq!(token_requests.load(Ordering::SeqCst), 1);
provider.token_state.write().await.expires_at =
Some(Instant::now() - Duration::from_secs(1));
let headers = provider.get_headers().await.unwrap();
assert_eq!(headers.get("authorization").unwrap(), "Bearer token-2");
assert_eq!(token_requests.load(Ordering::SeqCst), 2);
server.await.unwrap();
}
async fn spawn_oauth_server() -> (String, Arc<AtomicUsize>, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let token_requests = Arc::new(AtomicUsize::new(0));
let server_token_requests = Arc::clone(&token_requests);
let server = tokio::spawn(async move {
for _ in 0..3 {
let (mut stream, _) = listener.accept().await.unwrap();
let (request_line, body) = read_http_request(&mut stream).await;
if request_line.starts_with("GET /.well-known/openid-configuration ") {
let discovery = format!(r#"{{"token_endpoint":"http://{addr}/token"}}"#);
write_json_response(&mut stream, "200 OK", &discovery).await;
} else if request_line.starts_with("POST /token ") {
assert!(body.contains("grant_type=client_credentials"));
assert!(body.contains("client_id=client-id"));
assert!(body.contains("client_secret=secret"));
assert!(body.contains("scope=scope"));
let token_num = server_token_requests.fetch_add(1, Ordering::SeqCst) + 1;
let token = format!(
r#"{{"access_token":"token-{token_num}","expires_in":3600,"token_type":"Bearer"}}"#
);
write_json_response(&mut stream, "200 OK", &token).await;
} else {
write_json_response(&mut stream, "404 Not Found", "{}").await;
}
}
});
(issuer_url, token_requests, server)
}
async fn spawn_discovery_error_server() -> (String, JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let issuer_url = format!("http://{addr}");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
let (request_line, _) = read_http_request(&mut stream).await;
assert!(request_line.starts_with("GET /.well-known/openid-configuration "));
write_json_response(&mut stream, "503 Service Unavailable", "{}").await;
});
(issuer_url, server)
}
async fn read_http_request(stream: &mut TcpStream) -> (String, String) {
let mut buffer = Vec::new();
let mut header_end = None;
while header_end.is_none() {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request headers");
buffer.extend_from_slice(&chunk[..read]);
header_end = find_subsequence(&buffer, b"\r\n\r\n").map(|pos| pos + 4);
}
let header_end = header_end.unwrap();
let headers = String::from_utf8_lossy(&buffer[..header_end]).to_string();
let request_line = headers.lines().next().unwrap_or_default().to_string();
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while buffer.len() < header_end + content_length {
let mut chunk = [0; 1024];
let read = stream.read(&mut chunk).await.unwrap();
assert_ne!(read, 0, "connection closed before request body");
buffer.extend_from_slice(&chunk[..read]);
}
let body =
String::from_utf8_lossy(&buffer[header_end..header_end + content_length]).to_string();
(request_line, body)
}
fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
haystack
.windows(needle.len())
.position(|window| window == needle)
}
async fn write_json_response(stream: &mut TcpStream, status: &str, body: &str) {
let response = format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
stream.write_all(response.as_bytes()).await.unwrap();
}
}
+3 -150
View File
@@ -70,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
@@ -113,7 +102,6 @@ struct FreshnessState {
struct FreshnessHeaders {
min_version: Option<u64>,
min_timestamp: Option<SystemTime>,
min_read_version: Option<u64>,
}
impl FreshnessHeaders {
@@ -125,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
}
}
@@ -612,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));
@@ -900,7 +884,6 @@ impl<S: HttpSend> RemoteTable<S> {
self.client.read_consistency_interval,
SystemTime::now(),
),
min_read_version: state.min_read_version,
}
}
@@ -922,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,
@@ -969,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);
@@ -1587,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
@@ -1848,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 {
@@ -7168,7 +7124,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));
@@ -7193,7 +7148,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),
@@ -7205,7 +7159,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),
@@ -7350,106 +7303,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)]
+13 -98
View File
@@ -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)
@@ -590,45 +579,24 @@ fn array_to_f32_vec(arr: &Arc<dyn arrow_array::Array>) -> Result<Vec<f32>> {
})
}
/// Magic bytes that prefix (and suffix) the Arrow IPC *file* format.
const ARROW_IPC_FILE_MAGIC: &[u8] = b"ARROW1";
/// Parse Arrow IPC response from the namespace server.
///
/// The server may return either the Arrow IPC *file* format or the *stream*
/// format. REST/phalanx returns the file format (it begins with the `ARROW1`
/// magic); reading that with a `StreamReader` fails with "failed to fill whole
/// buffer". Detect the magic and pick the matching reader so both are handled.
async fn parse_arrow_ipc_response(bytes: bytes::Bytes) -> Result<DatasetRecordBatchStream> {
use arrow_ipc::reader::{FileReader, StreamReader};
use arrow_ipc::reader::StreamReader;
use std::io::Cursor;
let (schema, batches) = if bytes.starts_with(ARROW_IPC_FILE_MAGIC) {
let reader = FileReader::try_new(Cursor::new(bytes), None).map_err(|e| Error::Runtime {
message: format!("Failed to parse Arrow IPC file response: {}", e),
let cursor = Cursor::new(bytes);
let reader = StreamReader::try_new(cursor, None).map_err(|e| Error::Runtime {
message: format!("Failed to parse Arrow IPC response: {}", e),
})?;
// Collect all record batches
let schema = reader.schema();
let batches: Vec<_> = reader
.into_iter()
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|e| Error::Runtime {
message: format!("Failed to read Arrow IPC batches: {}", e),
})?;
let schema = reader.schema();
let batches = reader
.into_iter()
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|e| Error::Runtime {
message: format!("Failed to read Arrow IPC file batches: {}", e),
})?;
(schema, batches)
} else {
let reader =
StreamReader::try_new(Cursor::new(bytes), None).map_err(|e| Error::Runtime {
message: format!("Failed to parse Arrow IPC response: {}", e),
})?;
let schema = reader.schema();
let batches = reader
.into_iter()
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|e| Error::Runtime {
message: format!("Failed to read Arrow IPC batches: {}", e),
})?;
(schema, batches)
};
// Create a stream from the batches
let stream = futures::stream::iter(batches.into_iter().map(Ok));
@@ -656,59 +624,6 @@ mod tests {
FixedSizeListArray::try_new_from_values(Float32Array::from(values), dimension).unwrap()
}
#[tokio::test]
async fn test_parse_arrow_ipc_response_handles_file_and_stream() {
use arrow_array::{Int32Array, RecordBatch};
use arrow_ipc::writer::{FileWriter, StreamWriter};
use arrow_schema::{DataType, Field, Schema};
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef],
)
.unwrap();
// Arrow IPC *file* format -- what REST/phalanx returns. Previously this
// failed with "failed to fill whole buffer" because we used a StreamReader.
let mut file_buf = Vec::new();
{
let mut writer = FileWriter::try_new(&mut file_buf, &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
}
assert!(file_buf.starts_with(ARROW_IPC_FILE_MAGIC));
let rows: usize = parse_arrow_ipc_response(bytes::Bytes::from(file_buf))
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap()
.iter()
.map(|b| b.num_rows())
.sum();
assert_eq!(rows, 3);
// Arrow IPC *stream* format must still parse.
let mut stream_buf = Vec::new();
{
let mut writer = StreamWriter::try_new(&mut stream_buf, &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
}
assert!(!stream_buf.starts_with(ARROW_IPC_FILE_MAGIC));
let rows: usize = parse_arrow_ipc_response(bytes::Bytes::from(stream_buf))
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap()
.iter()
.map(|b| b.num_rows())
.sum();
assert_eq!(rows, 3);
}
#[test]
fn test_convert_to_namespace_query_vector() {
let query_vector = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0]));