mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-28 17:08:43 +00:00
Compare commits
45 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| a87cada90e | |||
| 0559108fa9 | |||
| 6ab3b9eb30 | |||
| c94d9a2a16 | |||
| 6c8aa22704 | |||
| 84f46df876 | |||
| 83cff3ab93 | |||
| b85776c22a | |||
| 9d3962686e | |||
| 25645d82d4 | |||
| 0dd9dfdfc7 | |||
| d24b2dcacc | |||
| 2deccf21cf | |||
| ead4d27bfc | |||
| 5153e5a023 | |||
| 79f626b09e | |||
| ae81d73563 | |||
| 8b7e13b0c6 | |||
| b78f2a5044 | |||
| 06872463cf | |||
| 2fbf6d6211 | |||
| 391cac9034 | |||
| 21530432a0 | |||
| 9b825c5f29 | |||
| 8083232dd5 | |||
| 302b21aa94 | |||
| 35b5d015ac | |||
| a57fb68891 | |||
| a614400755 | |||
| 1d880f11ff | |||
| ec4ad54ba2 | |||
| d0bcc6c6fe | |||
| 81c3f108ce | |||
| c988e4848d | |||
| 2fea7cd48d | |||
| 0e65123bd8 | |||
| 6ed3074d4c | |||
| c1a8c3f089 | |||
| fce45ba9fc | |||
| 5013c176dd | |||
| 71f85a8d9f | |||
| c72f5b2960 | |||
| 93f47b8aab | |||
| 105fd73bc6 | |||
| 94d484f539 |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.38.0-beta.6"
|
||||
current_version = "0.38.0-beta.12"
|
||||
parse = """(?x)
|
||||
(?P<major>0|[1-9]\\d*)\\.
|
||||
(?P<minor>0|[1-9]\\d*)\\.
|
||||
|
||||
@@ -44,3 +44,27 @@ updates:
|
||||
python-deps:
|
||||
patterns:
|
||||
- "*"
|
||||
|
||||
# The npm ecosystem covers pnpm lockfiles. There are two separate installs:
|
||||
# the bindings themselves and the examples, which have their own lockfile.
|
||||
# As with cargo and pip above, only bump the lockfile — the version ranges
|
||||
# in package.json are our consumers' constraints, not ours.
|
||||
- package-ecosystem: npm
|
||||
directory: /nodejs
|
||||
schedule:
|
||||
interval: weekly
|
||||
versioning-strategy: lockfile-only
|
||||
groups:
|
||||
nodejs-deps:
|
||||
patterns:
|
||||
- "*"
|
||||
|
||||
- package-ecosystem: npm
|
||||
directory: /nodejs/examples
|
||||
schedule:
|
||||
interval: weekly
|
||||
versioning-strategy: lockfile-only
|
||||
groups:
|
||||
nodejs-examples-deps:
|
||||
patterns:
|
||||
- "*"
|
||||
|
||||
@@ -29,12 +29,14 @@ jobs:
|
||||
steps:
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: "18"
|
||||
node-version: "24"
|
||||
- uses: pnpm/action-setup@v6
|
||||
with:
|
||||
version: 11.1.1
|
||||
# These rules are disabled because Github will always ensure there
|
||||
# is a blank line between the title and the body and Github will
|
||||
# word wrap the description field to ensure a reasonable max line
|
||||
# length.
|
||||
- run: npm install @commitlint/config-conventional
|
||||
- run: >
|
||||
echo 'module.exports = {
|
||||
"rules": {
|
||||
@@ -43,7 +45,11 @@ jobs:
|
||||
"body-leading-blank": [0, "always"]
|
||||
}
|
||||
}' > .commitlintrc.js
|
||||
- run: npx commitlint --extends @commitlint/config-conventional --verbose <<< $COMMIT_MSG
|
||||
- run: >
|
||||
pnpm dlx
|
||||
--package @commitlint/cli@21.2.2
|
||||
--package @commitlint/config-conventional@21.2.2
|
||||
commitlint --extends @commitlint/config-conventional --verbose <<< $COMMIT_MSG
|
||||
env:
|
||||
COMMIT_MSG: >
|
||||
${{ github.event.pull_request.title }}
|
||||
@@ -54,7 +60,7 @@ jobs:
|
||||
with:
|
||||
script: |
|
||||
const message = `**ACTION NEEDED**
|
||||
|
||||
|
||||
Lance follows the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/) for release automation.
|
||||
|
||||
The PR title and description are used as the merge commit message.\
|
||||
|
||||
@@ -56,7 +56,7 @@ jobs:
|
||||
uses: lycheeverse/lychee-action@e7477775783ea5526144ba13e8db5eec57747ce8 # v2.9.0
|
||||
with:
|
||||
# Restricted to http(s) on purpose. Much of docs/src is generated
|
||||
# API reference (the js/ tree comes from `npm run docs` in nodejs)
|
||||
# API reference (the js/ tree comes from `pnpm run docs` in nodejs)
|
||||
# and the hand-written pages use mkdocstrings cross-references and
|
||||
# nav-relative paths that only resolve in the site mkdocs builds,
|
||||
# not in this checkout, so relative links would be reported as
|
||||
|
||||
@@ -55,9 +55,7 @@ jobs:
|
||||
- name: Set up node
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
cache: 'npm'
|
||||
cache-dependency-path: docs/package-lock.json
|
||||
node-version: 24
|
||||
- name: Install node dependencies
|
||||
working-directory: nodejs
|
||||
run: |
|
||||
|
||||
@@ -47,9 +47,8 @@ jobs:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October. The library itself still supports Node >= 18
|
||||
# (see test matrix below).
|
||||
# Build on a supported LTS; the matrix job below covers every
|
||||
# Node version the library claims to support.
|
||||
node-version: 24
|
||||
cache: 'pnpm'
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
@@ -84,7 +83,7 @@ jobs:
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
node-version: [ "18", "20" ]
|
||||
node-version: [ "22", "24", "26" ]
|
||||
runs-on: "ubuntu-22.04"
|
||||
defaults:
|
||||
run:
|
||||
@@ -101,9 +100,9 @@ jobs:
|
||||
- uses: actions/setup-node@v6
|
||||
name: Setup Node.js 24 for build
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October. Build/install runs on Node 24; tests run on the
|
||||
# matrix version below using direct jest invocation.
|
||||
# Build and install once on a fixed version so the generated docs
|
||||
# are identical across matrix legs; the tests below then run on each
|
||||
# supported Node version.
|
||||
node-version: 24
|
||||
cache: 'pnpm'
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
@@ -152,9 +151,9 @@ jobs:
|
||||
S3_TEST: "1"
|
||||
# Newer @smithy/core uses dynamic ESM imports.
|
||||
NODE_OPTIONS: "--experimental-vm-modules"
|
||||
# Invoke jest directly because pnpm 11 itself requires Node 22+
|
||||
# while the matrix tests on older Node versions.
|
||||
run: npx jest --verbose
|
||||
# Invoke the installed jest binary directly; the pnpm shim is set up
|
||||
# against the build-phase Node, not the version selected above.
|
||||
run: node_modules/.bin/jest --verbose
|
||||
- name: Test examples
|
||||
working-directory: ./
|
||||
env:
|
||||
@@ -164,7 +163,7 @@ jobs:
|
||||
run: |
|
||||
python ci/mock_openai.py &
|
||||
cd nodejs/examples
|
||||
npx jest --testEnvironment jest-environment-node-single-context --verbose
|
||||
node_modules/.bin/jest --testEnvironment jest-environment-node-single-context --verbose
|
||||
macos:
|
||||
timeout-minutes: 30
|
||||
# macos-15 ships a newer linker; the older macos-14 linker fails to insert
|
||||
@@ -185,8 +184,7 @@ jobs:
|
||||
version: 11.1.1
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
# pnpm 11 requires Node >= 22.13.
|
||||
node-version: 24
|
||||
cache: 'pnpm'
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
|
||||
@@ -40,40 +40,31 @@ jobs:
|
||||
- target: aarch64-apple-darwin
|
||||
host: macos-latest
|
||||
features: fp16kernels
|
||||
# Fat LTO was ~111 of this job's ~113 minutes.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
brew install protobuf
|
||||
# Fat LTO (the workspace default in .cargo/config.toml) is
|
||||
# single-threaded and is the peak-memory step of the build. On
|
||||
# this runner it accounted for ~111 of the job's ~113 minutes,
|
||||
# making it the critical path of the entire publish pipeline.
|
||||
# ThinLTO parallelizes it across the runner's cores, for a few
|
||||
# percent of runtime performance.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-pc-windows-msvc
|
||||
host: windows-2025
|
||||
features: ","
|
||||
# The lower peak also keeps this on the standard 4-core runner.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc ninja nasm
|
||||
tail -n 1000 /c/ProgramData/chocolatey/logs/chocolatey.log
|
||||
# There is an issue where choco doesn't add nasm to the path
|
||||
export PATH="$PATH:/c/Program Files/NASM"
|
||||
nasm -v
|
||||
# See the ThinLTO note on aarch64-apple-darwin above. Keeping
|
||||
# peak memory down is also what lets this run on the standard
|
||||
# 4-core runner: the 8-core larger runner was only needed to
|
||||
# stop fat LTO from OOMing rustc-LLVM.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: aarch64-pc-windows-msvc
|
||||
host: windows-2025
|
||||
features: ","
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
choco install --no-progress protoc
|
||||
rustup target add aarch64-pc-windows-msvc
|
||||
# See the ThinLTO note on aarch64-apple-darwin above.
|
||||
export CARGO_PROFILE_RELEASE_LTO=thin
|
||||
export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16
|
||||
- target: x86_64-unknown-linux-gnu
|
||||
host: ubuntu-latest
|
||||
features: fp16kernels
|
||||
@@ -103,6 +94,14 @@ jobs:
|
||||
# https://github.com/napi-rs/napi-rs/blob/main/debian-aarch64.Dockerfile
|
||||
docker: ghcr.io/napi-rs/napi-rs/nodejs-rust:lts-debian-aarch64
|
||||
features: "fp16kernels"
|
||||
# Fat LTO OOM-killed rustc every nightly; even with lld it peaked
|
||||
# at 31391 MiB of the runner's 32 GiB.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
# arm64 Linux links through GNU `ld` where x86_64 defaults to
|
||||
# `rust-lld`, which is why only arm64 OOM'd. lld cut the largest
|
||||
# linker process 7.0 -> 4.0 GiB (lancedb/sophon#7313).
|
||||
linker: /tmp/aarch64-lld-clang
|
||||
pre_build: |-
|
||||
set -e &&
|
||||
apt-get update &&
|
||||
@@ -112,9 +111,30 @@ jobs:
|
||||
# AT_HWCAP2 (added in Linux 3.17). Define it for aws-lc-sys.
|
||||
export CFLAGS="$CFLAGS -DAT_HWCAP2=26" &&
|
||||
rustup target add aarch64-unknown-linux-gnu
|
||||
# Not `&&`-chained: in dash, errexit does not fire for a
|
||||
# non-final command in an `&&` list, so failures were ignored.
|
||||
#
|
||||
# A wrapper rather than `-C link-arg` because the per-target
|
||||
# rustflags variable does not reach every unit that links, while
|
||||
# the linker variable does. `clang` because GCC silently ignores
|
||||
# `-fuse-ld=lld` unless built with lld support. Two echoes
|
||||
# because printf's newline escape gets rewritten to `;` between
|
||||
# here and the container.
|
||||
echo '#!/bin/sh' > /tmp/aarch64-lld-clang
|
||||
echo 'exec clang --target=aarch64-unknown-linux-gnu --sysroot=/usr/aarch64-unknown-linux-gnu/aarch64-unknown-linux-gnu/sysroot --gcc-toolchain=/usr/aarch64-unknown-linux-gnu -fuse-ld=lld "$@"' >> /tmp/aarch64-lld-clang
|
||||
chmod 0755 /tmp/aarch64-lld-clang
|
||||
# Fail now, not at the cdylib link ~30 minutes later. Linking at
|
||||
# all also proves lld resolved; clang errors out when it cannot.
|
||||
echo 'int main(void){return 0;}' > /tmp/probe.c
|
||||
/tmp/aarch64-lld-clang /tmp/probe.c -o /tmp/probe
|
||||
readelf -h /tmp/probe | grep AArch64
|
||||
- target: aarch64-unknown-linux-musl
|
||||
host: ubuntu-2404-8x-x64
|
||||
features: ","
|
||||
# Fat LTO took the whole runner down. lld cannot help: it died
|
||||
# inside rustc's LLVM, before any linker was spawned.
|
||||
lto: thin
|
||||
codegen_units: 16
|
||||
pre_build: |-
|
||||
set -e &&
|
||||
sudo apt-get update &&
|
||||
@@ -123,6 +143,19 @@ jobs:
|
||||
export EXTRA_ARGS="-x"
|
||||
name: build - ${{ matrix.settings.target }}
|
||||
runs-on: ${{ matrix.settings.host }}
|
||||
# On the job, not exported from `pre_build`: `Swatinem/rust-cache` hashes
|
||||
# `CARGO_*` into its cache key before any step runs, so a step-local export
|
||||
# leaves the key unchanged while cargo still rebuilds cold. The ThinLTO
|
||||
# legs had been doing that every run.
|
||||
#
|
||||
# Not `RUSTFLAGS`: setting it, even to "", discards every config-file
|
||||
# rustflag, silently dropping .cargo/config.toml's `target-cpu` and
|
||||
# `target-feature` from the published binaries.
|
||||
env:
|
||||
CARGO_PROFILE_RELEASE_LTO: ${{ matrix.settings.lto || 'fat' }}
|
||||
CARGO_PROFILE_RELEASE_CODEGEN_UNITS: ${{ matrix.settings.codegen_units || '1' }}
|
||||
# Empty elsewhere: a per-target variable is only read for that triple.
|
||||
CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER: ${{ matrix.settings.linker }}
|
||||
defaults:
|
||||
run:
|
||||
working-directory: nodejs
|
||||
@@ -135,8 +168,7 @@ jobs:
|
||||
- name: Setup node
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
# pnpm 11 requires Node >= 22.13; use 24 since 22 hits EOL
|
||||
# in October.
|
||||
# pnpm 11 requires Node >= 22.13.
|
||||
node-version: 24
|
||||
cache: pnpm
|
||||
cache-dependency-path: nodejs/pnpm-lock.yaml
|
||||
@@ -169,19 +201,15 @@ jobs:
|
||||
# creating ref). The nightly cadence also keeps entries inside
|
||||
# GitHub's 7-day eviction window, which a tag-only trigger would not.
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
# Docker builds can use rust-cache too. `target/` already lives on the
|
||||
# host because the whole workspace is bind-mounted into the container, and
|
||||
# rust-cache's prune and save run host-side, so they can manage it -- which
|
||||
# is what keeps the entry to dependency artifacts rather than a multi-GB
|
||||
# copy of everything.
|
||||
# Docker builds can use rust-cache too: the workspace is bind-mounted, so
|
||||
# `target/` lives on the host and rust-cache's prune keeps the entry
|
||||
# small.
|
||||
#
|
||||
# Two differences from the native builds. The container's CARGO_HOME is
|
||||
# bind-mounted from `.cargo-cache` rather than the host's ~/.cargo, so that
|
||||
# has to be cached explicitly. And the key is derived from the *host* rustc
|
||||
# version, which is not the compiler that produced these artifacts; that is
|
||||
# safe because cargo fingerprints the real compiler and rebuilds on a
|
||||
# mismatch, it just means a base-image toolchain bump costs one cold build
|
||||
# instead of invalidating the key.
|
||||
# bind-mounted from `.cargo-cache` rather than ~/.cargo, so that is cached
|
||||
# explicitly. And the key uses the *host* rustc version, not the compiler
|
||||
# that built these artifacts -- safe, since cargo fingerprints the real
|
||||
# one; a base-image bump just costs one cold build.
|
||||
- name: Cache cargo (docker builds)
|
||||
uses: Swatinem/rust-cache@v2
|
||||
if: ${{ matrix.settings.docker }}
|
||||
@@ -210,14 +238,19 @@ jobs:
|
||||
# cache step above saves. Previously the registry mounts pointed at
|
||||
# `.cargo/...`, a path nothing cached, so the container re-downloaded
|
||||
# the whole crate registry on every run.
|
||||
#
|
||||
# `docker run` inherits nothing; `-e NAME` carries the job's `env:` in.
|
||||
options: "--user 0:0 -v ${{ github.workspace }}/.cargo-cache/git/db:/usr/local/cargo/git/db \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/cache:/usr/local/cargo/registry/cache \
|
||||
-v ${{ github.workspace }}/.cargo-cache/registry/index:/usr/local/cargo/registry/index \
|
||||
-e CARGO_PROFILE_RELEASE_LTO \
|
||||
-e CARGO_PROFILE_RELEASE_CODEGEN_UNITS \
|
||||
-e CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER \
|
||||
-v ${{ github.workspace }}:/build -w /build/nodejs"
|
||||
run: |
|
||||
set -e
|
||||
${{ matrix.settings.pre_build }}
|
||||
npx napi build --platform --release \
|
||||
node_modules/.bin/napi build --platform --release \
|
||||
--features ${{ matrix.settings.features }} \
|
||||
--target ${{ matrix.settings.target }} \
|
||||
--dts ../lancedb/native.d.ts \
|
||||
@@ -237,7 +270,7 @@ jobs:
|
||||
- name: Build
|
||||
run: |
|
||||
${{ matrix.settings.pre_build }}
|
||||
npx napi build --platform --release \
|
||||
node_modules/.bin/napi build --platform --release \
|
||||
--features ${{ matrix.settings.features }} \
|
||||
--target ${{ matrix.settings.target }} \
|
||||
--dts ../lancedb/native.d.ts \
|
||||
@@ -256,6 +289,18 @@ jobs:
|
||||
if: always()
|
||||
run: df -h
|
||||
shell: bash
|
||||
- name: Report peak memory
|
||||
if: always() && runner.os == 'Linux'
|
||||
shell: bash
|
||||
run: |
|
||||
peak=$(find /sys/fs/cgroup -name memory.peak -readable \
|
||||
-exec cat {} + 2>/dev/null | sort -n | tail -1)
|
||||
if [ -n "$peak" ]; then
|
||||
echo "peak memory: $((peak / 1024 / 1024)) MiB"
|
||||
else
|
||||
echo "peak memory: unavailable (no readable cgroup v2 memory.peak)"
|
||||
fi
|
||||
free -g || true
|
||||
- name: Upload artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
@@ -293,7 +338,7 @@ jobs:
|
||||
- target: aarch64-unknown-linux-gnu
|
||||
host: ubuntu-2404-8x-arm64
|
||||
node:
|
||||
- '20'
|
||||
- '22'
|
||||
runs-on: ${{ matrix.settings.host }}
|
||||
defaults:
|
||||
run:
|
||||
@@ -339,9 +384,9 @@ jobs:
|
||||
- name: Move built files
|
||||
run: cp dist/native.d.ts dist/native.js dist/*.node lancedb/
|
||||
- name: Test bindings
|
||||
# Invoke jest directly because pnpm 11 itself requires Node 22+
|
||||
# while the matrix tests on older Node versions.
|
||||
run: npx jest --verbose
|
||||
# Invoke the installed jest binary directly; the pnpm shim is set up
|
||||
# against the install-phase Node, not the version selected above.
|
||||
run: node_modules/.bin/jest --verbose
|
||||
publish:
|
||||
name: Publish
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
name: Update package-lock.json
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: ./.github/workflows/update_package_lock
|
||||
with:
|
||||
github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
|
||||
@@ -1,22 +0,0 @@
|
||||
name: Update NodeJs package-lock.json
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout
|
||||
uses: actions/checkout@v6
|
||||
with:
|
||||
ref: main
|
||||
persist-credentials: false
|
||||
fetch-depth: 0
|
||||
lfs: true
|
||||
- uses: ./.github/workflows/update_package_lock_nodejs
|
||||
with:
|
||||
github_token: ${{ secrets.LANCEDB_RELEASE_TOKEN }}
|
||||
@@ -20,7 +20,10 @@ repos:
|
||||
hooks:
|
||||
- id: local-biome-check
|
||||
name: biome check
|
||||
entry: npx @biomejs/biome@1.8.3 check --config-path nodejs/biome.json nodejs/
|
||||
# Use the biome from nodejs/package.json rather than a separately
|
||||
# pinned one: the two drifted apart and disagreed on formatting, so
|
||||
# this hook rejected code that `pnpm lint` accepted.
|
||||
entry: nodejs/node_modules/.bin/biome check --config-path nodejs/biome.json nodejs/
|
||||
language: system
|
||||
types: [text]
|
||||
files: "nodejs/.*"
|
||||
|
||||
@@ -38,7 +38,7 @@ Before committing changes, run formatting for every language you touched. At min
|
||||
* Rust changes: run `cargo fmt --all`.
|
||||
* Python changes: run `ruff format .` and `ruff check .` from the repository root,
|
||||
and run targeted tests through `cd python && uv run ...`.
|
||||
* TypeScript changes: run the relevant `npm`/`pnpm` lint, format, build, and docs commands in `nodejs`.
|
||||
* TypeScript changes: run the relevant `pnpm` lint, format, build, and docs commands in `nodejs`.
|
||||
|
||||
Before creating a PR, the exact value passed to `gh pr create --title` must follow
|
||||
Conventional Commits, such as `fix: support nested field paths in native index creation`
|
||||
@@ -101,12 +101,12 @@ Python bindings changes:
|
||||
TypeScript bindings changes:
|
||||
|
||||
1. Add napi-rs method binding on `Table` in `nodejs/src/table.rs`.
|
||||
2. Run `npm run build` to generate TypeScript definitions.
|
||||
2. Run `pnpm build` to generate TypeScript definitions.
|
||||
3. Add typescript method on abstract class `Table` in `nodejs/src/table.ts`.
|
||||
4. Add concrete method on `LocalTable` class in `nodejs/src/native_table.ts`.
|
||||
* Note: despite the name, this class is also used for remote tables.
|
||||
5. Add test in `nodejs/__test__/table.test.ts`.
|
||||
6. Run `npm run docs` to generate TypeScript documentation.
|
||||
6. Run `pnpm run docs` to generate TypeScript documentation.
|
||||
|
||||
## Python API reference
|
||||
|
||||
|
||||
Generated
+47
-47
@@ -1597,9 +1597,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chacha20"
|
||||
version = "0.10.0"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601"
|
||||
checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06"
|
||||
dependencies = [
|
||||
"cfg-if 1.0.4",
|
||||
"cpufeatures 0.3.0",
|
||||
@@ -3455,8 +3455,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
|
||||
|
||||
[[package]]
|
||||
name = "fsst"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"rand 0.9.5",
|
||||
@@ -4815,8 +4815,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
|
||||
|
||||
[[package]]
|
||||
name = "lance"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -4888,8 +4888,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-arrow"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4911,7 +4911,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-scalar"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4925,7 +4925,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "lance-arrow-stats"
|
||||
version = "58.0.0"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -4934,8 +4934,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-bitpacking"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrayref",
|
||||
"crunchy",
|
||||
@@ -4945,8 +4945,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-core"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -4983,8 +4983,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datafusion"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5013,8 +5013,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-datagen"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5031,8 +5031,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-derive"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
@@ -5041,8 +5041,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-encoding"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5075,8 +5075,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-file"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-arith",
|
||||
"arrow-array",
|
||||
@@ -5107,8 +5107,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arrow",
|
||||
@@ -5172,8 +5172,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-index-core"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5195,8 +5195,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-io"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5236,8 +5236,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-linalg"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5251,8 +5251,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
@@ -5264,8 +5264,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-impls"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-ipc",
|
||||
@@ -5318,8 +5318,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-select"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5333,8 +5333,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-table"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"arrow-array",
|
||||
@@ -5374,8 +5374,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-testing"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-schema",
|
||||
@@ -5388,8 +5388,8 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lance-tokenizer"
|
||||
version = "11.0.0-rc.1"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v11.0.0-rc.1#308c00db4f4dae6b2ea4a1d1d2ceb1c3ed88a159"
|
||||
version = "12.0.0-beta.2"
|
||||
source = "git+https://github.com/lance-format/lance.git?tag=v12.0.0-beta.2#dafa4642658d996b3e31dde91e02f72db7860d7e"
|
||||
dependencies = [
|
||||
"frostem",
|
||||
"icu_segmenter",
|
||||
@@ -5402,7 +5402,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"anyhow",
|
||||
@@ -5490,7 +5490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-nodejs"
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"arrow-buffer",
|
||||
@@ -5515,7 +5515,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
dependencies = [
|
||||
"arrow",
|
||||
"async-trait",
|
||||
|
||||
+14
-14
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
|
||||
rust-version = "1.91.0"
|
||||
|
||||
[workspace.dependencies]
|
||||
lance = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=11.0.0-rc.1", default-features = false, "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=11.0.0-rc.1", "tag" = "v11.0.0-rc.1", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-core = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datagen = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-file = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-io = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-index = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-linalg = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-namespace-impls = { "version" = "=12.0.0-beta.2", default-features = false, "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-table = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-testing = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-datafusion = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-encoding = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lance-arrow = { "version" = "=12.0.0-beta.2", "tag" = "v12.0.0-beta.2", "git" = "https://github.com/lance-format/lance.git" }
|
||||
lancedb = { path = "rust/lancedb", default-features = false }
|
||||
ahash = "0.8"
|
||||
# Note that this one does not include pyarrow
|
||||
|
||||
@@ -5,5 +5,5 @@ licenses:
|
||||
cd python && cargo about generate ../about.hbs -o RUST_THIRD_PARTY_LICENSES.html -c ../about.toml
|
||||
cd python && uv sync --all-extras && uv tool run pip-licenses --python .venv/bin/python --format=markdown --with-urls --output-file=PYTHON_THIRD_PARTY_LICENSES.md
|
||||
cd nodejs && cargo about generate ../about.hbs -o RUST_THIRD_PARTY_LICENSES.html -c ../about.toml
|
||||
cd nodejs && npx license-checker --markdown --out NODEJS_THIRD_PARTY_LICENSES.md
|
||||
cd nodejs && pnpm dlx license-checker@25 --markdown --out NODEJS_THIRD_PARTY_LICENSES.md
|
||||
cd java && ./mvnw license:aggregate-add-third-party -q
|
||||
|
||||
@@ -12,16 +12,12 @@ done
|
||||
# This updates the lockfile without building
|
||||
cargo metadata --quiet > /dev/null
|
||||
|
||||
pushd nodejs || exit 1
|
||||
npm install --package-lock-only --silent
|
||||
popd
|
||||
|
||||
if git diff --quiet --exit-code; then
|
||||
echo "No lockfile changes to commit; skipping amend."
|
||||
elif $AMEND; then
|
||||
git add Cargo.lock nodejs/package-lock.json
|
||||
git add Cargo.lock
|
||||
git commit --amend --no-edit
|
||||
else
|
||||
git add Cargo.lock nodejs/package-lock.json
|
||||
git add Cargo.lock
|
||||
git commit -m "Update lockfiles"
|
||||
fi
|
||||
|
||||
@@ -131,18 +131,13 @@ allow = [
|
||||
"BSD-3-Clause",
|
||||
"ISC",
|
||||
"Unicode-3.0",
|
||||
"Unicode-DFS-2016",
|
||||
"Zlib",
|
||||
"CC0-1.0",
|
||||
"MPL-2.0",
|
||||
"BSL-1.0",
|
||||
"OpenSSL",
|
||||
# 0BSD ("BSD Zero Clause") is effectively public domain — no attribution
|
||||
# required. Pulled in by `mock_instant`.
|
||||
"0BSD",
|
||||
# bzip2-1.0.6 is the permissive upstream bzip2 license (BSD-like). Pulled
|
||||
# in by `libbz2-rs-sys`, the pure-Rust bzip2 implementation.
|
||||
"bzip2-1.0.6",
|
||||
# CDLA-Permissive-2.0 is a permissive data license used by `webpki-roots`
|
||||
# for the Mozilla CA root bundle. Data-only, distribution-compatible.
|
||||
"CDLA-Permissive-2.0",
|
||||
@@ -150,12 +145,7 @@ allow = [
|
||||
confidence-threshold = 0.8
|
||||
# Per-crate license exceptions: allow a license for a specific crate only,
|
||||
# rather than globally via the `allow` list above.
|
||||
exceptions = [
|
||||
# CDDL-1.0 (copyleft) is pulled in only as a dev/profiling dependency via
|
||||
# `inferno` -> `pprof` -> `lance-testing`; it is a test dependency that we
|
||||
# do not distribute, so scope the allowance to `inferno` alone.
|
||||
{ allow = ["CDDL-1.0"], crate = "inferno" },
|
||||
]
|
||||
exceptions = []
|
||||
# Crates whose license cannot be determined from Cargo metadata but whose
|
||||
# license we've manually confirmed from upstream. Keep this list minimal.
|
||||
[[licenses.clarify]]
|
||||
|
||||
+11
-8
@@ -47,22 +47,24 @@ pytest -vv python/tests/docs
|
||||
|
||||
### Checking typescript examples
|
||||
|
||||
The `@lancedb/lancedb` package must be built before running the tests:
|
||||
The examples depend on `@lancedb/lancedb` at `file:../dist`, so the package must be
|
||||
built before running the tests. This uses pnpm; see the
|
||||
[Typescript contributing guide](../nodejs/CONTRIBUTING.md) for the toolchain setup.
|
||||
|
||||
```shell
|
||||
pushd nodejs
|
||||
npm ci
|
||||
npm run build
|
||||
pnpm install
|
||||
pnpm build
|
||||
popd
|
||||
```
|
||||
|
||||
Then you can run the examples by going to the `nodejs/examples` directory and
|
||||
running the tests like a normal npm package:
|
||||
Then you can run the examples by going to the `nodejs/examples` directory, which is a
|
||||
separate pnpm package with its own lockfile:
|
||||
|
||||
```shell
|
||||
pushd nodejs/examples
|
||||
npm ci
|
||||
npm test
|
||||
pnpm install
|
||||
pnpm test
|
||||
popd
|
||||
```
|
||||
|
||||
@@ -84,6 +86,7 @@ The new files should be checked into the repository.
|
||||
|
||||
```shell
|
||||
pushd nodejs
|
||||
npm run docs
|
||||
# `pnpm docs` would invoke pnpm's built-in `docs` command, not the script.
|
||||
pnpm run docs
|
||||
popd
|
||||
```
|
||||
|
||||
Generated
-135
@@ -1,135 +0,0 @@
|
||||
{
|
||||
"name": "lancedb-docs-test",
|
||||
"version": "1.0.0",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "lancedb-docs-test",
|
||||
"version": "1.0.0",
|
||||
"license": "Apache 2",
|
||||
"dependencies": {
|
||||
"apache-arrow": "file:../node/node_modules/apache-arrow",
|
||||
"vectordb": "file:../node"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^20.11.8",
|
||||
"typescript": "^5.3.3"
|
||||
}
|
||||
},
|
||||
"../node": {
|
||||
"name": "vectordb",
|
||||
"version": "0.21.2-beta.0",
|
||||
"cpu": [
|
||||
"x64",
|
||||
"arm64"
|
||||
],
|
||||
"license": "Apache-2.0",
|
||||
"os": [
|
||||
"darwin",
|
||||
"linux",
|
||||
"win32"
|
||||
],
|
||||
"dependencies": {
|
||||
"@neon-rs/load": "^0.0.74",
|
||||
"axios": "^1.4.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@neon-rs/cli": "^0.0.160",
|
||||
"@types/chai": "^4.3.4",
|
||||
"@types/chai-as-promised": "^7.1.5",
|
||||
"@types/mocha": "^10.0.1",
|
||||
"@types/node": "^18.16.2",
|
||||
"@types/sinon": "^10.0.15",
|
||||
"@types/temp": "^0.9.1",
|
||||
"@types/uuid": "^9.0.3",
|
||||
"@typescript-eslint/eslint-plugin": "^5.59.1",
|
||||
"apache-arrow-old": "npm:apache-arrow@13.0.0",
|
||||
"cargo-cp-artifact": "^0.1",
|
||||
"chai": "^4.3.7",
|
||||
"chai-as-promised": "^7.1.1",
|
||||
"eslint": "^8.39.0",
|
||||
"eslint-config-standard-with-typescript": "^34.0.1",
|
||||
"eslint-plugin-import": "^2.26.0",
|
||||
"eslint-plugin-n": "^15.7.0",
|
||||
"eslint-plugin-promise": "^6.1.1",
|
||||
"mocha": "^10.2.0",
|
||||
"openai": "^4.24.1",
|
||||
"sinon": "^15.1.0",
|
||||
"temp": "^0.9.4",
|
||||
"ts-node": "^10.9.1",
|
||||
"ts-node-dev": "^2.0.0",
|
||||
"typedoc": "^0.24.7",
|
||||
"typedoc-plugin-markdown": "^3.15.3",
|
||||
"typescript": "^5.1.0",
|
||||
"uuid": "^9.0.0"
|
||||
},
|
||||
"optionalDependencies": {
|
||||
"@lancedb/vectordb-darwin-arm64": "0.21.2-beta.0",
|
||||
"@lancedb/vectordb-darwin-x64": "0.21.2-beta.0",
|
||||
"@lancedb/vectordb-linux-arm64-gnu": "0.21.2-beta.0",
|
||||
"@lancedb/vectordb-linux-x64-gnu": "0.21.2-beta.0",
|
||||
"@lancedb/vectordb-win32-x64-msvc": "0.21.2-beta.0"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@apache-arrow/ts": "^14.0.2",
|
||||
"apache-arrow": "^14.0.2"
|
||||
}
|
||||
},
|
||||
"../node/node_modules/apache-arrow": {
|
||||
"version": "14.0.2",
|
||||
"license": "Apache-2.0",
|
||||
"dependencies": {
|
||||
"@types/command-line-args": "5.2.0",
|
||||
"@types/command-line-usage": "5.0.2",
|
||||
"@types/node": "20.3.0",
|
||||
"@types/pad-left": "2.1.1",
|
||||
"command-line-args": "5.2.1",
|
||||
"command-line-usage": "7.0.1",
|
||||
"flatbuffers": "23.5.26",
|
||||
"json-bignum": "^0.0.3",
|
||||
"pad-left": "^2.1.0",
|
||||
"tslib": "^2.5.3"
|
||||
},
|
||||
"bin": {
|
||||
"arrow2csv": "bin/arrow2csv.js"
|
||||
}
|
||||
},
|
||||
"node_modules/@types/node": {
|
||||
"version": "20.11.8",
|
||||
"resolved": "https://registry.npmjs.org/@types/node/-/node-20.11.8.tgz",
|
||||
"integrity": "sha512-i7omyekpPTNdv4Jb/Rgqg0RU8YqLcNsI12quKSDkRXNfx7Wxdm6HhK1awT3xTgEkgxPn3bvnSpiEAc7a7Lpyow==",
|
||||
"dev": true,
|
||||
"dependencies": {
|
||||
"undici-types": "~5.26.4"
|
||||
}
|
||||
},
|
||||
"node_modules/apache-arrow": {
|
||||
"resolved": "../node/node_modules/apache-arrow",
|
||||
"link": true
|
||||
},
|
||||
"node_modules/typescript": {
|
||||
"version": "5.3.3",
|
||||
"resolved": "https://registry.npmjs.org/typescript/-/typescript-5.3.3.tgz",
|
||||
"integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==",
|
||||
"dev": true,
|
||||
"bin": {
|
||||
"tsc": "bin/tsc",
|
||||
"tsserver": "bin/tsserver"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=14.17"
|
||||
}
|
||||
},
|
||||
"node_modules/undici-types": {
|
||||
"version": "5.26.5",
|
||||
"resolved": "https://registry.npmjs.org/undici-types/-/undici-types-5.26.5.tgz",
|
||||
"integrity": "sha512-JlCMO+ehdEIKqlFxk6IfVoAUVmgz7cU7zD/h9XZ0qzeosSHmUJVOzSQvvYSYWXkFXC+IfLKSIffhv0sVZup6pA==",
|
||||
"dev": true
|
||||
},
|
||||
"node_modules/vectordb": {
|
||||
"resolved": "../node",
|
||||
"link": true
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
{
|
||||
"name": "lancedb-docs-test",
|
||||
"version": "1.0.0",
|
||||
"description": "auto-generated tests from doc",
|
||||
"author": "dev@lancedb.com",
|
||||
"license": "Apache 2",
|
||||
"dependencies": {
|
||||
"apache-arrow": "file:../node/node_modules/apache-arrow",
|
||||
"vectordb": "file:../node"
|
||||
},
|
||||
"scripts": {
|
||||
"build": "tsc -b && cd ../node && npm run build-release",
|
||||
"example": "npm run build && node",
|
||||
"test": "npm run build && ls dist/*.js | xargs -n 1 node"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^20.11.8",
|
||||
"typescript": "^5.3.3"
|
||||
}
|
||||
}
|
||||
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
|
||||
<dependency>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-core</artifactId>
|
||||
<version>0.38.0-beta.6</version>
|
||||
<version>0.38.0-beta.12</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,518 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / AutoQuery
|
||||
|
||||
# Class: AutoQuery
|
||||
|
||||
A builder for automatic string searches.
|
||||
|
||||
Automatic search determines whether to use full-text or vector search from
|
||||
the table revision selected for each execution. This builder exposes the
|
||||
common operations supported by both query families.
|
||||
|
||||
## Extends
|
||||
|
||||
- `StandardQueryBase`<`NativeQuery` \| `NativeVectorQuery`>
|
||||
|
||||
## Properties
|
||||
|
||||
### inner
|
||||
|
||||
```ts
|
||||
protected inner: Query | VectorQuery | Promise<Query | VectorQuery>;
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.inner`
|
||||
|
||||
## Methods
|
||||
|
||||
### analyzePlan()
|
||||
|
||||
```ts
|
||||
analyzePlan(distributedMetrics?): Promise<string>
|
||||
```
|
||||
|
||||
Executes the query and returns the physical query plan annotated with runtime metrics.
|
||||
|
||||
This is useful for debugging and performance analysis, as it shows how the query was executed
|
||||
and includes metrics such as elapsed time, rows processed, and I/O statistics.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md)
|
||||
How distributed worker metrics are displayed for remote query plans.
|
||||
Defaults to `"aggregate"`.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
A query execution plan with runtime metrics for each step.
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
import * as lancedb from "@lancedb/lancedb"
|
||||
|
||||
const db = await lancedb.connect("./.lancedb");
|
||||
const table = await db.createTable("my_table", [
|
||||
{ vector: [1.1, 0.9], id: "1" },
|
||||
]);
|
||||
|
||||
const plan = await table.query().nearestTo([0.5, 0.2]).analyzePlan();
|
||||
|
||||
Example output (with runtime metrics inlined):
|
||||
AnalyzeExec verbose=true, metrics=[]
|
||||
ProjectionExec: expr=[id@3 as id, vector@0 as vector, _distance@2 as _distance], metrics=[output_rows=1, elapsed_compute=3.292µs]
|
||||
Take: columns="vector, _rowid, _distance, (id)", metrics=[output_rows=1, elapsed_compute=66.001µs, batches_processed=1, bytes_read=8, iops=1, requests=1]
|
||||
CoalesceBatchesExec: target_batch_size=1024, metrics=[output_rows=1, elapsed_compute=3.333µs]
|
||||
GlobalLimitExec: skip=0, fetch=10, metrics=[output_rows=1, elapsed_compute=167ns]
|
||||
FilterExec: _distance@2 IS NOT NULL, metrics=[output_rows=1, elapsed_compute=8.542µs]
|
||||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST], metrics=[output_rows=1, elapsed_compute=63.25µs, row_replacements=1]
|
||||
KNNVectorDistance: metric=l2, metrics=[output_rows=1, elapsed_compute=114.333µs, output_batches=1]
|
||||
LanceScan: uri=/path/to/data, projection=[vector], row_id=true, row_addr=false, ordered=false, metrics=[output_rows=1, elapsed_compute=103.626µs, bytes_read=549, iops=2, requests=2]
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.analyzePlan`
|
||||
|
||||
***
|
||||
|
||||
### execute()
|
||||
|
||||
```ts
|
||||
protected execute(options?): AsyncGenerator<RecordBatch<any>, void, unknown>
|
||||
```
|
||||
|
||||
Execute the query and return the results as an
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`AsyncGenerator`<`RecordBatch`<`any`>, `void`, `unknown`>
|
||||
|
||||
#### See
|
||||
|
||||
- AsyncIterator
|
||||
of
|
||||
- RecordBatch.
|
||||
|
||||
By default, LanceDb will use many threads to calculate results and, when
|
||||
the result set is large, multiple batches will be processed at one time.
|
||||
This readahead is limited however and backpressure will be applied if this
|
||||
stream is consumed slowly (this constrains the maximum memory used by a
|
||||
single query)
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.execute`
|
||||
|
||||
***
|
||||
|
||||
### explainPlan()
|
||||
|
||||
```ts
|
||||
explainPlan(verbose): Promise<string>
|
||||
```
|
||||
|
||||
Generates an explanation of the query execution plan.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **verbose**: `boolean` = `false`
|
||||
If true, provides a more detailed explanation. Defaults to false.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`string`>
|
||||
|
||||
A Promise that resolves to a string containing the query execution plan explanation.
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
import * as lancedb from "@lancedb/lancedb"
|
||||
const db = await lancedb.connect("./.lancedb");
|
||||
const table = await db.createTable("my_table", [
|
||||
{ vector: [1.1, 0.9], id: "1" },
|
||||
]);
|
||||
const plan = await table.query().nearestTo([0.5, 0.2]).explainPlan();
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.explainPlan`
|
||||
|
||||
***
|
||||
|
||||
### fastSearch()
|
||||
|
||||
```ts
|
||||
fastSearch(): this
|
||||
```
|
||||
|
||||
Skip searching un-indexed data. This can make search faster, but will miss
|
||||
any data that is not yet indexed.
|
||||
|
||||
Use [Table#optimize](Table.md#optimize) to index all un-indexed data.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.fastSearch`
|
||||
|
||||
***
|
||||
|
||||
### ~~filter()~~
|
||||
|
||||
```ts
|
||||
filter(predicate): this
|
||||
```
|
||||
|
||||
A filter statement to be applied to this query.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **predicate**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### See
|
||||
|
||||
where
|
||||
|
||||
#### Deprecated
|
||||
|
||||
Use `where` instead
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.filter`
|
||||
|
||||
***
|
||||
|
||||
### fullTextSearch()
|
||||
|
||||
```ts
|
||||
fullTextSearch(query, options?): this
|
||||
```
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **query**: `string` \| [`FullTextQuery`](../interfaces/FullTextQuery.md)
|
||||
|
||||
* **options?**: `Partial`<[`FullTextSearchOptions`](../interfaces/FullTextSearchOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.fullTextSearch`
|
||||
|
||||
***
|
||||
|
||||
### limit()
|
||||
|
||||
```ts
|
||||
limit(limit): this
|
||||
```
|
||||
|
||||
Set the maximum number of results to return.
|
||||
|
||||
By default, a plain search has no limit. If this method is not
|
||||
called then every valid row from the table will be returned.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **limit**: `number`
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.limit`
|
||||
|
||||
***
|
||||
|
||||
### offset()
|
||||
|
||||
```ts
|
||||
offset(offset): this
|
||||
```
|
||||
|
||||
Set the number of rows to skip before returning results.
|
||||
|
||||
This is useful for pagination.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **offset**: `number`
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.offset`
|
||||
|
||||
***
|
||||
|
||||
### orderBy()
|
||||
|
||||
```ts
|
||||
orderBy(ordering): this
|
||||
```
|
||||
|
||||
Sort the results by the specified column(s).
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **ordering**: [`ColumnOrdering`](../interfaces/ColumnOrdering.md) \| [`ColumnOrdering`](../interfaces/ColumnOrdering.md)[]
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
This query builder.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.orderBy`
|
||||
|
||||
***
|
||||
|
||||
### outputSchema()
|
||||
|
||||
```ts
|
||||
outputSchema(): Promise<Schema<any>>
|
||||
```
|
||||
|
||||
Returns the schema of the output that will be returned by this query.
|
||||
|
||||
This can be used to inspect the types and names of the columns that will be
|
||||
returned by the query before executing it.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Schema`<`any`>>
|
||||
|
||||
An Arrow Schema describing the output columns.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.outputSchema`
|
||||
|
||||
***
|
||||
|
||||
### select()
|
||||
|
||||
```ts
|
||||
select(columns): this
|
||||
```
|
||||
|
||||
Return only the specified columns.
|
||||
|
||||
By default a query will return all columns from the table. However, this can have
|
||||
a very significant impact on latency. LanceDb stores data in a columnar fashion. This
|
||||
means we can finely tune our I/O to select exactly the columns we need.
|
||||
|
||||
As a best practice you should always limit queries to the columns that you need. If you
|
||||
pass in an array of column names then only those columns will be returned.
|
||||
|
||||
You can also use this method to create new "dynamic" columns based on your existing columns.
|
||||
For example, you may not care about "a" or "b" but instead simply want "a + b". This is often
|
||||
seen in the SELECT clause of an SQL query (e.g. `SELECT a+b FROM my_table`).
|
||||
|
||||
To create dynamic columns you can pass in a Map<string, string>. A column will be returned
|
||||
for each entry in the map. The key provides the name of the column. The value is
|
||||
an SQL string used to specify how the column is calculated.
|
||||
|
||||
For example, an SQL query might state `SELECT a + b AS combined, c`. The equivalent
|
||||
input to this method would be:
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **columns**: `string` \| `string`[] \| `Record`<`string`, `string`> \| `Map`<`string`, `string`>
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
new Map([["combined", "a + b"], ["c", "c"]])
|
||||
|
||||
Columns will always be returned in the order given, even if that order is different than
|
||||
the order used when adding the data.
|
||||
|
||||
Note that you can pass in a `Record<string, string>` (e.g. an object literal). This method
|
||||
uses `Object.entries` which should preserve the insertion order of the object. However,
|
||||
object insertion order is easy to get wrong and `Map` is more foolproof.
|
||||
```
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.select`
|
||||
|
||||
***
|
||||
|
||||
### toArray()
|
||||
|
||||
```ts
|
||||
toArray(options?): Promise<any[]>
|
||||
```
|
||||
|
||||
Collect the results as an array of objects.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`any`[]>
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.toArray`
|
||||
|
||||
***
|
||||
|
||||
### toArrow()
|
||||
|
||||
```ts
|
||||
toArrow(options?): Promise<Table<any>>
|
||||
```
|
||||
|
||||
Collect the results as an Arrow
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **options?**: `Partial`<[`QueryExecutionOptions`](../interfaces/QueryExecutionOptions.md)>
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`Table`<`any`>>
|
||||
|
||||
#### See
|
||||
|
||||
ArrowTable.
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.toArrow`
|
||||
|
||||
***
|
||||
|
||||
### useLsm()
|
||||
|
||||
```ts
|
||||
useLsm(enable): this
|
||||
```
|
||||
|
||||
Control MemWAL read routing for this query.
|
||||
|
||||
By default (unset), when the table carries a MemWAL write spec (see
|
||||
[Table#setLsmWriteSpec](Table.md#setlsmwritespec)), reads are routed through the LSM scanner so
|
||||
they also return data written via the `mergeInsert` LSM path that has not yet
|
||||
been compacted into the base table (the active/frozen in-memory memtables and
|
||||
the flushed generations), deduplicated by primary key; a table without a spec
|
||||
reads the base table.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **enable**: `boolean`
|
||||
`true` forces the LSM scanner and errors if the table has no
|
||||
MemWAL write spec. `false` bypasses the MemWAL and reads the base table only,
|
||||
even when a spec is present.
|
||||
Note: the LSM scanner does not support every query shape (e.g. reranking,
|
||||
hybrid search, `orderBy`). On a MemWAL table those shapes error unless
|
||||
`useLsm(false)` is set, because a base-only read would silently exclude
|
||||
un-compacted MemWAL data.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.useLsm`
|
||||
|
||||
***
|
||||
|
||||
### where()
|
||||
|
||||
```ts
|
||||
where(predicate): this
|
||||
```
|
||||
|
||||
A filter statement to be applied to this query.
|
||||
|
||||
The filter should be supplied as an SQL query string. For example:
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **predicate**: `string`
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
x > 10
|
||||
y > 0 AND y < 100
|
||||
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
|
||||
|
||||
`StandardQueryBase.where`
|
||||
|
||||
***
|
||||
|
||||
### withRowId()
|
||||
|
||||
```ts
|
||||
withRowId(): this
|
||||
```
|
||||
|
||||
Whether to return the row id in the results.
|
||||
|
||||
This column can be used to match results between different queries. For
|
||||
example, to match results from a full text search and a vector search in
|
||||
order to perform hybrid search.
|
||||
|
||||
#### Returns
|
||||
|
||||
`this`
|
||||
|
||||
#### Inherited from
|
||||
|
||||
`StandardQueryBase.withRowId`
|
||||
@@ -584,6 +584,70 @@ Child namespace names and
|
||||
|
||||
***
|
||||
|
||||
### listTables()
|
||||
|
||||
#### listTables(options)
|
||||
|
||||
```ts
|
||||
abstract listTables(options?): Promise<ListTablesResponse>
|
||||
```
|
||||
|
||||
List a page of the tables in this database.
|
||||
|
||||
To retrieve the tables after the page, pass the `pageToken` the response
|
||||
carries back in. A page can be shorter than `limit` without being the last
|
||||
one, so walk until a response carries no page token:
|
||||
|
||||
```ts
|
||||
const names = [];
|
||||
let pageToken = undefined;
|
||||
do {
|
||||
const page = await conn.listTables({ pageToken, limit: 100 });
|
||||
names.push(...page.tables);
|
||||
pageToken = page.pageToken;
|
||||
} while (pageToken);
|
||||
```
|
||||
|
||||
##### Parameters
|
||||
|
||||
* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)>
|
||||
Pagination options
|
||||
(`pageToken`, `limit`).
|
||||
|
||||
##### Returns
|
||||
|
||||
`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)>
|
||||
|
||||
A page of table names and an
|
||||
optional token for the tables after it.
|
||||
|
||||
#### listTables(namespacePath, options)
|
||||
|
||||
```ts
|
||||
abstract listTables(namespacePath?, options?): Promise<ListTablesResponse>
|
||||
```
|
||||
|
||||
List a page of the tables in this database.
|
||||
|
||||
##### Parameters
|
||||
|
||||
* **namespacePath?**: `string`[]
|
||||
The namespace path to list tables from
|
||||
(defaults to root namespace)
|
||||
|
||||
* **options?**: `Partial`<[`ListTablesOptions`](../interfaces/ListTablesOptions.md)>
|
||||
Pagination options
|
||||
(`pageToken`, `limit`).
|
||||
|
||||
##### Returns
|
||||
|
||||
`Promise`<[`ListTablesResponse`](../interfaces/ListTablesResponse.md)>
|
||||
|
||||
A page of table names and an
|
||||
optional token for the tables after it.
|
||||
|
||||
***
|
||||
|
||||
### openMaterializedView()
|
||||
|
||||
```ts
|
||||
@@ -660,7 +724,7 @@ a "not supported" error.
|
||||
|
||||
***
|
||||
|
||||
### tableNames()
|
||||
### ~~tableNames()~~
|
||||
|
||||
#### tableNames(options)
|
||||
|
||||
@@ -682,6 +746,10 @@ Tables will be returned in lexicographical order.
|
||||
|
||||
`Promise`<`string`[]>
|
||||
|
||||
##### Deprecated
|
||||
|
||||
Use [Connection.listTables](Connection.md#listtables) instead.
|
||||
|
||||
#### tableNames(namespacePath, options)
|
||||
|
||||
```ts
|
||||
@@ -704,3 +772,7 @@ Tables will be returned in lexicographical order.
|
||||
##### Returns
|
||||
|
||||
`Promise`<`string`[]>
|
||||
|
||||
##### Deprecated
|
||||
|
||||
Use [Connection.listTables](Connection.md#listtables) instead.
|
||||
|
||||
@@ -942,7 +942,7 @@ Get the schema of the table.
|
||||
abstract search(
|
||||
query,
|
||||
queryType?,
|
||||
ftsColumns?): Query | VectorQuery
|
||||
ftsColumns?): Query | VectorQuery | AutoQuery
|
||||
```
|
||||
|
||||
Create a search query to find the nearest neighbors
|
||||
@@ -964,7 +964,7 @@ of the given query
|
||||
|
||||
#### Returns
|
||||
|
||||
[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md)
|
||||
[`Query`](Query.md) \| [`VectorQuery`](VectorQuery.md) \| [`AutoQuery`](AutoQuery.md)
|
||||
|
||||
***
|
||||
|
||||
@@ -1292,6 +1292,18 @@ abstract updateFieldMetadata(updates): Promise<UpdateFieldMetadataResult>
|
||||
|
||||
Update per-field (column) metadata.
|
||||
|
||||
The following keys are treated specially, by convention, and should be
|
||||
used when appropriate:
|
||||
|
||||
- `lancedb:description`: for a human-readable description of a field.
|
||||
- `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
- `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
`feature_v2` might be in the same logical column.
|
||||
- `lancedb:status`: for status options (`production`, `candidate`,
|
||||
`deprecated`, `archived`) to designate the current life cycle state of
|
||||
this column.
|
||||
|
||||
#### Parameters
|
||||
|
||||
* **updates**: [`FieldMetadataUpdate`](../interfaces/FieldMetadataUpdate.md)[]
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
|
||||
## Classes
|
||||
|
||||
- [AutoQuery](classes/AutoQuery.md)
|
||||
- [BooleanQuery](classes/BooleanQuery.md)
|
||||
- [BoostQuery](classes/BoostQuery.md)
|
||||
- [BranchContents](classes/BranchContents.md)
|
||||
@@ -100,6 +101,8 @@
|
||||
- [JobInfo](interfaces/JobInfo.md)
|
||||
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
|
||||
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
|
||||
- [ListTablesOptions](interfaces/ListTablesOptions.md)
|
||||
- [ListTablesResponse](interfaces/ListTablesResponse.md)
|
||||
- [LsmStats](interfaces/LsmStats.md)
|
||||
- [LsmWriteSpec](interfaces/LsmWriteSpec.md)
|
||||
- [MaterializedViewDefinition](interfaces/MaterializedViewDefinition.md)
|
||||
|
||||
@@ -17,7 +17,8 @@ metadata: Record<string, null | string>;
|
||||
```
|
||||
|
||||
Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
default; a value of `null` deletes that key.
|
||||
default; a value of `null` deletes that key. See
|
||||
[Table.updateFieldMetadata](../classes/Table.md#updatefieldmetadata) for the conventional `lancedb:*` keys.
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / ListTablesOptions
|
||||
|
||||
# Interface: ListTablesOptions
|
||||
|
||||
## Properties
|
||||
|
||||
### limit?
|
||||
|
||||
```ts
|
||||
optional limit: number;
|
||||
```
|
||||
|
||||
An upper bound on how many tables to return.
|
||||
|
||||
A page may hold fewer than this and still not be the last one, so keep
|
||||
going while the response carries a page token rather than while pages are
|
||||
full.
|
||||
|
||||
***
|
||||
|
||||
### pageToken?
|
||||
|
||||
```ts
|
||||
optional pageToken: string;
|
||||
```
|
||||
|
||||
Token from a previous response, to resume listing where it left off.
|
||||
|
||||
The token is opaque: it carries whatever the database needs to resume, and
|
||||
callers should not construct or interpret one.
|
||||
@@ -0,0 +1,23 @@
|
||||
[**@lancedb/lancedb**](../README.md) • **Docs**
|
||||
|
||||
***
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / ListTablesResponse
|
||||
|
||||
# Interface: ListTablesResponse
|
||||
|
||||
## Properties
|
||||
|
||||
### pageToken?
|
||||
|
||||
```ts
|
||||
optional pageToken: string;
|
||||
```
|
||||
|
||||
***
|
||||
|
||||
### tables
|
||||
|
||||
```ts
|
||||
tables: string[];
|
||||
```
|
||||
@@ -4,11 +4,16 @@
|
||||
|
||||
[@lancedb/lancedb](../globals.md) / TableNamesOptions
|
||||
|
||||
# Interface: TableNamesOptions
|
||||
# Interface: ~~TableNamesOptions~~
|
||||
|
||||
## Deprecated
|
||||
|
||||
Use [ListTablesOptions](ListTablesOptions.md) with [Connection.listTables](../classes/Connection.md#listtables)
|
||||
instead.
|
||||
|
||||
## Properties
|
||||
|
||||
### limit?
|
||||
### ~~limit?~~
|
||||
|
||||
```ts
|
||||
optional limit: number;
|
||||
@@ -18,7 +23,7 @@ An optional limit to the number of results to return.
|
||||
|
||||
***
|
||||
|
||||
### startAfter?
|
||||
### ~~startAfter?~~
|
||||
|
||||
```ts
|
||||
optional startAfter: string;
|
||||
|
||||
@@ -10,16 +10,12 @@
|
||||
function getRegistry(): EmbeddingFunctionRegistry
|
||||
```
|
||||
|
||||
Utility function to get the global instance of the registry
|
||||
Get the global embedding function registry.
|
||||
|
||||
LanceDB built-in providers are initialized when this public API is first
|
||||
used, so importing the root package does not change automatic search
|
||||
selection for tables without embedding metadata.
|
||||
|
||||
## Returns
|
||||
|
||||
[`EmbeddingFunctionRegistry`](../classes/EmbeddingFunctionRegistry.md)
|
||||
|
||||
`EmbeddingFunctionRegistry` The global instance of the registry
|
||||
|
||||
## Example
|
||||
|
||||
```ts
|
||||
const registry = getRegistry();
|
||||
const openai = registry.get("openai").create();
|
||||
|
||||
@@ -159,6 +159,8 @@ and combined with [BooleanQuery][lancedb.query.BooleanQuery].
|
||||
|
||||
::: lancedb.query.FullTextOperator
|
||||
|
||||
::: lancedb.query.DocumentGranularity
|
||||
|
||||
::: lancedb.query.Occur
|
||||
|
||||
## Embeddings
|
||||
@@ -221,9 +223,13 @@ tokens = list(
|
||||
Blob columns store large binary values out of line so they can be read lazily
|
||||
instead of being materialized with the rest of the row.
|
||||
|
||||
::: lancedb.blob
|
||||
`lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without
|
||||
pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns
|
||||
still work. Queries return descriptors. Call
|
||||
[`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or
|
||||
[`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes.
|
||||
|
||||
::: lancedb.BlobType
|
||||
::: lancedb.blob
|
||||
|
||||
::: lancedb._blob.BlobFile
|
||||
options:
|
||||
@@ -261,6 +267,8 @@ instead of being materialized with the rest of the row.
|
||||
|
||||
::: lancedb.streaming.StreamingDataset
|
||||
|
||||
::: lancedb.streaming.StreamingDataLoader
|
||||
|
||||
::: lancedb.permutation.permutation_builder
|
||||
|
||||
::: lancedb.permutation.PermutationBuilder
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
{
|
||||
"include": [
|
||||
"src/*.ts",
|
||||
],
|
||||
"compilerOptions": {
|
||||
"target": "es2022",
|
||||
"module": "nodenext",
|
||||
"declaration": true,
|
||||
"outDir": "./dist",
|
||||
"strict": true,
|
||||
"allowJs": true,
|
||||
"resolveJsonModule": true,
|
||||
},
|
||||
"exclude": [
|
||||
"./dist/*",
|
||||
]
|
||||
}
|
||||
@@ -8,7 +8,7 @@
|
||||
<parent>
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.6</version>
|
||||
<version>0.38.0-beta.12</version>
|
||||
<relativePath>../pom.xml</relativePath>
|
||||
</parent>
|
||||
|
||||
|
||||
+2
-2
@@ -6,7 +6,7 @@
|
||||
|
||||
<groupId>com.lancedb</groupId>
|
||||
<artifactId>lancedb-parent</artifactId>
|
||||
<version>0.38.0-beta.6</version>
|
||||
<version>0.38.0-beta.12</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>11.0.0-rc.1</lance-core.version>
|
||||
<lance-core.version>12.0.0-beta.2</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
@@ -1,7 +1,7 @@
|
||||
[package]
|
||||
name = "lancedb-nodejs"
|
||||
edition.workspace = true
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
publish = false
|
||||
license.workspace = true
|
||||
description.workspace = true
|
||||
|
||||
@@ -1,11 +1,16 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
import * as fs from "node:fs";
|
||||
import * as vm from "node:vm";
|
||||
import * as arrow15 from "apache-arrow-15";
|
||||
import * as arrow16 from "apache-arrow-16";
|
||||
import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
|
||||
import {
|
||||
Field as CurrentField,
|
||||
LargeBinary as CurrentLargeBinary,
|
||||
Schema as CurrentSchema,
|
||||
Vector as CurrentVector,
|
||||
convertToTable,
|
||||
tableFromIPC as currentTableFromIPC,
|
||||
@@ -36,6 +41,59 @@ function sampleRecords(): Array<Record<string, any>> {
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
it("serializes an Arrow Table created in another JavaScript realm", async () => {
|
||||
const context = vm.createContext({
|
||||
TextDecoder,
|
||||
TextEncoder,
|
||||
console,
|
||||
setTimeout,
|
||||
clearTimeout,
|
||||
});
|
||||
vm.runInContext(
|
||||
fs.readFileSync(
|
||||
require.resolve("apache-arrow-15/Arrow.es2015.min"),
|
||||
"utf8",
|
||||
),
|
||||
context,
|
||||
);
|
||||
const foreignTable: unknown = vm.runInContext(
|
||||
"Arrow.tableFromArrays({ id: new Int32Array([1, 2, 3]), text: ['foo', 'bar', 'baz'] })",
|
||||
context,
|
||||
);
|
||||
|
||||
const foreignMetadata = (
|
||||
foreignTable as { schema: { metadata: Map<string, string> } }
|
||||
).schema.metadata;
|
||||
expect(foreignMetadata).not.toBeInstanceOf(Map);
|
||||
|
||||
const buf = await fromDataToBuffer(
|
||||
foreignTable as Parameters<typeof fromDataToBuffer>[0],
|
||||
);
|
||||
const actual = currentTableFromIPC(buf);
|
||||
|
||||
expect(actual.numRows).toBe(3);
|
||||
expect(actual.getChild("id")?.toJSON()).toEqual([1, 2, 3]);
|
||||
expect(actual.getChild("text")?.toJSON()).toEqual(["foo", "bar", "baz"]);
|
||||
});
|
||||
|
||||
it("preserves field metadata from a provided schema", async function () {
|
||||
const jsonMetadata = new Map([["ARROW:extension:name", "lance.json"]]);
|
||||
const schema = new CurrentSchema([
|
||||
new CurrentField("meta", new CurrentLargeBinary(), true, jsonMetadata),
|
||||
]);
|
||||
|
||||
const table = makeArrowTable(
|
||||
[{ meta: Buffer.from(JSON.stringify({ source: "test" })) }],
|
||||
{ schema },
|
||||
);
|
||||
|
||||
expect(table.schema.fields[0].metadata).toEqual(jsonMetadata);
|
||||
|
||||
const roundTripped = currentTableFromIPC(await fromTableToBuffer(table));
|
||||
expect(roundTripped.schema.fields[0].metadata).toEqual(jsonMetadata);
|
||||
});
|
||||
|
||||
describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
"Arrow",
|
||||
(
|
||||
@@ -515,6 +573,137 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
);
|
||||
});
|
||||
|
||||
it("will allow matching inferred types across records", function () {
|
||||
expect(() =>
|
||||
makeArrowTable([{ value: 1 }, { value: 2 }]),
|
||||
).not.toThrow();
|
||||
});
|
||||
|
||||
it("will reject mismatched inferred types across records", function () {
|
||||
expect(() => makeArrowTable([{ value: 1 }, { value: "two" }])).toThrow(
|
||||
"Failed to infer schema for data. Previously inferred type Float64 but found Utf8 for field value at row 1. Consider providing an explicit schema.",
|
||||
);
|
||||
});
|
||||
|
||||
it("will ignore generated dictionary IDs when comparing inferred types", function () {
|
||||
const table = makeArrowTable([{ str: "a" }, { str: "b" }], {
|
||||
dictionaryEncodeStrings: true,
|
||||
});
|
||||
|
||||
expect(table.getChild("str")?.toJSON()).toEqual(["a", "b"]);
|
||||
});
|
||||
|
||||
it("will preserve null values without treating them as type mismatches", function () {
|
||||
for (const records of [
|
||||
[{ vector: [1, 2, 3] }, { vector: null }],
|
||||
[{ vector: null }, { vector: [1, 2, 3] }],
|
||||
]) {
|
||||
const table = makeArrowTable(records);
|
||||
|
||||
expect(table.numRows).toBe(2);
|
||||
expect(table.getChild("vector")?.nullCount).toBe(1);
|
||||
}
|
||||
});
|
||||
|
||||
it("will preserve empty variable-size lists", function () {
|
||||
for (const records of [
|
||||
[{ items: [1] }, { items: [] }],
|
||||
[{ items: [] }, { items: [1] }],
|
||||
]) {
|
||||
const table = makeArrowTable(records);
|
||||
expect(
|
||||
table
|
||||
.getChild("items")
|
||||
?.toJSON()
|
||||
.map((value) => value.toJSON()),
|
||||
).toEqual(records.map((record) => record.items));
|
||||
}
|
||||
});
|
||||
|
||||
it("will propagate deferred evidence through nested lists", function () {
|
||||
for (const records of [
|
||||
[{ items: [1] }, { items: [null] }],
|
||||
[{ items: [null] }, { items: [1] }],
|
||||
[{ items: [null, 1] }, { items: [2, null] }],
|
||||
]) {
|
||||
const table = makeArrowTable(records);
|
||||
expect(
|
||||
table
|
||||
.getChild("items")
|
||||
?.toJSON()
|
||||
.map((value) => value.toJSON()),
|
||||
).toEqual(records.map((record) => record.items));
|
||||
}
|
||||
|
||||
const nestedRecords = [{ items: [[1]] }, { items: [[null]] }];
|
||||
const nestedTable = makeArrowTable(nestedRecords);
|
||||
expect(
|
||||
nestedTable
|
||||
.getChild("items")
|
||||
?.toJSON()
|
||||
.map((value) =>
|
||||
value
|
||||
.toJSON()
|
||||
.map((nestedValue: { toJSON: () => unknown[] }) =>
|
||||
nestedValue.toJSON(),
|
||||
),
|
||||
),
|
||||
).toEqual(nestedRecords.map((record) => record.items));
|
||||
});
|
||||
|
||||
it("will reject incompatible deferred evidence within a list", function () {
|
||||
for (const items of [
|
||||
[[], 1],
|
||||
[1, []],
|
||||
[[null], 1],
|
||||
[1, [null]],
|
||||
]) {
|
||||
expect(() => makeArrowTable([{ items }])).toThrow(
|
||||
"Failed to infer data type for field items at row 0.",
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
it("will reject empty fixed-size lists", function () {
|
||||
expect(() =>
|
||||
makeArrowTable([{ vector: [1, 2, 3] }, { vector: [] }]),
|
||||
).toThrow(
|
||||
"Failed to infer schema for data. Previously inferred type FixedSizeList[3]<Float32> but found List[0] for field vector at row 1.",
|
||||
);
|
||||
});
|
||||
|
||||
it("will reject inferred leaf and branch shape changes", function () {
|
||||
expect(() =>
|
||||
makeArrowTable([{ value: 1 }, { value: { nested: 2 } }]),
|
||||
).toThrow(
|
||||
"Failed to infer schema for data. Previously inferred type Float64 but found Struct for field value at row 1.",
|
||||
);
|
||||
expect(() =>
|
||||
makeArrowTable([{ value: { nested: 1 } }, { value: 2 }]),
|
||||
).toThrow(
|
||||
"Failed to infer schema for data. Previously inferred type Struct but found Float64 for field value at row 1.",
|
||||
);
|
||||
});
|
||||
|
||||
it("will allow null values around inferred struct values", function () {
|
||||
for (const { records, nullIndex } of [
|
||||
{
|
||||
records: [{ value: null }, { value: { nested: 2 } }],
|
||||
nullIndex: 0,
|
||||
},
|
||||
{
|
||||
records: [{ value: { nested: 1 } }, { value: null }],
|
||||
nullIndex: 1,
|
||||
},
|
||||
]) {
|
||||
const table = makeArrowTable(records);
|
||||
const values = table.getChild("value");
|
||||
|
||||
expect(values?.nullCount).toBe(1);
|
||||
expect(values?.get(nullIndex)).toBeNull();
|
||||
}
|
||||
});
|
||||
|
||||
it("will allow a schema to be provided", async function () {
|
||||
await checkTableCreation(
|
||||
async (records, _, schema) =>
|
||||
|
||||
@@ -4,7 +4,13 @@
|
||||
import { readdirSync } from "fs";
|
||||
import { Field, Float64, Schema } from "apache-arrow";
|
||||
import * as tmp from "tmp";
|
||||
import { Connection, Table, connect, connectNamespace } from "../lancedb";
|
||||
import {
|
||||
Connection,
|
||||
ListTablesResponse,
|
||||
Table,
|
||||
connect,
|
||||
connectNamespace,
|
||||
} from "../lancedb";
|
||||
import { LocalTable } from "../lancedb/table";
|
||||
|
||||
describe("when connecting", () => {
|
||||
@@ -47,6 +53,7 @@ describe("given a connection", () => {
|
||||
await db.close();
|
||||
expect(db.isOpen()).toBe(false);
|
||||
await expect(db.tableNames()).rejects.toThrow("Connection is closed");
|
||||
await expect(db.listTables()).rejects.toThrow("Connection is closed");
|
||||
await expect(db.renameTable("a", "b")).rejects.toThrow(
|
||||
"Connection is closed",
|
||||
);
|
||||
@@ -129,6 +136,66 @@ describe("given a connection", () => {
|
||||
expect(tables).toEqual(["b", "c"]);
|
||||
});
|
||||
|
||||
it("should respect limit and page token when listing tables", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
await db.createTable("b", [{ id: 1 }]);
|
||||
await db.createTable("a", [{ id: 1 }]);
|
||||
await db.createTable("c", [{ id: 1 }]);
|
||||
|
||||
const all = await db.listTables();
|
||||
expect(all.tables).toEqual(["a", "b", "c"]);
|
||||
expect(all.pageToken).toBeUndefined();
|
||||
|
||||
const first = await db.listTables({ limit: 1 });
|
||||
expect(first.tables).toEqual(["a"]);
|
||||
expect(first.pageToken).toBeDefined();
|
||||
|
||||
const second = await db.listTables({
|
||||
limit: 1,
|
||||
pageToken: first.pageToken,
|
||||
});
|
||||
expect(second.tables).toEqual(["b"]);
|
||||
});
|
||||
|
||||
it("should visit every table exactly once when walking pages", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
|
||||
const created = ["a", "b", "c", "d", "e"];
|
||||
for (const name of created) {
|
||||
await db.createTable(name, [{ id: 1 }]);
|
||||
}
|
||||
|
||||
const seen: string[] = [];
|
||||
let pageToken: string | undefined = undefined;
|
||||
do {
|
||||
const page: ListTablesResponse = await db.listTables({
|
||||
limit: 2,
|
||||
pageToken,
|
||||
});
|
||||
seen.push(...page.tables);
|
||||
pageToken = page.pageToken;
|
||||
} while (pageToken);
|
||||
|
||||
expect(seen).toEqual(created);
|
||||
});
|
||||
|
||||
it("should list tables in a namespace", async () => {
|
||||
const db = await connect(tmpDir.name, {
|
||||
// biome-ignore lint/style/useNamingConvention: opaque backend property key, must match Rust
|
||||
namespaceClientProperties: { manifest_enabled: "true" },
|
||||
});
|
||||
await db.createNamespace(["child"]);
|
||||
await db.createTable("nested", [{ id: 1 }], ["child"]);
|
||||
|
||||
await expect(db.listTables(["child"])).resolves.toEqual(
|
||||
expect.objectContaining({ tables: ["nested"] }),
|
||||
);
|
||||
await expect(db.listTables()).resolves.toEqual(
|
||||
expect.objectContaining({ tables: [] }),
|
||||
);
|
||||
});
|
||||
|
||||
it("should create tables in v2 mode", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [...Array(10000).keys()].map((i) => ({ id: i }));
|
||||
|
||||
@@ -187,6 +187,58 @@ describe("embedding functions", () => {
|
||||
const vector0 = JSON.parse(JSON.stringify(arr[0].vector));
|
||||
expect(vector0).toEqual([1, 2, 3]);
|
||||
});
|
||||
it("should append multiple Python embeddings with the same alias", async () => {
|
||||
@register("python-mock")
|
||||
// biome-ignore lint/correctness/noUnusedVariables: the decorator registers this class
|
||||
class MockEmbeddingFunction extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType(): Float {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((value) =>
|
||||
value === "hello world" ? [1, 2, 3] : [4, 5, 6],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const metadata = new Map([
|
||||
[
|
||||
"embedding_functions",
|
||||
'[{"source_column":"text1","vector_column":"vector1","name":"python-mock","model":{}},{"source_column":"text2","vector_column":"vector2","name":"python-mock","model":{}}]',
|
||||
],
|
||||
]);
|
||||
const schema = new Schema(
|
||||
[
|
||||
new Field("text1", new Utf8(), true),
|
||||
new Field("text2", new Utf8(), true),
|
||||
new Field(
|
||||
"vector1",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
new Field(
|
||||
"vector2",
|
||||
new FixedSizeList(3, new Field("item", new Float32(), true)),
|
||||
true,
|
||||
),
|
||||
],
|
||||
metadata,
|
||||
);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const table = await db.createEmptyTable("test", schema);
|
||||
await table.add([{ text1: "hello world", text2: "goodbye world" }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector1))).toEqual([1, 2, 3]);
|
||||
expect(JSON.parse(JSON.stringify(rows[0].vector2))).toEqual([4, 5, 6]);
|
||||
});
|
||||
|
||||
it("should append generated vectors to a non-nullable schema", async () => {
|
||||
@register("non_nullable_schema_test")
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { execFileSync } from "node:child_process";
|
||||
import { resolve } from "node:path";
|
||||
|
||||
import type { OpenAIEmbeddingFunction } from "../lancedb/embedding/openai";
|
||||
import type { EmbeddingFunctionRegistry } from "../lancedb/embedding/registry";
|
||||
|
||||
type EmbeddingModule = typeof import("../lancedb/embedding");
|
||||
type OpenAIModule = typeof import("../lancedb/embedding/openai");
|
||||
type RegistryModule = typeof import("../lancedb/embedding/registry");
|
||||
|
||||
describe("embedding function registry", () => {
|
||||
const registries: EmbeddingFunctionRegistry[] = [];
|
||||
|
||||
afterEach(() => {
|
||||
for (const registry of registries) {
|
||||
registry.reset();
|
||||
}
|
||||
registries.length = 0;
|
||||
});
|
||||
|
||||
it("defers built-in providers until the public registry API is used", () => {
|
||||
jest.isolateModules(() => {
|
||||
const embedding = require("../lancedb/embedding") as EmbeddingModule;
|
||||
const { getRegistry: getInternalRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
const registry = getInternalRegistry();
|
||||
registries.push(registry);
|
||||
|
||||
expect(registry.length()).toBe(0);
|
||||
expect(embedding.getRegistry()).toBe(registry);
|
||||
expect(registry.get("openai")).toBeDefined();
|
||||
expect(registry.get("huggingface")).toBeDefined();
|
||||
});
|
||||
});
|
||||
|
||||
it("preserves automatic FTS search in a fresh process", () => {
|
||||
execFileSync(
|
||||
process.execPath,
|
||||
[resolve(__dirname, "fixtures", "auto_fts_search.cjs")],
|
||||
{ stdio: "pipe" },
|
||||
);
|
||||
});
|
||||
|
||||
it("shares registrations across duplicated provider module graphs", () => {
|
||||
let registeringRegistry: EmbeddingFunctionRegistry | undefined;
|
||||
let latestOpenAIConstructor: typeof OpenAIEmbeddingFunction | undefined;
|
||||
|
||||
jest.isolateModules(() => {
|
||||
require("../lancedb/embedding/openai");
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
registeringRegistry = getRegistry();
|
||||
registries.push(registeringRegistry);
|
||||
expect(registeringRegistry.get("openai")).toBeDefined();
|
||||
});
|
||||
|
||||
expect(() => {
|
||||
jest.isolateModules(() => {
|
||||
const { OpenAIEmbeddingFunction } =
|
||||
require("../lancedb/embedding/openai") as OpenAIModule;
|
||||
latestOpenAIConstructor = OpenAIEmbeddingFunction;
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding/registry") as RegistryModule;
|
||||
registries.push(getRegistry());
|
||||
});
|
||||
}).not.toThrow();
|
||||
|
||||
const previousApiKey = process.env.OPENAI_API_KEY;
|
||||
process.env.OPENAI_API_KEY = "test";
|
||||
try {
|
||||
const latestOpenAI = registeringRegistry!
|
||||
.get<OpenAIEmbeddingFunction>("openai")!
|
||||
.create();
|
||||
expect(latestOpenAI).toBeInstanceOf(latestOpenAIConstructor!);
|
||||
} finally {
|
||||
if (previousApiKey === undefined) {
|
||||
delete process.env.OPENAI_API_KEY;
|
||||
} else {
|
||||
process.env.OPENAI_API_KEY = previousApiKey;
|
||||
}
|
||||
}
|
||||
|
||||
jest.isolateModules(() => {
|
||||
const { getRegistry } =
|
||||
require("../lancedb/embedding") as EmbeddingModule;
|
||||
const publicRegistry = getRegistry();
|
||||
registries.push(publicRegistry);
|
||||
expect(publicRegistry).toBe(registeringRegistry);
|
||||
expect(publicRegistry.get("openai")).toBeDefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,33 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
const assert = require("node:assert/strict");
|
||||
const tmp = require("tmp");
|
||||
const { connect, embedding, Index } = require("../../dist");
|
||||
const { getRegistry } = require("../../dist/embedding/registry");
|
||||
|
||||
async function main() {
|
||||
assert.equal(typeof embedding.getRegistry, "function");
|
||||
assert.equal(getRegistry().length(), 0);
|
||||
assert.equal(embedding.getRegistry(), getRegistry());
|
||||
assert.equal(getRegistry().length(), 2);
|
||||
|
||||
const dir = tmp.dirSync({ unsafeCleanup: true });
|
||||
let db;
|
||||
try {
|
||||
db = await connect(dir.name);
|
||||
const table = await db.createTable("docs", [{ text: "hello world" }]);
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const rows = await table.search("hello").toArray();
|
||||
assert.equal(rows[0].text, "hello world");
|
||||
} finally {
|
||||
db?.close();
|
||||
dir.removeCallback();
|
||||
}
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
});
|
||||
@@ -5,8 +5,8 @@ import packageJson = require("../package.json");
|
||||
|
||||
describe("package metadata", () => {
|
||||
it("requires Node.js type declarations compatible with the runtime", () => {
|
||||
expect(packageJson.engines.node).toBe(">= 18");
|
||||
expect(packageJson.peerDependencies["@types/node"]).toBe(">=18");
|
||||
expect(packageJson.engines.node).toBe(">= 22");
|
||||
expect(packageJson.peerDependencies["@types/node"]).toBe(">=22");
|
||||
expect(packageJson.peerDependenciesMeta["@types/node"]).toEqual({
|
||||
optional: true,
|
||||
});
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import * as http from "http";
|
||||
import { RequestListener } from "http";
|
||||
import packageJson = require("../package.json");
|
||||
import {
|
||||
ClientConfig,
|
||||
Connection,
|
||||
@@ -70,7 +71,13 @@ async function withMockDatabase(
|
||||
try {
|
||||
await callback(db);
|
||||
} finally {
|
||||
server.close();
|
||||
// `close()` alone leaves the port bound until keep-alive sockets drain, so
|
||||
// a single failing test would cascade into EADDRINUSE for every test after
|
||||
// it. Destroy the connections and wait for the port to actually be free.
|
||||
await new Promise<void>((resolve) => {
|
||||
server.closeAllConnections();
|
||||
server.close(() => resolve());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -131,7 +138,7 @@ describe("remote connection", () => {
|
||||
(req, res) => {
|
||||
expect(req.headers["x-api-key"]).toEqual("fake");
|
||||
expect(req.headers["user-agent"]).toEqual(
|
||||
`LanceDB-Node-Client/${process.env.npm_package_version}`,
|
||||
`LanceDB-Node-Client/${packageJson.version}`,
|
||||
);
|
||||
|
||||
const body = JSON.stringify({ tables: [] });
|
||||
|
||||
@@ -11,10 +11,13 @@ import * as arrow17 from "apache-arrow-17";
|
||||
import * as arrow18 from "apache-arrow-18";
|
||||
|
||||
import {
|
||||
AutoQuery,
|
||||
Connection,
|
||||
MatchQuery,
|
||||
PhraseQuery,
|
||||
Query,
|
||||
Table,
|
||||
VectorQuery,
|
||||
connect,
|
||||
tokenize,
|
||||
} from "../lancedb";
|
||||
@@ -682,6 +685,56 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
},
|
||||
);
|
||||
|
||||
// https://github.com/lancedb/lancedb/issues/1963
|
||||
it("should query documents with LangChain PDF metadata", async () => {
|
||||
const tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
try {
|
||||
const db = await connect(tmpDir.name);
|
||||
const documents = [
|
||||
{
|
||||
text: "first page",
|
||||
vector: [1, 0],
|
||||
source: "first.pdf",
|
||||
loc: { pageNumber: 1, lines: { from: 1, to: 12 } },
|
||||
pdf: {
|
||||
version: "1.10.100",
|
||||
info: {
|
||||
format: "PDF 1.7",
|
||||
producer: "pdf.js",
|
||||
creator: "Writer",
|
||||
},
|
||||
totalPages: 2,
|
||||
},
|
||||
},
|
||||
{
|
||||
text: "second page",
|
||||
vector: [0, 1],
|
||||
source: "second.pdf",
|
||||
loc: { pageNumber: 2, lines: { from: 13, to: 24 } },
|
||||
pdf: {
|
||||
version: "1.10.100",
|
||||
info: {
|
||||
format: "PDF 1.7",
|
||||
producer: "pdf.js",
|
||||
creator: "Writer",
|
||||
},
|
||||
totalPages: 2,
|
||||
},
|
||||
},
|
||||
];
|
||||
const documentsTable = await db.createTable("documents", documents);
|
||||
|
||||
const results = await documentsTable.query().toArray();
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].source).toBe("first.pdf");
|
||||
expect(results[0].pdf.info.producer).toBe("pdf.js");
|
||||
expect(results[1].loc.pageNumber).toBe(2);
|
||||
} finally {
|
||||
tmpDir.removeCallback();
|
||||
}
|
||||
});
|
||||
|
||||
describe("merge insert", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
let table: Table;
|
||||
@@ -1777,6 +1830,194 @@ describe("Read consistency interval", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe("automatic search schema consistency", () => {
|
||||
let tmpDir: tmp.DirResult;
|
||||
|
||||
class SchemaRefreshEmbedding extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 2;
|
||||
}
|
||||
|
||||
embeddingDataType() {
|
||||
return new Float32();
|
||||
}
|
||||
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map((value) => [value.length, 1]);
|
||||
}
|
||||
|
||||
async computeQueryEmbeddings(value: string) {
|
||||
return [value.length, 1];
|
||||
}
|
||||
}
|
||||
|
||||
function embeddingSchema() {
|
||||
const func = new SchemaRefreshEmbedding();
|
||||
return LanceSchema({
|
||||
text: func.sourceField(new Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
getRegistry().reset();
|
||||
register("schema-refresh")(SchemaRefreshEmbedding);
|
||||
tmpDir = tmp.dirSync({ unsafeCleanup: true });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
getRegistry().reset();
|
||||
tmpDir.removeCallback();
|
||||
});
|
||||
|
||||
it("uses the schema refreshed from another connection", async () => {
|
||||
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
|
||||
try {
|
||||
const stale = await first.createTable("docs", [{ text: "before" }], {
|
||||
schema: embeddingSchema(),
|
||||
});
|
||||
const replacement = await second.createTable(
|
||||
"docs",
|
||||
[{ text: "after hello" }],
|
||||
{ mode: "overwrite" },
|
||||
);
|
||||
await replacement.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const search = stale.search("hello");
|
||||
expect(search).toBeInstanceOf(AutoQuery);
|
||||
expect(search).not.toBeInstanceOf(Query);
|
||||
expect(search).not.toBeInstanceOf(VectorQuery);
|
||||
expect("nprobes" in search).toBe(false);
|
||||
|
||||
const rows = await search.toArray();
|
||||
expect(rows[0].text).toBe("after hello");
|
||||
expect((await stale.schema()).metadata.has("embedding_functions")).toBe(
|
||||
false,
|
||||
);
|
||||
} finally {
|
||||
first.close();
|
||||
second.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("tracks embedding metadata across checkout and restore", async () => {
|
||||
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
|
||||
try {
|
||||
await first.createTable("docs", [{ text: "before" }], {
|
||||
schema: embeddingSchema(),
|
||||
});
|
||||
const table = await second.createTable(
|
||||
"docs",
|
||||
[{ text: "after hello" }],
|
||||
{ mode: "overwrite" },
|
||||
);
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
|
||||
await table.checkout(1);
|
||||
expect((await table.search("before").toArray())[0].text).toBe("before");
|
||||
|
||||
await table.checkoutLatest();
|
||||
expect((await table.search("hello").toArray())[0].text).toBe(
|
||||
"after hello",
|
||||
);
|
||||
|
||||
await table.checkout(1);
|
||||
await table.restore();
|
||||
expect((await table.search("before").toArray())[0].text).toBe("before");
|
||||
} finally {
|
||||
first.close();
|
||||
second.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("pins automatic search while computing an embedding", async () => {
|
||||
let markStarted!: () => void;
|
||||
let releaseEmbedding!: () => void;
|
||||
const started = new Promise<void>((resolve) => {
|
||||
markStarted = resolve;
|
||||
});
|
||||
const released = new Promise<void>((resolve) => {
|
||||
releaseEmbedding = resolve;
|
||||
});
|
||||
|
||||
class BlockingEmbedding extends SchemaRefreshEmbedding {
|
||||
async computeQueryEmbeddings(value: string) {
|
||||
markStarted();
|
||||
await released;
|
||||
return [value.length, 1];
|
||||
}
|
||||
}
|
||||
|
||||
register("schema-refresh-blocking")(BlockingEmbedding);
|
||||
const func = new BlockingEmbedding();
|
||||
const schema = LanceSchema({
|
||||
text: func.sourceField(new Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
|
||||
try {
|
||||
const table = await first.createTable(
|
||||
"docs",
|
||||
[{ text: "hello before" }],
|
||||
{ schema },
|
||||
);
|
||||
const pending = table.search("hello").toArray();
|
||||
await started;
|
||||
|
||||
const replacement = await second.createTable(
|
||||
"docs",
|
||||
[{ text: "hello after" }],
|
||||
{ mode: "overwrite" },
|
||||
);
|
||||
await replacement.createIndex("text", { config: Index.fts() });
|
||||
releaseEmbedding();
|
||||
|
||||
expect((await pending)[0].text).toBe("hello before");
|
||||
} finally {
|
||||
releaseEmbedding();
|
||||
first.close();
|
||||
second.close();
|
||||
}
|
||||
});
|
||||
|
||||
it("refreshes a reused automatic search for every execution", async () => {
|
||||
const first = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
const second = await connect(tmpDir.name, { readConsistencyInterval: 0 });
|
||||
|
||||
try {
|
||||
const table = await first.createTable("docs", [
|
||||
{ text: "hello before", marker: "before" },
|
||||
]);
|
||||
await table.createIndex("text", { config: Index.fts() });
|
||||
const search = table.search("hello").select(["text"]);
|
||||
|
||||
const before = (await search.toArray())[0];
|
||||
expect(before.text).toBe("hello before");
|
||||
expect(before.marker).toBeUndefined();
|
||||
|
||||
const replacement = await second.createTable(
|
||||
"docs",
|
||||
[{ text: "hello after", marker: "after" }],
|
||||
{ mode: "overwrite" },
|
||||
);
|
||||
await replacement.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const after = (await search.toArray())[0];
|
||||
expect(after.text).toBe("hello after");
|
||||
expect(after.marker).toBeUndefined();
|
||||
} finally {
|
||||
first.close();
|
||||
second.close();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe("schema evolution", function () {
|
||||
let tmpDir: tmp.DirResult;
|
||||
beforeEach(() => {
|
||||
@@ -2344,7 +2585,24 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
);
|
||||
});
|
||||
|
||||
test("full text search if no embedding function provided", async () => {
|
||||
test("full text search if only an unrelated embedding function is registered", async () => {
|
||||
register("unused")(
|
||||
class extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 3;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(_data: string) {
|
||||
return [1, 2, 3];
|
||||
}
|
||||
async computeSourceEmbeddings(data: string[]) {
|
||||
return data.map(() => [1, 2, 3]);
|
||||
}
|
||||
},
|
||||
);
|
||||
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
{ text: "hello world", vector: [0.1, 0.2, 0.3] },
|
||||
@@ -2366,6 +2624,306 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
|
||||
expect(results2[0].text).toBe(data[1].text);
|
||||
});
|
||||
|
||||
test("auto search stays consistent with the active revision", async () => {
|
||||
let initCalls = 0;
|
||||
let queryCalls = 0;
|
||||
let markStarted!: () => void;
|
||||
const started = new Promise<void>((resolve) => {
|
||||
markStarted = resolve;
|
||||
});
|
||||
let releaseEmbedding!: () => void;
|
||||
const embeddingReleased = new Promise<void>((resolve) => {
|
||||
releaseEmbedding = resolve;
|
||||
});
|
||||
|
||||
@register("refresh-test")
|
||||
class TestEmbedding extends EmbeddingFunction<string> {
|
||||
async init() {
|
||||
initCalls += 1;
|
||||
}
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings(value: string) {
|
||||
queryCalls += 1;
|
||||
if (value === "blocked") {
|
||||
markStarted();
|
||||
await embeddingReleased;
|
||||
}
|
||||
return value === "greetings" ? [0.1] : [0.2];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map((value) =>
|
||||
value === "hello world" ? [0.1] : [0.2],
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
await writer.createTable("test", [{ text: "plain", vector: [0.0] }]);
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("test");
|
||||
type SnapshotCountingNative = {
|
||||
querySnapshot: () => Promise<unknown>;
|
||||
};
|
||||
const native = (tracked as unknown as { inner: SnapshotCountingNative })
|
||||
.inner;
|
||||
const querySnapshot = native.querySnapshot.bind(native);
|
||||
let snapshotCalls = 0;
|
||||
native.querySnapshot = async () => {
|
||||
snapshotCalls += 1;
|
||||
return await querySnapshot();
|
||||
};
|
||||
const autoQuery = tracked.search("greetings").select(["text"]).limit(1);
|
||||
|
||||
const func = new TestEmbedding();
|
||||
const schema = LanceSchema({
|
||||
text: func.sourceField(new arrow.Utf8()),
|
||||
vector: func.vectorField(),
|
||||
});
|
||||
const data = [{ text: "hello world" }, { text: "goodbye world" }];
|
||||
await writer.createTable("test", data, { mode: "overwrite", schema });
|
||||
const baselineInitCalls = initCalls;
|
||||
|
||||
expect(
|
||||
(await tracked.schema()).metadata.get("embedding_functions"),
|
||||
).toBeDefined();
|
||||
const results = await autoQuery.toArray();
|
||||
expect(results[0].text).toBe(data[0].text);
|
||||
expect(initCalls).toBe(baselineInitCalls + 1);
|
||||
expect(queryCalls).toBe(1);
|
||||
expect(snapshotCalls).toBe(1);
|
||||
|
||||
const repeatedResults = await autoQuery.toArray();
|
||||
expect(repeatedResults[0].text).toBe(data[0].text);
|
||||
expect(initCalls).toBe(baselineInitCalls + 1);
|
||||
expect(queryCalls).toBe(1);
|
||||
expect(snapshotCalls).toBe(2);
|
||||
|
||||
const pending = tracked
|
||||
.search("blocked")
|
||||
.select(["text"])
|
||||
.limit(1)
|
||||
.toArray();
|
||||
await started;
|
||||
|
||||
const ftsData = [
|
||||
{ text: "greetings from full text", vector: [0.0] },
|
||||
{ text: "blocked from full text", vector: [0.0] },
|
||||
];
|
||||
const ftsTable = await writer.createTable("test", ftsData, {
|
||||
mode: "overwrite",
|
||||
});
|
||||
await ftsTable.createIndex("text", { config: Index.fts() });
|
||||
releaseEmbedding();
|
||||
|
||||
const pendingResults = await pending;
|
||||
expect(pendingResults[0].text).toBe(data[1].text);
|
||||
|
||||
expect(
|
||||
(await tracked.schema()).metadata.get("embedding_functions"),
|
||||
).toBeUndefined();
|
||||
const ftsResults = await autoQuery.toArray();
|
||||
expect(ftsResults[0].text).toBe(ftsData[0].text);
|
||||
});
|
||||
|
||||
test("auto search keeps newer preparation during a revision race", async () => {
|
||||
let aCalls = 0;
|
||||
let bCalls = 0;
|
||||
let markAStarted!: () => void;
|
||||
const aStarted = new Promise<void>((resolve) => {
|
||||
markAStarted = resolve;
|
||||
});
|
||||
let releaseA!: () => void;
|
||||
const aReleased = new Promise<void>((resolve) => {
|
||||
releaseA = resolve;
|
||||
});
|
||||
let markBStarted!: () => void;
|
||||
const bStarted = new Promise<void>((resolve) => {
|
||||
markBStarted = resolve;
|
||||
});
|
||||
let releaseB!: () => void;
|
||||
const bReleased = new Promise<void>((resolve) => {
|
||||
releaseB = resolve;
|
||||
});
|
||||
|
||||
@register("race-a")
|
||||
class EmbeddingA extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
aCalls += 1;
|
||||
markAStarted();
|
||||
await aReleased;
|
||||
return [0.1];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.1]);
|
||||
}
|
||||
}
|
||||
|
||||
@register("race-b")
|
||||
class EmbeddingB extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
bCalls += 1;
|
||||
markBStarted();
|
||||
await bReleased;
|
||||
return [0.2];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.2]);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
const embeddingA = new EmbeddingA();
|
||||
const schemaA = LanceSchema({
|
||||
text: embeddingA.sourceField(new arrow.Utf8()),
|
||||
vector: embeddingA.vectorField(),
|
||||
});
|
||||
await writer.createTable("race", [{ text: "revision a" }], {
|
||||
schema: schemaA,
|
||||
});
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("race");
|
||||
const query = tracked.search("query");
|
||||
|
||||
const first = query.toArray();
|
||||
await aStarted;
|
||||
|
||||
const embeddingB = new EmbeddingB();
|
||||
const schemaB = LanceSchema({
|
||||
text: embeddingB.sourceField(new arrow.Utf8()),
|
||||
vector: embeddingB.vectorField(),
|
||||
});
|
||||
await writer.createTable("race", [{ text: "revision b" }], {
|
||||
mode: "overwrite",
|
||||
schema: schemaB,
|
||||
});
|
||||
const second = query.toArray();
|
||||
await bStarted;
|
||||
|
||||
releaseA();
|
||||
releaseB();
|
||||
await Promise.all([first, second]);
|
||||
expect(aCalls).toBe(1);
|
||||
expect(bCalls).toBe(1);
|
||||
});
|
||||
|
||||
test("stale FTS routing keeps newer vector preparation", async () => {
|
||||
let vectorCalls = 0;
|
||||
let markVectorStarted!: () => void;
|
||||
const vectorStarted = new Promise<void>((resolve) => {
|
||||
markVectorStarted = resolve;
|
||||
});
|
||||
let releaseVector!: () => void;
|
||||
const vectorReleased = new Promise<void>((resolve) => {
|
||||
releaseVector = resolve;
|
||||
});
|
||||
|
||||
@register("stale-fts-race")
|
||||
class RaceEmbedding extends EmbeddingFunction<string> {
|
||||
ndims() {
|
||||
return 1;
|
||||
}
|
||||
embeddingDataType() {
|
||||
return new arrow.Float32();
|
||||
}
|
||||
async computeQueryEmbeddings() {
|
||||
vectorCalls += 1;
|
||||
markVectorStarted();
|
||||
await vectorReleased;
|
||||
return [0.1];
|
||||
}
|
||||
async computeSourceEmbeddings(values: string[]) {
|
||||
return values.map(() => [0.1]);
|
||||
}
|
||||
}
|
||||
|
||||
const writer = await connect(tmpDir.name);
|
||||
const ftsTable = await writer.createTable("stale_fts", [
|
||||
{ text: "hello", vector: [0.0] },
|
||||
]);
|
||||
await ftsTable.createIndex("text", { config: Index.fts() });
|
||||
|
||||
const reader = await connect(tmpDir.name, {
|
||||
readConsistencyInterval: 0,
|
||||
});
|
||||
const tracked = await reader.openTable("stale_fts");
|
||||
type Snapshot = {
|
||||
schema: () => Promise<Buffer>;
|
||||
};
|
||||
type NativeWithSnapshot = {
|
||||
querySnapshot: () => Promise<Snapshot>;
|
||||
};
|
||||
const native = (tracked as unknown as { inner: NativeWithSnapshot })
|
||||
.inner;
|
||||
const querySnapshot = native.querySnapshot.bind(native);
|
||||
let snapshotCalls = 0;
|
||||
let markStaleSchemaStarted!: () => void;
|
||||
const staleSchemaStarted = new Promise<void>((resolve) => {
|
||||
markStaleSchemaStarted = resolve;
|
||||
});
|
||||
let releaseStaleSchema!: () => void;
|
||||
const staleSchemaReleased = new Promise<void>((resolve) => {
|
||||
releaseStaleSchema = resolve;
|
||||
});
|
||||
native.querySnapshot = async () => {
|
||||
const snapshot = await querySnapshot();
|
||||
snapshotCalls += 1;
|
||||
if (snapshotCalls === 1) {
|
||||
const schema = snapshot.schema.bind(snapshot);
|
||||
snapshot.schema = async () => {
|
||||
markStaleSchemaStarted();
|
||||
await staleSchemaReleased;
|
||||
return await schema();
|
||||
};
|
||||
}
|
||||
return snapshot;
|
||||
};
|
||||
|
||||
const query = tracked.search("hello");
|
||||
const staleFtsExecution = query.toArray();
|
||||
await staleSchemaStarted;
|
||||
|
||||
const embedding = new RaceEmbedding();
|
||||
const vectorSchema = LanceSchema({
|
||||
text: embedding.sourceField(new arrow.Utf8()),
|
||||
vector: embedding.vectorField(),
|
||||
});
|
||||
await writer.createTable("stale_fts", [{ text: "hello" }], {
|
||||
mode: "overwrite",
|
||||
schema: vectorSchema,
|
||||
});
|
||||
|
||||
const vectorExecution = query.toArray();
|
||||
await vectorStarted;
|
||||
releaseStaleSchema();
|
||||
await staleFtsExecution;
|
||||
releaseVector();
|
||||
await vectorExecution;
|
||||
|
||||
await query.toArray();
|
||||
expect(vectorCalls).toBe(1);
|
||||
});
|
||||
|
||||
test("tokenizes FTS queries by column or index name", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [
|
||||
@@ -2916,6 +3474,30 @@ describe("column name options", () => {
|
||||
expect(results[1].query_index).toBe(1);
|
||||
});
|
||||
|
||||
test("observes promised additional vectors while the query is pending", async () => {
|
||||
const initialVector = new Promise<number[]>(() => undefined);
|
||||
const query = table.query().nearestTo(initialVector);
|
||||
const unhandled: unknown[] = [];
|
||||
const onUnhandled = (reason: unknown) => unhandled.push(reason);
|
||||
process.on("unhandledRejection", onUnhandled);
|
||||
|
||||
try {
|
||||
query.addQueryVector(Promise.reject(new Error("extra vector failed")));
|
||||
await new Promise<void>((resolve) => setImmediate(resolve));
|
||||
expect(unhandled).toEqual([]);
|
||||
|
||||
const rejectedQuery = table
|
||||
.query()
|
||||
.nearestTo([0.1, 0.2])
|
||||
.addQueryVector(Promise.reject(new Error("consumed vector failed")));
|
||||
await expect(rejectedQuery.toArray()).rejects.toThrow(
|
||||
"consumed vector failed",
|
||||
);
|
||||
} finally {
|
||||
process.off("unhandledRejection", onUnhandled);
|
||||
}
|
||||
});
|
||||
|
||||
test("index and search multivectors", async () => {
|
||||
const db = await connect(tmpDir.name);
|
||||
const data = [];
|
||||
@@ -2979,6 +3561,27 @@ describe("when creating an empty table", () => {
|
||||
expect((actualSchema.fields[1].type as Float64).precision).toBe(2);
|
||||
});
|
||||
|
||||
it("can add and query JSON data", async () => {
|
||||
const schema = new Schema([
|
||||
new Field("id", new Int32(), true),
|
||||
new Field(
|
||||
"meta",
|
||||
new Utf8(),
|
||||
true,
|
||||
new Map([["ARROW:extension:name", "arrow.json"]]),
|
||||
),
|
||||
]);
|
||||
const table = await con.createEmptyTable("json", schema);
|
||||
const meta = JSON.stringify({ x: 1 });
|
||||
|
||||
await table.add([{ id: 1, meta }]);
|
||||
|
||||
const rows = await table.query().toArray();
|
||||
expect(rows).toHaveLength(1);
|
||||
expect(rows[0].id).toBe(1);
|
||||
expect(rows[0].meta).toBe(meta);
|
||||
});
|
||||
|
||||
it("can create an empty table from schema that specifies field types by name", async () => {
|
||||
const schemaLike = {
|
||||
fields: [
|
||||
|
||||
@@ -170,7 +170,7 @@ test("basic table examples", async () => {
|
||||
// --8<-- [end:create_index]
|
||||
|
||||
// --8<-- [start:delete_rows]
|
||||
await tbl.delete('item = "fizz"');
|
||||
await tbl.delete("item = 'fizz'");
|
||||
// --8<-- [end:delete_rows]
|
||||
|
||||
// --8<-- [start:drop_table]
|
||||
|
||||
@@ -8,7 +8,8 @@
|
||||
"//1": "--experimental-vm-modules is needed to run jest with sentence-transformers",
|
||||
"//2": "--testEnvironment is needed to run jest with sentence-transformers",
|
||||
"//3": "See: https://github.com/huggingface/transformers.js/issues/57",
|
||||
"test": "node --experimental-vm-modules node_modules/.bin/jest --testEnvironment jest-environment-node-single-context --verbose",
|
||||
"//4": "jest is invoked by its JS entry, not node_modules/.bin/jest: under pnpm that path is a shell shim, which `node` cannot execute",
|
||||
"test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --testEnvironment jest-environment-node-single-context --verbose",
|
||||
"lint": "biome check *.ts && biome format *.ts",
|
||||
"lint-ci": "biome ci .",
|
||||
"lint-fix": "biome check --write *.ts && pnpm format",
|
||||
|
||||
+35
-309
@@ -5,7 +5,6 @@ import {
|
||||
Data as ArrowData,
|
||||
Table as ArrowTable,
|
||||
Binary,
|
||||
Bool,
|
||||
BufferType,
|
||||
DataType,
|
||||
DateUnit,
|
||||
@@ -18,12 +17,7 @@ import {
|
||||
FixedSizeList,
|
||||
Float,
|
||||
Float32,
|
||||
Float64,
|
||||
Int,
|
||||
Int8,
|
||||
Int16,
|
||||
Int32,
|
||||
Int64,
|
||||
LargeBinary,
|
||||
List,
|
||||
Null,
|
||||
@@ -36,17 +30,16 @@ import {
|
||||
Struct,
|
||||
Timestamp,
|
||||
Type,
|
||||
Uint8,
|
||||
Uint16,
|
||||
Uint32,
|
||||
Utf8,
|
||||
Vector,
|
||||
makeVector as arrowMakeVector,
|
||||
util as arrowUtil,
|
||||
vectorFromArray as badVectorFromArray,
|
||||
makeBuilder,
|
||||
makeData,
|
||||
} from "apache-arrow";
|
||||
import { Buffers } from "apache-arrow/data";
|
||||
import { typedArrayToArrowType } from "./arrow_type";
|
||||
import { type EmbeddingFunction } from "./embedding/embedding_function";
|
||||
import {
|
||||
EmbeddingFunctionConfig,
|
||||
@@ -59,14 +52,7 @@ import {
|
||||
sanitizeTable,
|
||||
sanitizeType,
|
||||
} from "./sanitize";
|
||||
|
||||
/**
|
||||
* Check if a field name indicates a vector column.
|
||||
*/
|
||||
function nameSuggestsVectorColumn(fieldName: string): boolean {
|
||||
const nameLower = fieldName.toLowerCase();
|
||||
return nameLower.includes("vector") || nameLower.includes("embedding");
|
||||
}
|
||||
import { inferSchema } from "./schema";
|
||||
|
||||
export * from "apache-arrow";
|
||||
export type SchemaLike =
|
||||
@@ -86,8 +72,7 @@ export type FieldLike =
|
||||
};
|
||||
|
||||
export type DataLike =
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
| import("apache-arrow").Data<Struct<any>>
|
||||
| import("apache-arrow").Data
|
||||
| {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
type: any;
|
||||
@@ -96,6 +81,7 @@ export type DataLike =
|
||||
stride: number;
|
||||
nullable: boolean;
|
||||
children: DataLike[];
|
||||
dictionary?: { data: readonly DataLike[] };
|
||||
get nullCount(): number;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: <explanation>
|
||||
values: Buffers<any>[BufferType.DATA];
|
||||
@@ -459,110 +445,6 @@ export function makeArrowTable(
|
||||
return new ArrowTable(inferredSchema, finalColumns);
|
||||
}
|
||||
|
||||
function inferSchema(
|
||||
data: Array<Record<string, unknown>>,
|
||||
schema: Schema | undefined,
|
||||
opts: MakeArrowTableOptions,
|
||||
): Schema {
|
||||
// We will collect all fields we see in the data.
|
||||
const pathTree = new PathTree<DataType>();
|
||||
|
||||
for (const [rowI, row] of data.entries()) {
|
||||
for (const [path, value] of rowPathsAndValues(row)) {
|
||||
if (!pathTree.has(path)) {
|
||||
// First time seeing this field.
|
||||
if (schema !== undefined) {
|
||||
const field = getFieldForPath(schema, path);
|
||||
if (field === undefined) {
|
||||
throw new Error(
|
||||
`Found field not in schema: ${path.join(".")} at row ${rowI}`,
|
||||
);
|
||||
} else {
|
||||
pathTree.set(path, field.type);
|
||||
}
|
||||
} else {
|
||||
const inferredType = inferType(value, path, opts);
|
||||
if (inferredType === undefined) {
|
||||
throw new Error(`Failed to infer data type for field ${path.join(
|
||||
".",
|
||||
)} at row ${rowI}. \
|
||||
Consider providing an explicit schema.`);
|
||||
}
|
||||
pathTree.set(path, inferredType);
|
||||
}
|
||||
} else if (schema === undefined) {
|
||||
const currentType = pathTree.get(path);
|
||||
const newType = inferType(value, path, opts);
|
||||
if (currentType !== newType) {
|
||||
new Error(`Failed to infer schema for data. Previously inferred type \
|
||||
${currentType} but found ${newType} at row ${rowI}. Consider \
|
||||
providing an explicit schema.`);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (schema === undefined) {
|
||||
function fieldsFromPathTree(pathTree: PathTree<DataType>): Field[] {
|
||||
const fields = [];
|
||||
for (const [name, value] of pathTree.map.entries()) {
|
||||
if (value instanceof PathTree) {
|
||||
const children = fieldsFromPathTree(value);
|
||||
fields.push(new Field(name, new Struct(children), true));
|
||||
} else {
|
||||
fields.push(new Field(name, value, true));
|
||||
}
|
||||
}
|
||||
return fields;
|
||||
}
|
||||
const fields = fieldsFromPathTree(pathTree);
|
||||
return new Schema(fields);
|
||||
} else {
|
||||
function takeMatchingFields(
|
||||
fields: Field[],
|
||||
pathTree: PathTree<DataType>,
|
||||
): Field[] {
|
||||
const outFields = [];
|
||||
for (const field of fields) {
|
||||
if (pathTree.map.has(field.name)) {
|
||||
const value = pathTree.get([field.name]);
|
||||
if (value instanceof PathTree) {
|
||||
const struct = field.type as Struct;
|
||||
const children = takeMatchingFields(struct.children, value);
|
||||
outFields.push(
|
||||
new Field(field.name, new Struct(children), field.nullable),
|
||||
);
|
||||
} else {
|
||||
outFields.push(
|
||||
new Field(field.name, value as DataType, field.nullable),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
return outFields;
|
||||
}
|
||||
const fields = takeMatchingFields(schema.fields, pathTree);
|
||||
return new Schema(fields);
|
||||
}
|
||||
}
|
||||
|
||||
function* rowPathsAndValues(
|
||||
row: Record<string, unknown>,
|
||||
basePath: string[] = [],
|
||||
): Generator<[string[], unknown]> {
|
||||
for (const [key, value] of Object.entries(row)) {
|
||||
if (isObject(value)) {
|
||||
yield* rowPathsAndValues(value, [...basePath, key]);
|
||||
} else {
|
||||
// Skip undefined values - they should be treated the same as missing fields
|
||||
// for embedding function purposes
|
||||
if (value !== undefined) {
|
||||
yield [[...basePath, key], value];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function isObject(value: unknown): value is Record<string, unknown> {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
@@ -577,146 +459,19 @@ function isObject(value: unknown): value is Record<string, unknown> {
|
||||
);
|
||||
}
|
||||
|
||||
function getFieldForPath(schema: Schema, path: string[]): Field | undefined {
|
||||
let current: Field | Schema = schema;
|
||||
function valueAtPath(datum: Record<string, unknown>, path: string[]): unknown {
|
||||
let current: unknown = datum;
|
||||
for (const key of path) {
|
||||
if (current instanceof Schema) {
|
||||
const field: Field | undefined = current.fields.find(
|
||||
(f) => f.name === key,
|
||||
);
|
||||
if (field === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
current = field;
|
||||
} else if (current instanceof Field && DataType.isStruct(current.type)) {
|
||||
const struct: Struct = current.type;
|
||||
const field = struct.children.find((f) => f.name === key);
|
||||
if (field === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
current = field;
|
||||
if (current == null) {
|
||||
return null;
|
||||
}
|
||||
if (isObject(current) && (Object.hasOwn(current, key) || key in current)) {
|
||||
current = current[key];
|
||||
} else {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
if (current instanceof Field) {
|
||||
return current;
|
||||
} else {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Try to infer which Arrow type to use for a given value.
|
||||
*
|
||||
* May return undefined if the type cannot be inferred.
|
||||
*/
|
||||
function inferType(
|
||||
value: unknown,
|
||||
path: string[],
|
||||
opts: MakeArrowTableOptions,
|
||||
): DataType | undefined {
|
||||
if (typeof value === "bigint") {
|
||||
return new Int64();
|
||||
} else if (typeof value === "number") {
|
||||
// Even if it's an integer, it's safer to assume Float64. Users can
|
||||
// always provide an explicit schema or use BigInt if they mean integer.
|
||||
return new Float64();
|
||||
} else if (typeof value === "string") {
|
||||
if (opts.dictionaryEncodeStrings) {
|
||||
return new Dictionary(new Utf8(), new Int32());
|
||||
} else {
|
||||
return new Utf8();
|
||||
}
|
||||
} else if (typeof value === "boolean") {
|
||||
return new Bool();
|
||||
} else if (value instanceof Buffer) {
|
||||
return new Binary();
|
||||
} else if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
|
||||
const info = typedArrayToArrowType(value);
|
||||
if (info !== undefined) {
|
||||
const child = new Field("item", info.elementType, true);
|
||||
return new FixedSizeList(info.length, child);
|
||||
}
|
||||
return undefined;
|
||||
} else if (Array.isArray(value)) {
|
||||
if (value.length === 0) {
|
||||
return undefined; // Without any values we can't infer the type
|
||||
}
|
||||
if (path.length === 1 && Object.hasOwn(opts.vectorColumns, path[0])) {
|
||||
const floatType = sanitizeType(opts.vectorColumns[path[0]].type);
|
||||
return new FixedSizeList(
|
||||
value.length,
|
||||
new Field("item", floatType, true),
|
||||
);
|
||||
}
|
||||
const valueType = inferType(value[0], path, opts);
|
||||
if (valueType === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
// Try to automatically detect embedding columns.
|
||||
if (nameSuggestsVectorColumn(path[path.length - 1])) {
|
||||
// Check if value is a Uint8Array for integer vector type determination
|
||||
if (value instanceof Uint8Array) {
|
||||
// For integer vectors, we default to Uint8 (matching Python implementation)
|
||||
const child = new Field("item", new Uint8(), true);
|
||||
return new FixedSizeList(value.length, child);
|
||||
} else {
|
||||
// For float vectors, we default to Float32
|
||||
const child = new Field("item", new Float32(), true);
|
||||
return new FixedSizeList(value.length, child);
|
||||
}
|
||||
} else {
|
||||
const child = new Field("item", valueType, true);
|
||||
return new List(child);
|
||||
}
|
||||
} else {
|
||||
// TODO: timestamp
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
class PathTree<V> {
|
||||
map: Map<string, V | PathTree<V>>;
|
||||
|
||||
constructor(entries?: [string[], V][]) {
|
||||
this.map = new Map();
|
||||
if (entries !== undefined) {
|
||||
for (const [path, value] of entries) {
|
||||
this.set(path, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
has(path: string[]): boolean {
|
||||
let ref: PathTree<V> = this;
|
||||
for (const part of path) {
|
||||
if (!(ref instanceof PathTree) || !ref.map.has(part)) {
|
||||
return false;
|
||||
}
|
||||
ref = ref.map.get(part) as PathTree<V>;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
get(path: string[]): V | undefined {
|
||||
let ref: PathTree<V> = this;
|
||||
for (const part of path) {
|
||||
if (!(ref instanceof PathTree) || !ref.map.has(part)) {
|
||||
return undefined;
|
||||
}
|
||||
ref = ref.map.get(part) as PathTree<V>;
|
||||
}
|
||||
return ref as V;
|
||||
}
|
||||
set(path: string[], value: V): void {
|
||||
let ref: PathTree<V> = this;
|
||||
for (const part of path.slice(0, path.length - 1)) {
|
||||
if (!ref.map.has(part)) {
|
||||
ref.map.set(part, new PathTree<V>());
|
||||
}
|
||||
ref = ref.map.get(part) as PathTree<V>;
|
||||
}
|
||||
ref.map.set(path[path.length - 1], value);
|
||||
}
|
||||
return current;
|
||||
}
|
||||
|
||||
function transposeData(
|
||||
@@ -724,37 +479,26 @@ function transposeData(
|
||||
field: Field,
|
||||
path: string[] = [],
|
||||
): Vector {
|
||||
const valuesPath = [...path, field.name];
|
||||
const values = data.map((datum) => valueAtPath(datum, valuesPath));
|
||||
if (field.type instanceof Struct) {
|
||||
const childFields = field.type.children;
|
||||
const fullPath = [...path, field.name];
|
||||
const childVectors = childFields.map((child) => {
|
||||
return transposeData(data, child, fullPath);
|
||||
return transposeData(data, child, valuesPath);
|
||||
});
|
||||
const nullCount = values.filter((value) => value === null).length;
|
||||
const structData = makeData({
|
||||
type: field.type,
|
||||
length: values.length,
|
||||
nullCount,
|
||||
nullBitmap:
|
||||
nullCount > 0
|
||||
? arrowUtil.packBools(values.map((value) => value !== null))
|
||||
: undefined,
|
||||
children: childVectors as unknown as ArrowData<DataType>[],
|
||||
});
|
||||
return arrowMakeVector(structData);
|
||||
} else {
|
||||
const valuesPath = [...path, field.name];
|
||||
const values = data.map((datum) => {
|
||||
let current: unknown = datum;
|
||||
for (const key of valuesPath) {
|
||||
if (current == null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (
|
||||
isObject(current) &&
|
||||
(Object.hasOwn(current, key) || key in current)
|
||||
) {
|
||||
current = current[key];
|
||||
} else {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
return current;
|
||||
});
|
||||
return makeVector(values, field.type, undefined, field.nullable);
|
||||
}
|
||||
}
|
||||
@@ -797,32 +541,6 @@ function makeListVector(lists: unknown[][]): Vector<unknown> {
|
||||
return listBuilder.finish().toVector();
|
||||
}
|
||||
|
||||
/**
|
||||
* Map a JS TypedArray instance to the corresponding Arrow element DataType
|
||||
* and its length. Returns undefined if the value is not a recognized TypedArray.
|
||||
*/
|
||||
function typedArrayToArrowType(
|
||||
value: ArrayBufferView,
|
||||
): { elementType: DataType; length: number } | undefined {
|
||||
if (value instanceof Float32Array)
|
||||
return { elementType: new Float32(), length: value.length };
|
||||
if (value instanceof Float64Array)
|
||||
return { elementType: new Float64(), length: value.length };
|
||||
if (value instanceof Uint8Array)
|
||||
return { elementType: new Uint8(), length: value.length };
|
||||
if (value instanceof Uint16Array)
|
||||
return { elementType: new Uint16(), length: value.length };
|
||||
if (value instanceof Uint32Array)
|
||||
return { elementType: new Uint32(), length: value.length };
|
||||
if (value instanceof Int8Array)
|
||||
return { elementType: new Int8(), length: value.length };
|
||||
if (value instanceof Int16Array)
|
||||
return { elementType: new Int16(), length: value.length };
|
||||
if (value instanceof Int32Array)
|
||||
return { elementType: new Int32(), length: value.length };
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/** Helper function to convert an Array of JS values to an Arrow Vector */
|
||||
function makeVector(
|
||||
values: unknown[],
|
||||
@@ -1462,8 +1180,12 @@ export function ensureNestedFieldsExist(
|
||||
completeRow[field.name] = row[field.name];
|
||||
}
|
||||
} else {
|
||||
// Field is missing from the data - set to null
|
||||
completeRow[field.name] = null;
|
||||
// Keep a missing struct valid while filling each of its children with
|
||||
// null. This is distinct from an explicitly null struct value.
|
||||
completeRow[field.name] =
|
||||
field.type.constructor.name === "Struct"
|
||||
? ensureStructFieldsExist({}, field.type as Struct)
|
||||
: null;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1498,8 +1220,12 @@ function ensureStructFieldsExist(
|
||||
completeStruct[childField.name] = data[childField.name];
|
||||
}
|
||||
} else {
|
||||
// Field is missing - set to null
|
||||
completeStruct[childField.name] = null;
|
||||
// Keep a missing struct valid while filling each of its children with
|
||||
// null. This is distinct from an explicitly null struct value.
|
||||
completeStruct[childField.name] =
|
||||
childField.type.constructor.name === "Struct"
|
||||
? ensureStructFieldsExist({}, childField.type as Struct)
|
||||
: null;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
type DataType,
|
||||
Float32,
|
||||
Float64,
|
||||
Int8,
|
||||
Int16,
|
||||
Int32,
|
||||
Uint8,
|
||||
Uint16,
|
||||
Uint32,
|
||||
} from "apache-arrow";
|
||||
|
||||
/**
|
||||
* Map a JS TypedArray instance to the corresponding Arrow element type and
|
||||
* length. Returns undefined when the view is not a supported TypedArray.
|
||||
*/
|
||||
export function typedArrayToArrowType(
|
||||
value: ArrayBufferView,
|
||||
): { elementType: DataType; length: number } | undefined {
|
||||
if (value instanceof Float32Array)
|
||||
return { elementType: new Float32(), length: value.length };
|
||||
if (value instanceof Float64Array)
|
||||
return { elementType: new Float64(), length: value.length };
|
||||
if (value instanceof Uint8Array)
|
||||
return { elementType: new Uint8(), length: value.length };
|
||||
if (value instanceof Uint16Array)
|
||||
return { elementType: new Uint16(), length: value.length };
|
||||
if (value instanceof Uint32Array)
|
||||
return { elementType: new Uint32(), length: value.length };
|
||||
if (value instanceof Int8Array)
|
||||
return { elementType: new Int8(), length: value.length };
|
||||
if (value instanceof Int16Array)
|
||||
return { elementType: new Int16(), length: value.length };
|
||||
if (value instanceof Int32Array)
|
||||
return { elementType: new Int32(), length: value.length };
|
||||
return undefined;
|
||||
}
|
||||
@@ -31,12 +31,14 @@ import type {
|
||||
JobDescription,
|
||||
JobInfo,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
} from "./native";
|
||||
export type {
|
||||
CreateNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
};
|
||||
import { sanitizeTable } from "./sanitize";
|
||||
import { LocalTable, Table } from "./table";
|
||||
@@ -134,6 +136,10 @@ export interface OpenTableOptions {
|
||||
indexCacheSize?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Use {@link ListTablesOptions} with {@link Connection.listTables}
|
||||
* instead.
|
||||
*/
|
||||
export interface TableNamesOptions {
|
||||
/**
|
||||
* If present, only return names that come lexicographically after the
|
||||
@@ -147,6 +153,24 @@ export interface TableNamesOptions {
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface ListTablesOptions {
|
||||
/**
|
||||
* Token from a previous response, to resume listing where it left off.
|
||||
*
|
||||
* The token is opaque: it carries whatever the database needs to resume, and
|
||||
* callers should not construct or interpret one.
|
||||
*/
|
||||
pageToken?: string;
|
||||
/**
|
||||
* An upper bound on how many tables to return.
|
||||
*
|
||||
* A page may hold fewer than this and still not be the last one, so keep
|
||||
* going while the response carries a page token rather than while pages are
|
||||
* full.
|
||||
*/
|
||||
limit?: number;
|
||||
}
|
||||
|
||||
export interface ListNamespacesOptions {
|
||||
/** Token from a previous response for pagination. */
|
||||
pageToken?: string;
|
||||
@@ -231,6 +255,7 @@ export abstract class Connection {
|
||||
* @param {Partial<TableNamesOptions>} options - options to control the
|
||||
* paging / start point (backwards compatibility)
|
||||
*
|
||||
* @deprecated Use {@link Connection.listTables} instead.
|
||||
*/
|
||||
abstract tableNames(options?: Partial<TableNamesOptions>): Promise<string[]>;
|
||||
/**
|
||||
@@ -241,12 +266,53 @@ export abstract class Connection {
|
||||
* @param {Partial<TableNamesOptions>} options - options to control the
|
||||
* paging / start point
|
||||
*
|
||||
* @deprecated Use {@link Connection.listTables} instead.
|
||||
*/
|
||||
abstract tableNames(
|
||||
namespacePath?: string[],
|
||||
options?: Partial<TableNamesOptions>,
|
||||
): Promise<string[]>;
|
||||
|
||||
/**
|
||||
* List a page of the tables in this database.
|
||||
*
|
||||
* To retrieve the tables after the page, pass the `pageToken` the response
|
||||
* carries back in. A page can be shorter than `limit` without being the last
|
||||
* one, so walk until a response carries no page token:
|
||||
*
|
||||
* ```ts
|
||||
* const names = [];
|
||||
* let pageToken = undefined;
|
||||
* do {
|
||||
* const page = await conn.listTables({ pageToken, limit: 100 });
|
||||
* names.push(...page.tables);
|
||||
* pageToken = page.pageToken;
|
||||
* } while (pageToken);
|
||||
* ```
|
||||
*
|
||||
* @param {Partial<ListTablesOptions>} options - Pagination options
|
||||
* (`pageToken`, `limit`).
|
||||
* @returns {Promise<ListTablesResponse>} A page of table names and an
|
||||
* optional token for the tables after it.
|
||||
*/
|
||||
abstract listTables(
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse>;
|
||||
/**
|
||||
* List a page of the tables in this database.
|
||||
*
|
||||
* @param {string[]} namespacePath - The namespace path to list tables from
|
||||
* (defaults to root namespace)
|
||||
* @param {Partial<ListTablesOptions>} options - Pagination options
|
||||
* (`pageToken`, `limit`).
|
||||
* @returns {Promise<ListTablesResponse>} A page of table names and an
|
||||
* optional token for the tables after it.
|
||||
*/
|
||||
abstract listTables(
|
||||
namespacePath?: string[],
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse>;
|
||||
|
||||
/**
|
||||
* Open a table in the database.
|
||||
* @param {string} name - The name of the table
|
||||
@@ -601,6 +667,25 @@ export class LocalConnection extends Connection {
|
||||
return await this.inner.listMaterializedViews();
|
||||
}
|
||||
|
||||
async listTables(
|
||||
namespacePathOrOptions?: string[] | Partial<ListTablesOptions>,
|
||||
options?: Partial<ListTablesOptions>,
|
||||
): Promise<ListTablesResponse> {
|
||||
// Detect if first argument is namespacePath array or options object
|
||||
const namespacePath = Array.isArray(namespacePathOrOptions)
|
||||
? namespacePathOrOptions
|
||||
: undefined;
|
||||
const listTablesOptions = Array.isArray(namespacePathOrOptions)
|
||||
? options
|
||||
: namespacePathOrOptions;
|
||||
|
||||
return this.inner.listTables(
|
||||
namespacePath ?? [],
|
||||
listTablesOptions?.pageToken,
|
||||
listTablesOptions?.limit,
|
||||
);
|
||||
}
|
||||
|
||||
async openTable(
|
||||
name: string,
|
||||
namespacePath?: string[],
|
||||
|
||||
@@ -4,7 +4,15 @@
|
||||
import { Field, Schema } from "../arrow";
|
||||
import { sanitizeType } from "../sanitize";
|
||||
import { EmbeddingFunction } from "./embedding_function";
|
||||
import { EmbeddingFunctionConfig, getRegistry } from "./registry";
|
||||
import {
|
||||
EmbeddingFunctionConfig,
|
||||
EmbeddingFunctionRegistry,
|
||||
getRegistry as getGlobalRegistry,
|
||||
registerBuiltIn,
|
||||
} from "./registry";
|
||||
|
||||
type OpenAIModule = typeof import("./openai");
|
||||
type TransformersModule = typeof import("./transformers");
|
||||
|
||||
export {
|
||||
FieldOptions,
|
||||
@@ -14,7 +22,39 @@ export {
|
||||
EmbeddingFunctionConstructor,
|
||||
} from "./embedding_function";
|
||||
|
||||
export * from "./registry";
|
||||
export {
|
||||
EmbeddingFunctionRegistry,
|
||||
parseEmbeddingMetadata,
|
||||
register,
|
||||
} from "./registry";
|
||||
export type {
|
||||
CreateReturnType,
|
||||
EmbeddingFunctionConfig,
|
||||
EmbeddingFunctionCreate,
|
||||
EmbeddingMetadataEntry,
|
||||
ResolvedEmbeddingFunctionConfig,
|
||||
} from "./registry";
|
||||
|
||||
function initializeBuiltInProviders() {
|
||||
const { OpenAIEmbeddingFunction } = require("./openai") as OpenAIModule;
|
||||
const { TransformersEmbeddingFunction } =
|
||||
require("./transformers") as TransformersModule;
|
||||
|
||||
registerBuiltIn("openai", OpenAIEmbeddingFunction);
|
||||
registerBuiltIn("huggingface", TransformersEmbeddingFunction);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the global embedding function registry.
|
||||
*
|
||||
* LanceDB built-in providers are initialized when this public API is first
|
||||
* used, so importing the root package does not change automatic search
|
||||
* selection for tables without embedding metadata.
|
||||
*/
|
||||
export function getRegistry(): EmbeddingFunctionRegistry {
|
||||
initializeBuiltInProviders();
|
||||
return getGlobalRegistry();
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a schema with embedding functions.
|
||||
|
||||
@@ -5,14 +5,13 @@ import type OpenAI from "openai";
|
||||
import type { EmbeddingCreateParams } from "openai/resources/index";
|
||||
import { Float, Float32 } from "../arrow";
|
||||
import { EmbeddingFunction } from "./embedding_function";
|
||||
import { register } from "./registry";
|
||||
import { registerBuiltIn } from "./registry";
|
||||
|
||||
export type OpenAIOptions = {
|
||||
apiKey: string;
|
||||
model: EmbeddingCreateParams["model"];
|
||||
};
|
||||
|
||||
@register("openai")
|
||||
export class OpenAIEmbeddingFunction extends EmbeddingFunction<
|
||||
string,
|
||||
Partial<OpenAIOptions>
|
||||
@@ -100,3 +99,5 @@ export class OpenAIEmbeddingFunction extends EmbeddingFunction<
|
||||
return response.data[0].embedding;
|
||||
}
|
||||
}
|
||||
|
||||
registerBuiltIn("openai", OpenAIEmbeddingFunction);
|
||||
|
||||
@@ -7,6 +7,10 @@ import {
|
||||
} from "./embedding_function";
|
||||
import "reflect-metadata";
|
||||
|
||||
const builtInFunctionsKey = Symbol.for(
|
||||
"@lancedb/lancedb::embedding-built-in-functions::v1",
|
||||
);
|
||||
|
||||
export type CreateReturnType<T> = T extends { init: () => Promise<void> }
|
||||
? Promise<T>
|
||||
: T;
|
||||
@@ -59,6 +63,15 @@ export class EmbeddingFunctionRegistry {
|
||||
};
|
||||
}
|
||||
|
||||
/** @ignore */
|
||||
setBuiltIn<
|
||||
T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor,
|
||||
>(name: string, ctor: T): T {
|
||||
this.#functions.set(name, ctor);
|
||||
Reflect.defineMetadata("lancedb::embedding::name", name, ctor);
|
||||
return ctor;
|
||||
}
|
||||
|
||||
get<T extends EmbeddingFunction<unknown>>(
|
||||
name: string,
|
||||
): EmbeddingFunctionCreate<T> | undefined;
|
||||
@@ -96,6 +109,7 @@ export class EmbeddingFunctionRegistry {
|
||||
*/
|
||||
reset(this: EmbeddingFunctionRegistry) {
|
||||
this.#functions.clear();
|
||||
getBuiltInFunctions(this).clear();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -183,12 +197,56 @@ export class EmbeddingFunctionRegistry {
|
||||
}
|
||||
}
|
||||
|
||||
const _REGISTRY = new EmbeddingFunctionRegistry();
|
||||
function getBuiltInFunctions(registry: EmbeddingFunctionRegistry): Set<string> {
|
||||
const registryWithBuiltIns = registry as EmbeddingFunctionRegistry & {
|
||||
[key: symbol]: Set<string> | undefined;
|
||||
};
|
||||
let builtInFunctions = registryWithBuiltIns[builtInFunctionsKey];
|
||||
if (builtInFunctions === undefined) {
|
||||
builtInFunctions = new Set<string>();
|
||||
registryWithBuiltIns[builtInFunctionsKey] = builtInFunctions;
|
||||
}
|
||||
return builtInFunctions;
|
||||
}
|
||||
|
||||
// Server bundlers can load the side-effect embedding entry points and the public
|
||||
// embedding API from separate module graphs. Keep their registry shared.
|
||||
const registryKey = Symbol.for(
|
||||
"@lancedb/lancedb::embedding-function-registry::v1",
|
||||
);
|
||||
const registryGlobal = globalThis as typeof globalThis & {
|
||||
[key: symbol]: EmbeddingFunctionRegistry | undefined;
|
||||
};
|
||||
|
||||
function getGlobalRegistry(): EmbeddingFunctionRegistry {
|
||||
const existingRegistry = registryGlobal[registryKey];
|
||||
if (existingRegistry !== undefined) {
|
||||
return existingRegistry;
|
||||
}
|
||||
const registry = new EmbeddingFunctionRegistry();
|
||||
registryGlobal[registryKey] = registry;
|
||||
return registry;
|
||||
}
|
||||
|
||||
const _REGISTRY = getGlobalRegistry();
|
||||
|
||||
export function register(name?: string) {
|
||||
return _REGISTRY.register(name);
|
||||
}
|
||||
|
||||
/** @ignore */
|
||||
export function registerBuiltIn<
|
||||
T extends EmbeddingFunctionConstructor = EmbeddingFunctionConstructor,
|
||||
>(name: string, ctor: T): T {
|
||||
const builtInFunctions = getBuiltInFunctions(_REGISTRY);
|
||||
if (builtInFunctions.has(name)) {
|
||||
return _REGISTRY.setBuiltIn(name, ctor);
|
||||
}
|
||||
_REGISTRY.register(name)(ctor);
|
||||
builtInFunctions.add(name);
|
||||
return ctor;
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility function to get the global instance of the registry
|
||||
* @returns `EmbeddingFunctionRegistry` The global instance of the registry
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
|
||||
import { Float, Float32 } from "../arrow";
|
||||
import { EmbeddingFunction } from "./embedding_function";
|
||||
import { register } from "./registry";
|
||||
import { registerBuiltIn } from "./registry";
|
||||
|
||||
export type XenovaTransformerOptions = {
|
||||
/** The wasm compatible model to use */
|
||||
@@ -31,7 +31,6 @@ export type XenovaTransformerOptions = {
|
||||
};
|
||||
};
|
||||
|
||||
@register("huggingface")
|
||||
export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
string,
|
||||
Partial<XenovaTransformerOptions>
|
||||
@@ -158,6 +157,8 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
}
|
||||
}
|
||||
|
||||
registerBuiltIn("huggingface", TransformersEmbeddingFunction);
|
||||
|
||||
const tensorDiv = (
|
||||
src: import("@huggingface/transformers").Tensor,
|
||||
divBy: number,
|
||||
|
||||
@@ -81,11 +81,13 @@ export {
|
||||
Connection,
|
||||
CreateTableOptions,
|
||||
TableNamesOptions,
|
||||
ListTablesOptions,
|
||||
OpenTableOptions,
|
||||
ListNamespacesOptions,
|
||||
CreateNamespaceOptions,
|
||||
DropNamespaceOptions,
|
||||
ListNamespacesResponse,
|
||||
ListTablesResponse,
|
||||
CreateNamespaceResponse,
|
||||
DropNamespaceResponse,
|
||||
DescribeNamespaceResponse,
|
||||
@@ -101,6 +103,7 @@ export {
|
||||
} from "./native.js";
|
||||
|
||||
export {
|
||||
AutoQuery,
|
||||
ExecutableQuery,
|
||||
Query,
|
||||
QueryBase,
|
||||
|
||||
+205
-106
@@ -100,6 +100,29 @@ export interface FullTextSearchOptions {
|
||||
columns?: string | string[];
|
||||
}
|
||||
|
||||
function nearestToNative(
|
||||
inner: NativeQuery,
|
||||
vector: Awaited<IntoVector>,
|
||||
): NativeVectorQuery {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
return inner.nearestToRaw(raw.data, raw.dtype);
|
||||
}
|
||||
return inner.nearestTo(Float32Array.from(vector as number[]));
|
||||
}
|
||||
|
||||
function addQueryVectorToNative(
|
||||
inner: NativeVectorQuery,
|
||||
vector: Awaited<IntoVector>,
|
||||
) {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
inner.addQueryVectorRaw(raw.data, raw.dtype);
|
||||
} else {
|
||||
inner.addQueryVector(Float32Array.from(vector as number[]));
|
||||
}
|
||||
}
|
||||
|
||||
/** Common methods supported by all query types
|
||||
*
|
||||
* @see {@link Query}
|
||||
@@ -111,13 +134,15 @@ export class QueryBase<
|
||||
NativeQueryType extends NativeQuery | NativeVectorQuery | NativeTakeQuery,
|
||||
> implements AsyncIterable<RecordBatch>
|
||||
{
|
||||
protected inner!: NativeQueryType | Promise<NativeQueryType>;
|
||||
|
||||
/**
|
||||
* @hidden
|
||||
*/
|
||||
protected constructor(
|
||||
protected inner: NativeQueryType | Promise<NativeQueryType>,
|
||||
) {
|
||||
// intentionally empty
|
||||
protected constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
|
||||
if (inner !== undefined) {
|
||||
this.inner = inner;
|
||||
}
|
||||
}
|
||||
|
||||
// call a function on the inner (either a promise or the actual object)
|
||||
@@ -135,6 +160,15 @@ export class QueryBase<
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Return the native query used by the next terminal operation.
|
||||
*
|
||||
* @hidden
|
||||
*/
|
||||
protected async getInner(): Promise<NativeQueryType> {
|
||||
return this.inner;
|
||||
}
|
||||
|
||||
/**
|
||||
* Return only the specified columns.
|
||||
*
|
||||
@@ -207,16 +241,11 @@ export class QueryBase<
|
||||
/**
|
||||
* @hidden
|
||||
*/
|
||||
protected nativeExecute(
|
||||
protected async nativeExecute(
|
||||
options?: Partial<QueryExecutionOptions>,
|
||||
): Promise<NativeBatchIterator> {
|
||||
if (this.inner instanceof Promise) {
|
||||
return this.inner.then((inner) =>
|
||||
inner.execute(options?.maxBatchLength, options?.timeoutMs),
|
||||
);
|
||||
} else {
|
||||
return this.inner.execute(options?.maxBatchLength, options?.timeoutMs);
|
||||
}
|
||||
const inner = await this.getInner();
|
||||
return inner.execute(options?.maxBatchLength, options?.timeoutMs);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -245,12 +274,7 @@ export class QueryBase<
|
||||
/** Collect the results as an Arrow @see {@link ArrowTable}. */
|
||||
async toArrow(options?: Partial<QueryExecutionOptions>): Promise<ArrowTable> {
|
||||
const batches = [];
|
||||
let inner;
|
||||
if (this.inner instanceof Promise) {
|
||||
inner = await this.inner;
|
||||
} else {
|
||||
inner = this.inner;
|
||||
}
|
||||
const inner = await this.getInner();
|
||||
for await (const batch of new RecordBatchIterable(inner, options)) {
|
||||
batches.push(batch);
|
||||
}
|
||||
@@ -279,11 +303,8 @@ export class QueryBase<
|
||||
* @returns A Promise that resolves to a string containing the query execution plan explanation.
|
||||
*/
|
||||
async explainPlan(verbose = false): Promise<string> {
|
||||
if (this.inner instanceof Promise) {
|
||||
return this.inner.then((inner) => inner.explainPlan(verbose));
|
||||
} else {
|
||||
return this.inner.explainPlan(verbose);
|
||||
}
|
||||
const inner = await this.getInner();
|
||||
return inner.explainPlan(verbose);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -321,13 +342,8 @@ export class QueryBase<
|
||||
distributedMetrics?: AnalyzePlanDistributedMetrics,
|
||||
): Promise<string> {
|
||||
const distributedMetricsMode = distributedMetrics ?? "aggregate";
|
||||
if (this.inner instanceof Promise) {
|
||||
return this.inner.then((inner) =>
|
||||
inner.analyzePlan(distributedMetricsMode),
|
||||
);
|
||||
} else {
|
||||
return this.inner.analyzePlan(distributedMetricsMode);
|
||||
}
|
||||
const inner = await this.getInner();
|
||||
return inner.analyzePlan(distributedMetricsMode);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -339,12 +355,8 @@ export class QueryBase<
|
||||
* @returns An Arrow Schema describing the output columns.
|
||||
*/
|
||||
async outputSchema(): Promise<import("./arrow").Schema> {
|
||||
let schemaBuffer: Buffer;
|
||||
if (this.inner instanceof Promise) {
|
||||
schemaBuffer = await this.inner.then((inner) => inner.outputSchema());
|
||||
} else {
|
||||
schemaBuffer = await this.inner.outputSchema();
|
||||
}
|
||||
const inner = await this.getInner();
|
||||
const schemaBuffer = await inner.outputSchema();
|
||||
const schema = tableFromIPC(schemaBuffer).schema;
|
||||
return schema;
|
||||
}
|
||||
@@ -356,7 +368,7 @@ export class StandardQueryBase<
|
||||
extends QueryBase<NativeQueryType>
|
||||
implements ExecutableQuery
|
||||
{
|
||||
constructor(inner: NativeQueryType | Promise<NativeQueryType>) {
|
||||
constructor(inner?: NativeQueryType | Promise<NativeQueryType>) {
|
||||
super(inner);
|
||||
}
|
||||
|
||||
@@ -510,6 +522,13 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
super(inner);
|
||||
}
|
||||
|
||||
/**
|
||||
* @hidden
|
||||
*/
|
||||
protected doVectorCall(fn: (inner: NativeVectorQuery) => void) {
|
||||
super.doCall(fn);
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the number of partitions to search (probe)
|
||||
*
|
||||
@@ -537,7 +556,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* the minimum and maximum to the same value.
|
||||
*/
|
||||
nprobes(nprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.nprobes(nprobes));
|
||||
this.doVectorCall((inner) => inner.nprobes(nprobes));
|
||||
|
||||
return this;
|
||||
}
|
||||
@@ -551,7 +570,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* but will also increase latency.
|
||||
*/
|
||||
minimumNprobes(minimumNprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.minimumNprobes(minimumNprobes));
|
||||
this.doVectorCall((inner) => inner.minimumNprobes(minimumNprobes));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -565,7 +584,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* potential false negatives.
|
||||
*/
|
||||
maximumNprobes(maximumNprobes: number): VectorQuery {
|
||||
super.doCall((inner) => inner.maximumNprobes(maximumNprobes));
|
||||
this.doVectorCall((inner) => inner.maximumNprobes(maximumNprobes));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -578,7 +597,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* `undefined` means no lower or upper bound.
|
||||
*/
|
||||
distanceRange(lowerBound?: number, upperBound?: number): VectorQuery {
|
||||
super.doCall((inner) => inner.distanceRange(lowerBound, upperBound));
|
||||
this.doVectorCall((inner) => inner.distanceRange(lowerBound, upperBound));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -592,7 +611,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* also increase the latency of your query. The default value is 1.5*limit.
|
||||
*/
|
||||
ef(ef: number): VectorQuery {
|
||||
super.doCall((inner) => inner.ef(ef));
|
||||
this.doVectorCall((inner) => inner.ef(ef));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -606,7 +625,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* whose data type is a fixed-size-list of floats.
|
||||
*/
|
||||
column(column: string): VectorQuery {
|
||||
super.doCall((inner) => inner.column(column));
|
||||
this.doVectorCall((inner) => inner.column(column));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -627,7 +646,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
distanceType(
|
||||
distanceType: Required<IvfPqOptions>["distanceType"],
|
||||
): VectorQuery {
|
||||
super.doCall((inner) => inner.distanceType(distanceType));
|
||||
this.doVectorCall((inner) => inner.distanceType(distanceType));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -661,7 +680,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* distance between the query vector and the actual uncompressed vector.
|
||||
*/
|
||||
refineFactor(refineFactor: number): VectorQuery {
|
||||
super.doCall((inner) => inner.refineFactor(refineFactor));
|
||||
this.doVectorCall((inner) => inner.refineFactor(refineFactor));
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -686,7 +705,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* factor can often help restore some of the results lost by post filtering.
|
||||
*/
|
||||
postfilter(): VectorQuery {
|
||||
super.doCall((inner) => inner.postfilter());
|
||||
this.doVectorCall((inner) => inner.postfilter());
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -700,7 +719,7 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* calculate your recall to select an appropriate value for nprobes.
|
||||
*/
|
||||
bypassVectorIndex(): VectorQuery {
|
||||
super.doCall((inner) => inner.bypassVectorIndex());
|
||||
this.doVectorCall((inner) => inner.bypassVectorIndex());
|
||||
return this;
|
||||
}
|
||||
|
||||
@@ -708,43 +727,39 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
* Add a query vector to the search
|
||||
*
|
||||
* This method can be called multiple times to add multiple query vectors
|
||||
* to the search. If multiple query vectors are added, then they will be searched
|
||||
* in parallel, and the results will be concatenated. A column called `query_index`
|
||||
* will be added to indicate the index of the query vector that produced the result.
|
||||
*
|
||||
* Performance wise, this is equivalent to running multiple queries concurrently.
|
||||
* to the search. A column called `query_index` will be added to indicate the index
|
||||
* of the query vector that produced the result. Flat searches share one table scan
|
||||
* across the query vectors, avoiding the scan and memory amplification of running
|
||||
* multiple queries concurrently. Indexed searches may still perform per-vector
|
||||
* index work.
|
||||
*/
|
||||
addQueryVector(vector: IntoVector): VectorQuery {
|
||||
if (vector instanceof Promise) {
|
||||
// Observe the promise as soon as it is accepted. The existing native
|
||||
// query may still be pending, and delaying observation until it resolves
|
||||
// can otherwise surface a fast rejection as unhandled.
|
||||
const settledVector = vector.then(
|
||||
(value) => ({ status: "fulfilled" as const, value }),
|
||||
(reason) => ({ status: "rejected" as const, reason }),
|
||||
);
|
||||
const res = (async () => {
|
||||
try {
|
||||
const v = await vector;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
|
||||
const value: any = this.addQueryVector(v);
|
||||
const inner = value.inner as
|
||||
| NativeVectorQuery
|
||||
| Promise<NativeVectorQuery>;
|
||||
return inner;
|
||||
} catch (e) {
|
||||
return Promise.reject(e);
|
||||
const inner = await this.getInner();
|
||||
const outcome = await settledVector;
|
||||
if (outcome.status === "rejected") {
|
||||
throw outcome.reason;
|
||||
}
|
||||
addQueryVectorToNative(inner, outcome.value);
|
||||
return inner;
|
||||
})();
|
||||
return new VectorQuery(res);
|
||||
} else {
|
||||
super.doCall((inner) => {
|
||||
const raw = Array.isArray(vector) ? null : extractVectorBuffer(vector);
|
||||
if (raw) {
|
||||
inner.addQueryVectorRaw(raw.data, raw.dtype);
|
||||
} else {
|
||||
inner.addQueryVector(Float32Array.from(vector as number[]));
|
||||
}
|
||||
});
|
||||
this.doVectorCall((inner) => addQueryVectorToNative(inner, vector));
|
||||
return this;
|
||||
}
|
||||
}
|
||||
|
||||
rerank(reranker: Reranker): VectorQuery {
|
||||
super.doCall((inner) =>
|
||||
this.doVectorCall((inner) =>
|
||||
inner.rerank(async (args) => {
|
||||
const vecResults = await fromBufferToRecordBatch(args.vecResults);
|
||||
const ftsResults = await fromBufferToRecordBatch(args.ftsResults);
|
||||
@@ -763,6 +778,71 @@ export class VectorQuery extends StandardQueryBase<NativeVectorQuery> {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a string query whose vector/FTS routing is resolved against the active
|
||||
* table schema when the query executes.
|
||||
*
|
||||
* @hidden
|
||||
*/
|
||||
export function createAutoQuery(
|
||||
table: NativeTable,
|
||||
query: string,
|
||||
columns: string[] | null,
|
||||
getVector: (metadata: string) => Promise<Awaited<IntoVector>>,
|
||||
): AutoQuery {
|
||||
type RouteSnapshot = {
|
||||
table: NativeTable;
|
||||
embeddingMetadata: string | undefined;
|
||||
};
|
||||
type CachedPreparation = {
|
||||
metadata: string;
|
||||
vector: Promise<Awaited<IntoVector>>;
|
||||
};
|
||||
|
||||
let cachedPreparation: CachedPreparation | undefined;
|
||||
|
||||
const snapshotRoute = async (): Promise<RouteSnapshot> => {
|
||||
const snapshot = await table.querySnapshot();
|
||||
const schema = tableFromIPC(await snapshot.schema()).schema;
|
||||
return {
|
||||
table: snapshot,
|
||||
embeddingMetadata: schema.metadata.get("embedding_functions"),
|
||||
};
|
||||
};
|
||||
|
||||
const createInner = async (): Promise<NativeQuery | NativeVectorQuery> => {
|
||||
const route = await snapshotRoute();
|
||||
if (route.embeddingMetadata === undefined) {
|
||||
const inner = route.table.query();
|
||||
inner.fullTextSearch({ query, columns });
|
||||
return inner;
|
||||
}
|
||||
|
||||
const metadata = route.embeddingMetadata;
|
||||
if (cachedPreparation?.metadata !== metadata) {
|
||||
cachedPreparation = {
|
||||
metadata,
|
||||
vector: Promise.resolve().then(() => getVector(metadata)),
|
||||
};
|
||||
}
|
||||
|
||||
const preparation = cachedPreparation;
|
||||
let vector: Awaited<IntoVector>;
|
||||
try {
|
||||
vector = await preparation.vector;
|
||||
} catch (error) {
|
||||
if (cachedPreparation === preparation) {
|
||||
cachedPreparation = undefined;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
|
||||
return nearestToNative(route.table.query(), vector);
|
||||
};
|
||||
|
||||
return new AutoQuery(createInner);
|
||||
}
|
||||
|
||||
/**
|
||||
* A query that returns a subset of the rows in the table.
|
||||
*
|
||||
@@ -788,6 +868,51 @@ export class TakeQuery extends QueryBase<NativeTakeQuery> {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A builder for automatic string searches.
|
||||
*
|
||||
* Automatic search determines whether to use full-text or vector search from
|
||||
* the table revision selected for each execution. This builder exposes the
|
||||
* common operations supported by both query families.
|
||||
*
|
||||
* @hideconstructor
|
||||
*/
|
||||
export class AutoQuery extends StandardQueryBase<
|
||||
NativeQuery | NativeVectorQuery
|
||||
> {
|
||||
private readonly calls: Array<
|
||||
(inner: NativeQuery | NativeVectorQuery) => void
|
||||
> = [];
|
||||
|
||||
/** @hidden */
|
||||
constructor(
|
||||
private readonly createInner: () => Promise<
|
||||
NativeQuery | NativeVectorQuery
|
||||
>,
|
||||
) {
|
||||
super();
|
||||
}
|
||||
|
||||
/** @hidden */
|
||||
protected override doCall(
|
||||
fn: (inner: NativeQuery | NativeVectorQuery) => void,
|
||||
) {
|
||||
this.calls.push(fn);
|
||||
}
|
||||
|
||||
/** @hidden */
|
||||
protected override async getInner(): Promise<
|
||||
NativeQuery | NativeVectorQuery
|
||||
> {
|
||||
const calls = [...this.calls];
|
||||
const inner = await this.createInner();
|
||||
for (const call of calls) {
|
||||
call(inner);
|
||||
}
|
||||
return inner;
|
||||
}
|
||||
}
|
||||
|
||||
/** A builder for LanceDB queries.
|
||||
*
|
||||
* @see {@link Table#query}, {@link Table#search}
|
||||
@@ -840,45 +965,19 @@ export class Query extends StandardQueryBase<NativeQuery> {
|
||||
* a default `limit` of 10 will be used. @see {@link Query#limit}
|
||||
*/
|
||||
nearestTo(vector: IntoVector): VectorQuery {
|
||||
const callNearestTo = (
|
||||
inner: NativeQuery,
|
||||
resolved: Float32Array | Float64Array | Uint8Array | number[],
|
||||
): NativeVectorQuery => {
|
||||
const raw = Array.isArray(resolved)
|
||||
? null
|
||||
: extractVectorBuffer(resolved);
|
||||
if (raw) {
|
||||
return inner.nearestToRaw(raw.data, raw.dtype);
|
||||
}
|
||||
return inner.nearestTo(Float32Array.from(resolved as number[]));
|
||||
};
|
||||
|
||||
if (this.inner instanceof Promise) {
|
||||
const nativeQuery = this.inner.then(async (inner) => {
|
||||
const resolved = vector instanceof Promise ? await vector : vector;
|
||||
return callNearestTo(inner, resolved);
|
||||
});
|
||||
const inner = this.inner;
|
||||
if (inner instanceof Promise) {
|
||||
const nativeQuery = inner.then(async (resolvedInner) =>
|
||||
nearestToNative(resolvedInner, await vector),
|
||||
);
|
||||
return new VectorQuery(nativeQuery);
|
||||
}
|
||||
if (vector instanceof Promise) {
|
||||
const res = (async () => {
|
||||
try {
|
||||
const v = await vector;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: we need to get the `inner`, but js has no package scoping
|
||||
const value: any = this.nearestTo(v);
|
||||
const inner = value.inner as
|
||||
| NativeVectorQuery
|
||||
| Promise<NativeVectorQuery>;
|
||||
return inner;
|
||||
} catch (e) {
|
||||
return Promise.reject(e);
|
||||
}
|
||||
})();
|
||||
return new VectorQuery(res);
|
||||
} else {
|
||||
const vectorQuery = callNearestTo(this.inner, vector);
|
||||
return new VectorQuery(vectorQuery);
|
||||
return new VectorQuery(
|
||||
vector.then((resolvedVector) => nearestToNative(inner, resolvedVector)),
|
||||
);
|
||||
}
|
||||
return new VectorQuery(nearestToNative(inner, vector));
|
||||
}
|
||||
|
||||
nearestToText(query: string | FullTextQuery, columns?: string[]): Query {
|
||||
|
||||
@@ -94,17 +94,24 @@ export function sanitizeMetadata(
|
||||
if (metadataLike === undefined || metadataLike === null) {
|
||||
return undefined;
|
||||
}
|
||||
if (!(metadataLike instanceof Map)) {
|
||||
|
||||
let entries: IterableIterator<[unknown, unknown]>;
|
||||
try {
|
||||
entries = Map.prototype.entries.call(metadataLike);
|
||||
} catch {
|
||||
throw Error("Expected metadata, if present, to be a Map<string, string>");
|
||||
}
|
||||
for (const item of metadataLike) {
|
||||
if (typeof item[0] !== "string" || typeof item[1] !== "string") {
|
||||
|
||||
const metadata = new Map<string, string>();
|
||||
for (const [key, value] of entries) {
|
||||
if (typeof key !== "string" || typeof value !== "string") {
|
||||
throw Error(
|
||||
"Expected metadata, if present, to be a Map<string, string> but it had non-string keys or values",
|
||||
);
|
||||
}
|
||||
metadata.set(key, value);
|
||||
}
|
||||
return metadataLike as Map<string, string>;
|
||||
return metadata;
|
||||
}
|
||||
|
||||
export function sanitizeInt(typeLike: object) {
|
||||
|
||||
@@ -0,0 +1,567 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import {
|
||||
Binary,
|
||||
Bool,
|
||||
DataType,
|
||||
Dictionary,
|
||||
Field,
|
||||
FixedSizeList,
|
||||
Float32,
|
||||
Float64,
|
||||
Int32,
|
||||
Int64,
|
||||
List,
|
||||
Schema,
|
||||
Struct,
|
||||
Utf8,
|
||||
util as arrowUtil,
|
||||
} from "apache-arrow";
|
||||
import { typedArrayToArrowType } from "./arrow_type";
|
||||
import { sanitizeType } from "./sanitize";
|
||||
|
||||
type InferenceOptions = {
|
||||
dictionaryEncodeStrings: boolean;
|
||||
vectorColumns: Record<string, { type: unknown }>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Infer the Arrow schema represented by a set of records.
|
||||
*
|
||||
* This is the intentionally small interface to schema inference. The stateful
|
||||
* details of combining partial type evidence are encapsulated below so callers
|
||||
* only need to provide records, an optional schema, and inference options.
|
||||
*/
|
||||
export function inferSchema(
|
||||
data: Array<Record<string, unknown>>,
|
||||
schema: Schema | undefined,
|
||||
options: InferenceOptions,
|
||||
): Schema {
|
||||
return new SchemaInferrer(schema, options).infer(data);
|
||||
}
|
||||
|
||||
class SchemaInferrer {
|
||||
private readonly fields = new FieldTree();
|
||||
|
||||
constructor(
|
||||
private readonly providedSchema: Schema | undefined,
|
||||
private readonly options: InferenceOptions,
|
||||
) {}
|
||||
|
||||
infer(data: Array<Record<string, unknown>>): Schema {
|
||||
for (const [row, record] of data.entries()) {
|
||||
for (const [path, value] of recordPathsAndValues(record)) {
|
||||
this.observe(path, value, row);
|
||||
}
|
||||
}
|
||||
|
||||
return this.providedSchema === undefined
|
||||
? new Schema(fieldsFromTree(this.fields))
|
||||
: new Schema(matchingFields(this.providedSchema.fields, this.fields));
|
||||
}
|
||||
|
||||
private observe(path: string[], value: unknown, row: number): void {
|
||||
const current = this.fields.get(path);
|
||||
if (current === undefined) {
|
||||
this.addField(path, value, row);
|
||||
} else if (this.providedSchema === undefined) {
|
||||
this.updateInferredField(path, value, row, current);
|
||||
}
|
||||
}
|
||||
|
||||
private addField(path: string[], value: unknown, row: number): void {
|
||||
if (this.providedSchema !== undefined) {
|
||||
this.addSchemaField(this.providedSchema, path, row);
|
||||
return;
|
||||
}
|
||||
|
||||
const evidence =
|
||||
this.inferType(value, path) ?? DeferredTypeEvidence.from(value, row);
|
||||
if (evidence === undefined) {
|
||||
throw typeInferenceError(path, row);
|
||||
}
|
||||
|
||||
const conflict = this.fields.set(
|
||||
path,
|
||||
evidence,
|
||||
(existing) =>
|
||||
existing instanceof DeferredTypeEvidence && existing.isOnlyNulls(),
|
||||
);
|
||||
if (conflict !== undefined) {
|
||||
throw branchConflictError(conflict, row, "Struct");
|
||||
}
|
||||
}
|
||||
|
||||
private addSchemaField(schema: Schema, path: string[], row: number): void {
|
||||
const field = fieldAtPath(schema, path);
|
||||
if (field === undefined) {
|
||||
throw new Error(
|
||||
`Found field not in schema: ${path.join(".")} at row ${row}`,
|
||||
);
|
||||
}
|
||||
|
||||
const conflict = this.fields.set(path, field.type);
|
||||
if (conflict !== undefined) {
|
||||
throw branchConflictError(conflict, row, "Struct");
|
||||
}
|
||||
}
|
||||
|
||||
private updateInferredField(
|
||||
path: string[],
|
||||
value: unknown,
|
||||
row: number,
|
||||
current: FieldNode,
|
||||
): void {
|
||||
const newType = this.inferType(value, path);
|
||||
const deferred = DeferredTypeEvidence.from(value, row);
|
||||
|
||||
if (current instanceof FieldTree) {
|
||||
if (deferred?.isOnlyNulls()) {
|
||||
return;
|
||||
}
|
||||
throw schemaInferenceError(
|
||||
path,
|
||||
row,
|
||||
"Struct",
|
||||
describeEvidence(newType ?? deferred),
|
||||
);
|
||||
}
|
||||
|
||||
if (current instanceof DeferredTypeEvidence) {
|
||||
this.resolveDeferredField(path, row, current, newType, deferred);
|
||||
return;
|
||||
}
|
||||
|
||||
if (newType !== undefined) {
|
||||
if (!inferredTypesEqual(current, newType)) {
|
||||
throw schemaInferenceError(
|
||||
path,
|
||||
row,
|
||||
describeEvidence(current),
|
||||
describeEvidence(newType),
|
||||
);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (deferred === undefined || !deferred.matches(current)) {
|
||||
throw schemaInferenceError(
|
||||
path,
|
||||
row,
|
||||
describeEvidence(current),
|
||||
describeEvidence(deferred),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
private resolveDeferredField(
|
||||
path: string[],
|
||||
row: number,
|
||||
current: DeferredTypeEvidence,
|
||||
newType: DataType | undefined,
|
||||
deferred: DeferredTypeEvidence | undefined,
|
||||
): void {
|
||||
if (newType !== undefined) {
|
||||
if (!current.matches(newType)) {
|
||||
throw schemaInferenceError(
|
||||
path,
|
||||
row,
|
||||
current.describe(),
|
||||
describeEvidence(newType),
|
||||
);
|
||||
}
|
||||
this.fields.set(path, newType);
|
||||
return;
|
||||
}
|
||||
|
||||
if (deferred !== undefined) {
|
||||
this.fields.set(path, current.merge(deferred));
|
||||
return;
|
||||
}
|
||||
|
||||
throw schemaInferenceError(
|
||||
path,
|
||||
row,
|
||||
current.describe(),
|
||||
describeEvidence(newType),
|
||||
);
|
||||
}
|
||||
|
||||
private inferType(value: unknown, path: string[]): DataType | undefined {
|
||||
if (typeof value === "bigint") {
|
||||
return new Int64();
|
||||
}
|
||||
if (typeof value === "number") {
|
||||
return new Float64();
|
||||
}
|
||||
if (typeof value === "string") {
|
||||
return this.options.dictionaryEncodeStrings
|
||||
? new Dictionary(new Utf8(), new Int32())
|
||||
: new Utf8();
|
||||
}
|
||||
if (typeof value === "boolean") {
|
||||
return new Bool();
|
||||
}
|
||||
if (value instanceof Buffer) {
|
||||
return new Binary();
|
||||
}
|
||||
if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
|
||||
const typedArray = typedArrayToArrowType(value);
|
||||
return typedArray === undefined
|
||||
? undefined
|
||||
: new FixedSizeList(
|
||||
typedArray.length,
|
||||
new Field("item", typedArray.elementType, true),
|
||||
);
|
||||
}
|
||||
if (!Array.isArray(value) || value.length === 0) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const configuredVector =
|
||||
path.length === 1 ? this.options.vectorColumns[path[0]] : undefined;
|
||||
if (configuredVector !== undefined) {
|
||||
return new FixedSizeList(
|
||||
value.length,
|
||||
new Field("item", sanitizeType(configuredVector.type), true),
|
||||
);
|
||||
}
|
||||
|
||||
const itemType = this.inferArrayItemType(value, path);
|
||||
if (itemType === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
return nameSuggestsVectorColumn(path[path.length - 1])
|
||||
? new FixedSizeList(value.length, new Field("item", new Float32(), true))
|
||||
: new List(new Field("item", itemType, true));
|
||||
}
|
||||
|
||||
private inferArrayItemType(
|
||||
values: unknown[],
|
||||
path: string[],
|
||||
): DataType | undefined {
|
||||
let itemType: DataType | undefined;
|
||||
const deferredItems: unknown[] = [];
|
||||
|
||||
for (const value of values) {
|
||||
const candidate = this.inferType(value, path);
|
||||
if (candidate === undefined) {
|
||||
if (!isDeferredValue(value)) {
|
||||
return undefined;
|
||||
}
|
||||
deferredItems.push(value);
|
||||
} else if (itemType === undefined) {
|
||||
itemType = candidate;
|
||||
} else if (!inferredTypesEqual(itemType, candidate)) {
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
|
||||
if (itemType === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
return deferredItems.every((value) =>
|
||||
deferredValueMatchesType(value, itemType),
|
||||
)
|
||||
? itemType
|
||||
: undefined;
|
||||
}
|
||||
}
|
||||
|
||||
/** Nulls and empty/all-null lists that do not determine a type by themselves. */
|
||||
class DeferredTypeEvidence {
|
||||
private constructor(
|
||||
private readonly values: Array<{ value: unknown; row: number }>,
|
||||
) {}
|
||||
|
||||
static from(value: unknown, row: number): DeferredTypeEvidence | undefined {
|
||||
return isDeferredValue(value)
|
||||
? new DeferredTypeEvidence([{ value, row }])
|
||||
: undefined;
|
||||
}
|
||||
|
||||
isOnlyNulls(): boolean {
|
||||
return this.values.every(({ value }) => value == null);
|
||||
}
|
||||
|
||||
matches(type: DataType): boolean {
|
||||
return this.values.every(({ value }) =>
|
||||
deferredValueMatchesType(value, type),
|
||||
);
|
||||
}
|
||||
|
||||
merge(other: DeferredTypeEvidence): DeferredTypeEvidence {
|
||||
return new DeferredTypeEvidence([...this.values, ...other.values]);
|
||||
}
|
||||
|
||||
describe(): string {
|
||||
const list = this.values.find(({ value }) => Array.isArray(value));
|
||||
return list === undefined
|
||||
? "null"
|
||||
: `List[${(list.value as unknown[]).length}]`;
|
||||
}
|
||||
|
||||
firstRow(): number {
|
||||
return this.values[0].row;
|
||||
}
|
||||
}
|
||||
|
||||
type FieldNode = DataType | DeferredTypeEvidence | FieldTree;
|
||||
type LeafNode = Exclude<FieldNode, FieldTree>;
|
||||
type FieldConflict = { path: string[]; value: FieldNode };
|
||||
|
||||
/** Nested field state, kept separate from Arrow's eventual Struct types. */
|
||||
class FieldTree {
|
||||
private readonly children = new Map<string, FieldNode>();
|
||||
|
||||
get(path: string[]): FieldNode | undefined {
|
||||
let current: FieldNode = this;
|
||||
for (const part of path) {
|
||||
if (!(current instanceof FieldTree)) {
|
||||
return undefined;
|
||||
}
|
||||
const child = current.children.get(part);
|
||||
if (child === undefined) {
|
||||
return undefined;
|
||||
}
|
||||
current = child;
|
||||
}
|
||||
return current;
|
||||
}
|
||||
|
||||
set(
|
||||
path: string[],
|
||||
value: LeafNode,
|
||||
canReplaceLeaf: (value: LeafNode) => boolean = () => false,
|
||||
): FieldConflict | undefined {
|
||||
let branch: FieldTree = this;
|
||||
for (const [index, part] of path.slice(0, -1).entries()) {
|
||||
const child = branch.children.get(part);
|
||||
if (child === undefined || (isLeaf(child) && canReplaceLeaf(child))) {
|
||||
const nextBranch = new FieldTree();
|
||||
branch.children.set(part, nextBranch);
|
||||
branch = nextBranch;
|
||||
} else if (child instanceof FieldTree) {
|
||||
branch = child;
|
||||
} else {
|
||||
return { path: path.slice(0, index + 1), value: child };
|
||||
}
|
||||
}
|
||||
|
||||
const name = path[path.length - 1];
|
||||
const current = branch.children.get(name);
|
||||
if (current instanceof FieldTree) {
|
||||
return { path, value: current };
|
||||
}
|
||||
branch.children.set(name, value);
|
||||
return undefined;
|
||||
}
|
||||
|
||||
entries(): IterableIterator<[string, FieldNode]> {
|
||||
return this.children.entries();
|
||||
}
|
||||
|
||||
has(name: string): boolean {
|
||||
return this.children.has(name);
|
||||
}
|
||||
}
|
||||
|
||||
function isLeaf(value: FieldNode): value is LeafNode {
|
||||
return !(value instanceof FieldTree);
|
||||
}
|
||||
|
||||
function fieldsFromTree(tree: FieldTree, path: string[] = []): Field[] {
|
||||
const fields: Field[] = [];
|
||||
for (const [name, value] of tree.entries()) {
|
||||
if (value instanceof FieldTree) {
|
||||
fields.push(
|
||||
new Field(
|
||||
name,
|
||||
new Struct(fieldsFromTree(value, [...path, name])),
|
||||
true,
|
||||
),
|
||||
);
|
||||
} else if (value instanceof DeferredTypeEvidence) {
|
||||
throw typeInferenceError([...path, name], value.firstRow());
|
||||
} else {
|
||||
fields.push(new Field(name, value, true));
|
||||
}
|
||||
}
|
||||
return fields;
|
||||
}
|
||||
|
||||
function matchingFields(fields: Field[], tree: FieldTree): Field[] {
|
||||
const matches: Field[] = [];
|
||||
for (const field of fields) {
|
||||
if (!tree.has(field.name)) {
|
||||
continue;
|
||||
}
|
||||
const value = tree.get([field.name]);
|
||||
if (value instanceof FieldTree) {
|
||||
const struct = field.type as Struct;
|
||||
matches.push(
|
||||
new Field(
|
||||
field.name,
|
||||
new Struct(matchingFields(struct.children, value)),
|
||||
field.nullable,
|
||||
field.metadata,
|
||||
),
|
||||
);
|
||||
} else {
|
||||
matches.push(field);
|
||||
}
|
||||
}
|
||||
return matches;
|
||||
}
|
||||
|
||||
function* recordPathsAndValues(
|
||||
record: Record<string, unknown>,
|
||||
path: string[] = [],
|
||||
): Generator<[string[], unknown]> {
|
||||
for (const [name, value] of Object.entries(record)) {
|
||||
if (isRecord(value)) {
|
||||
yield* recordPathsAndValues(value, [...path, name]);
|
||||
} else if (value !== undefined) {
|
||||
yield [[...path, name], value];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return (
|
||||
typeof value === "object" &&
|
||||
value !== null &&
|
||||
!Array.isArray(value) &&
|
||||
!(value instanceof RegExp) &&
|
||||
!(value instanceof Date) &&
|
||||
!(value instanceof Set) &&
|
||||
!(value instanceof Map) &&
|
||||
!(value instanceof Buffer) &&
|
||||
!ArrayBuffer.isView(value)
|
||||
);
|
||||
}
|
||||
|
||||
function fieldAtPath(schema: Schema, path: string[]): Field | undefined {
|
||||
let fields = schema.fields;
|
||||
let field: Field | undefined;
|
||||
for (const [index, name] of path.entries()) {
|
||||
field = fields.find((candidate) => candidate.name === name);
|
||||
if (field === undefined || index === path.length - 1) {
|
||||
return field;
|
||||
}
|
||||
if (!DataType.isStruct(field.type)) {
|
||||
return undefined;
|
||||
}
|
||||
fields = field.type.children;
|
||||
}
|
||||
return field;
|
||||
}
|
||||
|
||||
function isDeferredValue(value: unknown): boolean {
|
||||
return (
|
||||
value == null || (Array.isArray(value) && value.every(isDeferredValue))
|
||||
);
|
||||
}
|
||||
|
||||
function deferredValueMatchesType(value: unknown, type: DataType): boolean {
|
||||
if (value == null) {
|
||||
return true;
|
||||
}
|
||||
if (!Array.isArray(value)) {
|
||||
return false;
|
||||
}
|
||||
if (DataType.isList(type)) {
|
||||
return value.every((item) =>
|
||||
deferredValueMatchesType(item, type.valueType),
|
||||
);
|
||||
}
|
||||
if (DataType.isFixedSizeList(type)) {
|
||||
return (
|
||||
value.length === type.listSize &&
|
||||
value.every((item) => deferredValueMatchesType(item, type.valueType))
|
||||
);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function inferredTypesEqual(current: DataType, candidate: DataType): boolean {
|
||||
if (DataType.isDictionary(current)) {
|
||||
return (
|
||||
DataType.isDictionary(candidate) &&
|
||||
current.isOrdered === candidate.isOrdered &&
|
||||
inferredTypesEqual(current.indices, candidate.indices) &&
|
||||
inferredTypesEqual(current.dictionary, candidate.dictionary)
|
||||
);
|
||||
}
|
||||
if (DataType.isList(current)) {
|
||||
return (
|
||||
DataType.isList(candidate) &&
|
||||
current.valueField.name === candidate.valueField.name &&
|
||||
current.valueField.nullable === candidate.valueField.nullable &&
|
||||
inferredTypesEqual(current.valueType, candidate.valueType)
|
||||
);
|
||||
}
|
||||
if (DataType.isFixedSizeList(current)) {
|
||||
return (
|
||||
DataType.isFixedSizeList(candidate) &&
|
||||
current.listSize === candidate.listSize &&
|
||||
current.valueField.name === candidate.valueField.name &&
|
||||
current.valueField.nullable === candidate.valueField.nullable &&
|
||||
inferredTypesEqual(current.valueType, candidate.valueType)
|
||||
);
|
||||
}
|
||||
return arrowUtil.compareTypes(current, candidate);
|
||||
}
|
||||
|
||||
function describeEvidence(
|
||||
evidence: DataType | DeferredTypeEvidence | undefined,
|
||||
): string {
|
||||
if (evidence === undefined) {
|
||||
return "an unsupported value";
|
||||
}
|
||||
return evidence instanceof DeferredTypeEvidence
|
||||
? evidence.describe()
|
||||
: evidence.toString();
|
||||
}
|
||||
|
||||
function branchConflictError(
|
||||
conflict: FieldConflict,
|
||||
row: number,
|
||||
candidate: string,
|
||||
): Error {
|
||||
return schemaInferenceError(
|
||||
conflict.path,
|
||||
row,
|
||||
conflict.value instanceof FieldTree
|
||||
? "Struct"
|
||||
: describeEvidence(conflict.value),
|
||||
candidate,
|
||||
);
|
||||
}
|
||||
|
||||
function schemaInferenceError(
|
||||
path: string[],
|
||||
row: number,
|
||||
currentType: string,
|
||||
newType: string,
|
||||
): Error {
|
||||
return new Error(
|
||||
`Failed to infer schema for data. Previously inferred type ${currentType} ` +
|
||||
`but found ${newType} for field ${path.join(".")} at row ${row}. ` +
|
||||
"Consider providing an explicit schema.",
|
||||
);
|
||||
}
|
||||
|
||||
function typeInferenceError(path: string[], row: number): Error {
|
||||
return new Error(
|
||||
`Failed to infer data type for field ${path.join(".")} at row ${row}. ` +
|
||||
"Consider providing an explicit schema.",
|
||||
);
|
||||
}
|
||||
|
||||
function nameSuggestsVectorColumn(name: string): boolean {
|
||||
const normalized = name.toLowerCase();
|
||||
return normalized.includes("vector") || normalized.includes("embedding");
|
||||
}
|
||||
+45
-15
@@ -43,10 +43,12 @@ import {
|
||||
Table as _NativeTable,
|
||||
} from "./native";
|
||||
import {
|
||||
AutoQuery,
|
||||
FullTextQuery,
|
||||
Query,
|
||||
TakeQuery,
|
||||
VectorQuery,
|
||||
createAutoQuery,
|
||||
instanceOfFullTextQuery,
|
||||
} from "./query";
|
||||
import { sanitizeType } from "./sanitize";
|
||||
@@ -523,7 +525,7 @@ export abstract class Table {
|
||||
query: string | IntoVector | MultiVector | FullTextQuery,
|
||||
queryType?: string,
|
||||
ftsColumns?: string | string[],
|
||||
): VectorQuery | Query;
|
||||
): VectorQuery | Query | AutoQuery;
|
||||
/**
|
||||
* Search the table with a given query vector.
|
||||
*
|
||||
@@ -628,6 +630,18 @@ export abstract class Table {
|
||||
|
||||
/**
|
||||
* Update per-field (column) metadata.
|
||||
*
|
||||
* The following keys are treated specially, by convention, and should be
|
||||
* used when appropriate:
|
||||
*
|
||||
* - `lancedb:description`: for a human-readable description of a field.
|
||||
* - `lancedb:tag:<name>`: for a user-defined key-value tag, where the suffix
|
||||
* names the tag category; e.g. `lancedb:tag:model: "clip"`.
|
||||
* - `lancedb:logical-column`: for a column grouping; e.g. `feature_v1` and
|
||||
* `feature_v2` might be in the same logical column.
|
||||
* - `lancedb:status`: for status options (`production`, `candidate`,
|
||||
* `deprecated`, `archived`) to designate the current life cycle state of
|
||||
* this column.
|
||||
* @param {FieldMetadataUpdate[]} updates One or more per-field updates. Each
|
||||
* update's metadata is merged into the field's existing metadata by default;
|
||||
* a value of `null` deletes that key, and `replace: true` swaps the whole map.
|
||||
@@ -975,10 +989,11 @@ export class LocalTable extends Table {
|
||||
return this.inner.display();
|
||||
}
|
||||
|
||||
private async getEmbeddingFunctions(): Promise<
|
||||
Map<string, EmbeddingFunctionConfig>
|
||||
> {
|
||||
const schema = await this.schema();
|
||||
private async getEmbeddingFunctions(
|
||||
inner: _NativeTable = this.inner,
|
||||
): Promise<Map<string, EmbeddingFunctionConfig>> {
|
||||
const schemaBuf = await inner.schema();
|
||||
const schema = tableFromIPC(schemaBuf).schema;
|
||||
const registry = getRegistry();
|
||||
return registry.parseFunctions(schema.metadata);
|
||||
}
|
||||
@@ -1160,7 +1175,7 @@ export class LocalTable extends Table {
|
||||
query: string | IntoVector | MultiVector | FullTextQuery,
|
||||
queryType: string = "auto",
|
||||
ftsColumns?: string | string[],
|
||||
): VectorQuery | Query {
|
||||
): VectorQuery | Query | AutoQuery {
|
||||
if (typeof query !== "string" && !instanceOfFullTextQuery(query)) {
|
||||
if (queryType === "fts") {
|
||||
throw new Error("Cannot perform full text search on a vector query");
|
||||
@@ -1175,14 +1190,28 @@ export class LocalTable extends Table {
|
||||
});
|
||||
}
|
||||
|
||||
// The query type is auto or vector
|
||||
// fall back to full text search if no embedding functions are defined and the query is a string
|
||||
if (
|
||||
queryType === "auto" &&
|
||||
(getRegistry().length() === 0 || instanceOfFullTextQuery(query))
|
||||
) {
|
||||
return this.query().fullTextSearch(query, {
|
||||
columns: ftsColumns,
|
||||
if (queryType === "auto") {
|
||||
if (instanceOfFullTextQuery(query)) {
|
||||
return this.query().fullTextSearch(query, {
|
||||
columns: ftsColumns,
|
||||
});
|
||||
}
|
||||
|
||||
const columns =
|
||||
typeof ftsColumns === "string" ? [ftsColumns] : (ftsColumns ?? null);
|
||||
return createAutoQuery(this.inner, query, columns, async (metadata) => {
|
||||
const functions = await getRegistry().parseFunctions(
|
||||
new Map([["embedding_functions", metadata]]),
|
||||
);
|
||||
// TODO: Support multiple embedding functions
|
||||
const embeddingFunc: EmbeddingFunctionConfig | undefined = functions
|
||||
.values()
|
||||
.next().value;
|
||||
// The route only calls this callback when embedding metadata exists.
|
||||
// parseFunctions either yields a provider or reports malformed metadata.
|
||||
if (!embeddingFunc)
|
||||
throw new Error("Invalid embedding function metadata");
|
||||
return await embeddingFunc.function.computeQueryEmbeddings(query);
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1538,7 +1567,8 @@ export interface FieldMetadataUpdate {
|
||||
path: string;
|
||||
/**
|
||||
* Metadata key/value pairs. Merged into the field's existing metadata by
|
||||
* default; a value of `null` deletes that key.
|
||||
* default; a value of `null` deletes that key. See
|
||||
* {@link Table.updateFieldMetadata} for the conventional `lancedb:*` keys.
|
||||
*/
|
||||
metadata: Record<string, string | null>;
|
||||
/** If true, replace the field's entire metadata map instead of merging. */
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-darwin-arm64",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["darwin"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.darwin-arm64.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-gnu",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-arm64-musl",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["arm64"],
|
||||
"main": "lancedb.linux-arm64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-gnu",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-gnu.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-linux-x64-musl",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["linux"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.linux-x64-musl.node",
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-x64-msvc",
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"os": ["win32"],
|
||||
"cpu": ["x64"],
|
||||
"main": "lancedb.win32-x64-msvc.node",
|
||||
|
||||
Generated
-11106
File diff suppressed because it is too large
Load Diff
+3
-3
@@ -11,7 +11,7 @@
|
||||
"ann"
|
||||
],
|
||||
"private": false,
|
||||
"version": "0.38.0-beta.6",
|
||||
"version": "0.38.0-beta.12",
|
||||
"main": "dist/index.js",
|
||||
"exports": {
|
||||
".": "./dist/index.js",
|
||||
@@ -67,7 +67,7 @@
|
||||
"timeout": "3m"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">= 18"
|
||||
"node": ">= 22"
|
||||
},
|
||||
"packageManager": "pnpm@11.1.1",
|
||||
"cpu": ["x64", "arm64"],
|
||||
@@ -101,7 +101,7 @@
|
||||
"openai": "4.29.2"
|
||||
},
|
||||
"peerDependencies": {
|
||||
"@types/node": ">=18",
|
||||
"@types/node": ">=22",
|
||||
"apache-arrow": ">=15.0.0 <=18.1.0"
|
||||
},
|
||||
"peerDependenciesMeta": {
|
||||
|
||||
@@ -17,6 +17,7 @@ use lancedb::connection::{ConnectBuilder, Connection as LanceDBConnection, conne
|
||||
|
||||
use lance_namespace::models::{
|
||||
CreateNamespaceRequest, DescribeNamespaceRequest, DropNamespaceRequest, ListNamespacesRequest,
|
||||
ListTablesRequest,
|
||||
};
|
||||
use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema};
|
||||
|
||||
@@ -36,6 +37,12 @@ pub struct ListNamespacesResponse {
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct ListTablesResponse {
|
||||
pub tables: Vec<String>,
|
||||
pub page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[napi(object)]
|
||||
pub struct CreateNamespaceResponse {
|
||||
pub properties: Option<HashMap<String, String>>,
|
||||
@@ -206,6 +213,33 @@ impl Connection {
|
||||
op.execute().await.default_error()
|
||||
}
|
||||
|
||||
/// List a page of tables in the database.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn list_tables(
|
||||
&self,
|
||||
namespace_path: Option<Vec<String>>,
|
||||
page_token: Option<String>,
|
||||
limit: Option<u32>,
|
||||
) -> napi::Result<ListTablesResponse> {
|
||||
let request = ListTablesRequest {
|
||||
// The root namespace is an empty path, not an absent one: a namespace-backed
|
||||
// database rejects a request that names no namespace.
|
||||
id: Some(namespace_path.unwrap_or_default()),
|
||||
page_token,
|
||||
limit: limit.map(|limit| i32::try_from(limit).unwrap_or(i32::MAX)),
|
||||
..Default::default()
|
||||
};
|
||||
let response = self
|
||||
.get_inner()?
|
||||
.list_tables(request)
|
||||
.await
|
||||
.default_error()?;
|
||||
Ok(ListTablesResponse {
|
||||
tables: response.tables,
|
||||
page_token: response.page_token,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create table from a Apache Arrow IPC (file) buffer.
|
||||
///
|
||||
/// Parameters:
|
||||
|
||||
@@ -278,6 +278,13 @@ impl Table {
|
||||
Ok(Query::new(self.inner_ref()?.query()))
|
||||
}
|
||||
|
||||
/// Return a read-only table handle pinned to the current query revision.
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn query_snapshot(&self) -> napi::Result<Self> {
|
||||
let snapshot = self.inner_ref()?.query_snapshot().await.default_error()?;
|
||||
Ok(Self::new(snapshot))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub fn take_offsets(&self, offsets: Vec<i64>) -> napi::Result<TakeQuery> {
|
||||
Ok(TakeQuery::new(
|
||||
@@ -554,6 +561,12 @@ impl Table {
|
||||
.default_error()
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn checkout_current(&self) -> napi::Result<Self> {
|
||||
let table = self.inner_ref()?.checkout_current().await.default_error()?;
|
||||
Ok(Self::new(table))
|
||||
}
|
||||
|
||||
#[napi(catch_unwind)]
|
||||
pub async fn checkout(&self, version: i64) -> napi::Result<()> {
|
||||
self.inner_ref()?
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb-python"
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
publish = false
|
||||
edition.workspace = true
|
||||
description = "Python bindings for LanceDB"
|
||||
|
||||
@@ -101,9 +101,12 @@ azure = ["adlfs>=2024.2.0"]
|
||||
[tool.maturin]
|
||||
python-source = "python"
|
||||
module-name = "lancedb._lancedb"
|
||||
# uv installs the project as an editable package before `uv run`, so keep that
|
||||
# bootstrap build consistent with `maturin develop`.
|
||||
editable-profile = "dev"
|
||||
|
||||
[build-system]
|
||||
requires = ["maturin>=1.9.4"]
|
||||
requires = ["maturin>=1.10"]
|
||||
build-backend = "maturin"
|
||||
|
||||
[tool.ruff.lint]
|
||||
|
||||
@@ -6,7 +6,7 @@ import importlib.metadata
|
||||
import os
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
from typing import Dict, Optional, Union, Any, List, Iterable
|
||||
from typing import Dict, Optional, Union, Any, List, Iterable, TYPE_CHECKING
|
||||
|
||||
__version__ = importlib.metadata.version("lancedb")
|
||||
|
||||
@@ -20,7 +20,7 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection
|
||||
from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .expr import Expr, col, lit, func
|
||||
from .schema import blob, vector, BlobType
|
||||
from .schema import blob, vector
|
||||
from .job import AsyncJob, Job
|
||||
from .functions import (
|
||||
FunctionArtifactRequest as FunctionArtifactRequest,
|
||||
@@ -49,6 +49,19 @@ from .namespace import (
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lance.blob import BlobType as BlobType
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "BlobType":
|
||||
from .schema import BlobType
|
||||
|
||||
globals()["BlobType"] = BlobType
|
||||
return BlobType
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
def _check_s3_bucket_with_dots(
|
||||
uri: str, storage_options: Optional[Dict[str, str]]
|
||||
) -> None:
|
||||
@@ -179,6 +192,18 @@ def connect(
|
||||
... },
|
||||
... )
|
||||
|
||||
For Azure Blob Storage, credentials can be passed directly without setting
|
||||
environment variables:
|
||||
|
||||
>>> azure_storage_options = {
|
||||
... "account_name": "some-account",
|
||||
... "account_key": "some-key",
|
||||
... }
|
||||
>>> db = lancedb.connect( # doctest: +SKIP
|
||||
... "az://my-container/my-database",
|
||||
... storage_options=azure_storage_options,
|
||||
... )
|
||||
|
||||
For tests and temporary data, use an in-memory database:
|
||||
|
||||
>>> db = lancedb.connect("memory://")
|
||||
@@ -465,6 +490,10 @@ async def connect_async(
|
||||
--------
|
||||
|
||||
>>> import lancedb
|
||||
>>> azure_storage_options = {
|
||||
... "account_name": "some-account",
|
||||
... "account_key": "some-key",
|
||||
... }
|
||||
>>> async def doctest_example():
|
||||
... # For a local directory, provide a path to the database
|
||||
... db = await lancedb.connect_async("~/.lancedb")
|
||||
@@ -472,6 +501,11 @@ async def connect_async(
|
||||
... db = await lancedb.connect_async("s3://my-bucket/lancedb",
|
||||
... storage_options={
|
||||
... "aws_access_key_id": "***"})
|
||||
... # Azure credentials can also be passed directly
|
||||
... db = await lancedb.connect_async(
|
||||
... "az://my-container/my-database",
|
||||
... storage_options=azure_storage_options,
|
||||
... )
|
||||
... # For tests and temporary data, use an in-memory database
|
||||
... db = await lancedb.connect_async("memory://")
|
||||
... # Connect to LanceDB cloud
|
||||
|
||||
@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union
|
||||
import pyarrow as pa
|
||||
|
||||
from .expr import Expr
|
||||
from .schema import blob_v2_column_paths
|
||||
from .schema import row_addressable_blob_v2_paths
|
||||
from .types import BlobMode, QueryProjection, QueryProjectionSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -119,7 +119,7 @@ def blob_v2_projection_sources(
|
||||
schema: pa.Schema,
|
||||
projection: QueryProjection,
|
||||
) -> dict[str, str]:
|
||||
blob_columns = blob_v2_column_paths(schema)
|
||||
blob_columns = row_addressable_blob_v2_paths(schema)
|
||||
if not blob_columns:
|
||||
return {}
|
||||
columns = set(blob_columns)
|
||||
@@ -140,7 +140,9 @@ def v2_projection_needs_row_id(
|
||||
) -> bool:
|
||||
if with_row_id:
|
||||
return False
|
||||
return projection_includes_blob_column(projection, blob_v2_column_paths(schema))
|
||||
return projection_includes_blob_column(
|
||||
projection, row_addressable_blob_v2_paths(schema)
|
||||
)
|
||||
|
||||
|
||||
def blob_auto_row_id_for_scan(
|
||||
@@ -270,7 +272,8 @@ def _iter_projection_pairs(
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
source = expr._column_name()
|
||||
yield name, source if source is not None else expr.to_sql()
|
||||
return
|
||||
for column in projection:
|
||||
if isinstance(column, str):
|
||||
@@ -280,7 +283,8 @@ def _iter_projection_pairs(
|
||||
if isinstance(expr, str):
|
||||
yield name, expr
|
||||
elif isinstance(expr, Expr):
|
||||
yield name, expr.to_sql()
|
||||
source = expr._column_name()
|
||||
yield name, source if source is not None else expr.to_sql()
|
||||
|
||||
|
||||
def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table:
|
||||
|
||||
@@ -87,6 +87,7 @@ class PyExpr:
|
||||
def contains(self, substr: "PyExpr") -> "PyExpr": ...
|
||||
def isin(self, values: List["PyExpr"]) -> "PyExpr": ...
|
||||
def cast(self, data_type: pa.DataType) -> "PyExpr": ...
|
||||
def column_name(self) -> Optional[str]: ...
|
||||
def to_sql(self) -> str: ...
|
||||
|
||||
def expr_col(name: str) -> PyExpr: ...
|
||||
@@ -283,6 +284,7 @@ class Table:
|
||||
mode: Literal["append", "overwrite"],
|
||||
progress: Optional[Any] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult: ...
|
||||
async def update(
|
||||
self, updates: Dict[str, str], where: Optional[str]
|
||||
@@ -607,6 +609,7 @@ class PyQueryRequest:
|
||||
filter: Optional[Union[str, bytes]]
|
||||
full_text_search: Optional[FullTextQuery]
|
||||
select: Optional[Union[str, List[str]]]
|
||||
select_source_columns: Optional[Dict[str, str]]
|
||||
fast_search: Optional[bool]
|
||||
with_row_id: Optional[bool]
|
||||
use_lsm: Optional[bool]
|
||||
|
||||
@@ -249,6 +249,10 @@ class Expr:
|
||||
|
||||
# ── utilities ────────────────────────────────────────────────────────────
|
||||
|
||||
def _column_name(self) -> str | None:
|
||||
"""Return the source name when this is a bare column expression."""
|
||||
return self._inner.column_name()
|
||||
|
||||
def to_sql(self) -> str:
|
||||
"""Render the expression as a SQL string (useful for debugging)."""
|
||||
return self._inner.to_sql()
|
||||
@@ -312,7 +316,7 @@ def func(name: str, *args: ExprLike) -> Expr:
|
||||
--------
|
||||
>>> from lancedb.expr import col, func
|
||||
>>> func("lower", col("name"))
|
||||
Expr(lower(name))
|
||||
Expr(lower(`name`))
|
||||
"""
|
||||
inner_args = [_coerce(a)._inner for a in args]
|
||||
return Expr(expr_func(name, inner_args))
|
||||
|
||||
+223
-102
@@ -4,7 +4,7 @@
|
||||
"""Canonical Function values exchanged with LanceDB Enterprise services.
|
||||
|
||||
These immutable models contain client/wire state only. Catalog persistence,
|
||||
environment bake, secret resolution, and execution are owned by Sophon.
|
||||
environment bake, and execution are owned by Sophon.
|
||||
``RefreshColumnResult`` is also the backend-neutral result of a local
|
||||
expression-backed refresh job.
|
||||
"""
|
||||
@@ -12,17 +12,19 @@ expression-backed refresh job.
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import builtins
|
||||
import base64
|
||||
import functools
|
||||
import hashlib
|
||||
import importlib
|
||||
import inspect
|
||||
import symtable
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sys
|
||||
import textwrap
|
||||
import types
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from datetime import date, datetime
|
||||
from typing import (
|
||||
@@ -220,13 +222,14 @@ class PythonEnvironmentSpec(_RemoteValue):
|
||||
|
||||
kind: str
|
||||
packages: tuple[str, ...] = ()
|
||||
channels: tuple[str, ...] = ()
|
||||
path: Optional[str] = None
|
||||
modules: tuple[str, ...] = ()
|
||||
image: Optional[str] = None
|
||||
|
||||
|
||||
class PythonRuntimeSpec(_RemoteValue):
|
||||
"""Remote runtime definition with non-secret environment values.
|
||||
"""Remote runtime definition with environment values.
|
||||
|
||||
V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
|
||||
their unknown payload fields are intentionally not retained by the client.
|
||||
@@ -265,7 +268,6 @@ class FunctionVersion(_RemoteValue):
|
||||
runtime: PythonRuntimeSpec
|
||||
runtime_digest: str
|
||||
environment_digest: str
|
||||
required_secrets: tuple[str, ...] = ()
|
||||
created_at: str
|
||||
|
||||
def __call__(self, **inputs: Any) -> FunctionApplication:
|
||||
@@ -273,7 +275,7 @@ class FunctionVersion(_RemoteValue):
|
||||
|
||||
Every input must be a direct [lancedb.col][lancedb.expr.col]
|
||||
reference. The returned application is immutable and retains a
|
||||
named-struct output as one sibling group, so every row's sibling values
|
||||
named-struct output as one binding, so every row's sibling values
|
||||
come from one logical Function evaluation. Map result fields to table
|
||||
columns with
|
||||
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename],
|
||||
@@ -323,22 +325,16 @@ class FunctionVersion(_RemoteValue):
|
||||
function=FunctionVersionRef(name=self.name, version=self.version),
|
||||
inputs=tuple(bindings),
|
||||
output=self.signature.output,
|
||||
group_id=f"fg_{uuid.uuid4().hex}",
|
||||
)
|
||||
|
||||
|
||||
class FunctionRegistrationRequest(_RemoteValue):
|
||||
"""Stable remote registration envelope produced by :func:`udf`.
|
||||
|
||||
Only secret names are represented. Secret values are resolved inside the
|
||||
remote service and have no client request field.
|
||||
"""
|
||||
"""Stable remote registration envelope produced by :func:`udf`."""
|
||||
|
||||
name: str
|
||||
artifact: FunctionArtifactRequest
|
||||
signature: FunctionSignature
|
||||
runtime: PythonRuntimeSpec
|
||||
required_secrets: tuple[str, ...] = ()
|
||||
|
||||
|
||||
class FunctionVersionRef(_OpenRemoteValue):
|
||||
@@ -367,7 +363,7 @@ class ApplicationInput(_OpenRemoteValue):
|
||||
class FunctionApplication(_OpenRemoteValue):
|
||||
"""Immutable pre-declaration application of an exact Function version.
|
||||
|
||||
A named-struct output remains one grouped application through table
|
||||
A named-struct output remains one application through table
|
||||
declaration and execution.
|
||||
[FunctionApplication.rename][lancedb.functions.FunctionApplication.rename]
|
||||
records the result-field to table-column mapping without splitting sibling
|
||||
@@ -377,7 +373,6 @@ class FunctionApplication(_OpenRemoteValue):
|
||||
function: FunctionVersionRef
|
||||
inputs: tuple[ApplicationInput, ...]
|
||||
output: FunctionOutput
|
||||
group_id: str
|
||||
columns: Mapping[str, str] = Field(default_factory=dict)
|
||||
|
||||
def _known_dict(self) -> dict[str, Any]:
|
||||
@@ -449,12 +444,10 @@ class OutputMapping(_RemoteValue):
|
||||
|
||||
|
||||
class FunctionBinding(_RemoteValue):
|
||||
"""Immutable grouped binding persisted by the Enterprise table service."""
|
||||
"""Immutable Function binding persisted by the Enterprise table service."""
|
||||
|
||||
binding_id: str
|
||||
revision: _UInt64
|
||||
function: FunctionVersionRef
|
||||
group_id: str
|
||||
inputs: tuple[InputBinding, ...]
|
||||
outputs: tuple[OutputMapping, ...]
|
||||
input_schema: Optional[Mapping[str, Any]] = None
|
||||
@@ -486,62 +479,60 @@ class RefreshColumnResult(_RemoteValue):
|
||||
|
||||
|
||||
_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")
|
||||
_SECRET_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
|
||||
|
||||
_GRAMMAR_PRIMITIVES = (
|
||||
(pa.bool_(), "bool"),
|
||||
(pa.int8(), "int8"),
|
||||
(pa.int16(), "int16"),
|
||||
(pa.int32(), "int32"),
|
||||
(pa.int64(), "int64"),
|
||||
(pa.uint8(), "uint8"),
|
||||
(pa.uint16(), "uint16"),
|
||||
(pa.uint32(), "uint32"),
|
||||
(pa.uint64(), "uint64"),
|
||||
(pa.float16(), "float16"),
|
||||
(pa.float32(), "float32"),
|
||||
(pa.float64(), "float64"),
|
||||
(pa.string(), "utf8"),
|
||||
(pa.binary(), "binary"),
|
||||
(pa.date32(), "date32"),
|
||||
(pa.date64(), "date64"),
|
||||
)
|
||||
|
||||
|
||||
def _canonical_arrow_type(data_type: pa.DataType) -> str:
|
||||
primitive_types = (
|
||||
(pa.bool_(), "bool"),
|
||||
(pa.int8(), "int8"),
|
||||
(pa.int16(), "int16"),
|
||||
(pa.int32(), "int32"),
|
||||
(pa.int64(), "int64"),
|
||||
(pa.uint8(), "uint8"),
|
||||
(pa.uint16(), "uint16"),
|
||||
(pa.uint32(), "uint32"),
|
||||
(pa.uint64(), "uint64"),
|
||||
(pa.float16(), "float16"),
|
||||
(pa.float32(), "float32"),
|
||||
(pa.float64(), "float64"),
|
||||
(pa.string(), "utf8"),
|
||||
(pa.large_utf8(), "large_utf8"),
|
||||
(pa.binary(), "binary"),
|
||||
(pa.large_binary(), "large_binary"),
|
||||
(pa.date32(), "date32"),
|
||||
(pa.date64(), "date64"),
|
||||
)
|
||||
for candidate, name in primitive_types:
|
||||
"""The server's V1 Function type grammar. Anything outside it is rejected
|
||||
here rather than at registration."""
|
||||
for candidate, name in _GRAMMAR_PRIMITIVES:
|
||||
if data_type == candidate:
|
||||
return name
|
||||
if pa.types.is_fixed_size_binary(data_type):
|
||||
return f"fixed_size_binary[{data_type.byte_width}]"
|
||||
if pa.types.is_list(data_type):
|
||||
return f"list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
if pa.types.is_large_list(data_type):
|
||||
return f"large_list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
if pa.types.is_fixed_size_list(data_type):
|
||||
if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
|
||||
prefix = "list" if pa.types.is_list(data_type) else "large_list"
|
||||
return f"{prefix}<{_canonical_list_item(data_type)}>"
|
||||
if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
|
||||
return (
|
||||
f"fixed_size_list<{_canonical_arrow_type(data_type.value_type)}>"
|
||||
f"[{data_type.list_size}]"
|
||||
f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
|
||||
)
|
||||
if pa.types.is_struct(data_type):
|
||||
fields = ",".join(
|
||||
f"{field.name}:{_canonical_arrow_type(field.type)}" for field in data_type
|
||||
)
|
||||
return f"struct<{fields}>"
|
||||
if pa.types.is_timestamp(data_type):
|
||||
timezone = f",tz={data_type.tz}" if data_type.tz is not None else ""
|
||||
return f"timestamp[{data_type.unit}{timezone}]"
|
||||
if pa.types.is_time32(data_type) or pa.types.is_time64(data_type):
|
||||
return f"time[{data_type.unit}]"
|
||||
if pa.types.is_duration(data_type):
|
||||
return f"duration[{data_type.unit}]"
|
||||
if pa.types.is_decimal(data_type):
|
||||
bit_width = data_type.bit_width
|
||||
return f"decimal{bit_width}({data_type.precision},{data_type.scale})"
|
||||
raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
|
||||
|
||||
|
||||
def _canonical_list_item(data_type: pa.DataType) -> str:
|
||||
"""The grammar names only the item type; it always means a non-nullable
|
||||
child called `item`, so any other child metadata cannot be represented."""
|
||||
child = data_type.value_field
|
||||
if child.name != "item" or child.nullable or child.metadata:
|
||||
raise TypeError(
|
||||
"unsupported Arrow type for Function signature: list items must be a "
|
||||
f"non-nullable field named 'item', got {child}"
|
||||
)
|
||||
return _canonical_arrow_type(child.type)
|
||||
|
||||
|
||||
def _list_of(item: pa.DataType) -> pa.DataType:
|
||||
return pa.list_(pa.field("item", item, nullable=False))
|
||||
|
||||
|
||||
def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
|
||||
nullable = False
|
||||
origin = get_origin(annotation)
|
||||
@@ -589,7 +580,7 @@ def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
|
||||
value_type, value_nullable = _annotation_type(arguments[0])
|
||||
if value_nullable:
|
||||
raise TypeError("nullable Function list elements are not supported")
|
||||
return pa.list_(value_type), nullable
|
||||
return _list_of(value_type), nullable
|
||||
raise TypeError(f"unsupported Function annotation: {annotation!r}")
|
||||
|
||||
|
||||
@@ -736,6 +727,104 @@ def _literal_source(value: Any) -> str:
|
||||
)
|
||||
|
||||
|
||||
_DYNAMIC_NAMESPACE_ACCESS = frozenset(
|
||||
{"globals", "locals", "vars", "eval", "exec", "compile", "__import__"}
|
||||
)
|
||||
# Modules that hand out namespaces (`sys.modules`, `builtins`, importers,
|
||||
# introspection). The artifact's module namespace holds only the names it was
|
||||
# packaged with, so reaching around it cannot be represented.
|
||||
_NAMESPACE_MODULES = frozenset(
|
||||
{"sys", "builtins", "importlib", "inspect", "gc", "ctypes", "types"}
|
||||
)
|
||||
|
||||
|
||||
def _namespace_acquisition(
|
||||
definition: ast.FunctionDef, references: set[str]
|
||||
) -> list[str]:
|
||||
found = set(references & _DYNAMIC_NAMESPACE_ACCESS)
|
||||
for node in ast.walk(definition):
|
||||
if isinstance(node, ast.Import):
|
||||
found.update(
|
||||
alias.name
|
||||
for alias in node.names
|
||||
if alias.name.split(".")[0] in _NAMESPACE_MODULES
|
||||
)
|
||||
elif isinstance(node, ast.ImportFrom) and node.module:
|
||||
if node.module.split(".")[0] in _NAMESPACE_MODULES:
|
||||
found.add(node.module)
|
||||
return sorted(found)
|
||||
|
||||
|
||||
def _module_references(module_source: str) -> set[str]:
|
||||
"""Names any scope in `module_source` binds or loads at module scope.
|
||||
Python's own scope analysis on the exact text that ships: free variables
|
||||
belong to an enclosing scope inside the function, and postponed
|
||||
annotations are not runtime loads."""
|
||||
|
||||
def visit(table: symtable.SymbolTable, found: set[str]) -> None:
|
||||
for symbol in table.get_symbols():
|
||||
if symbol.is_global() and (
|
||||
symbol.is_referenced() or symbol.is_declared_global()
|
||||
):
|
||||
found.add(symbol.get_name())
|
||||
for child in table.get_children():
|
||||
visit(child, found)
|
||||
|
||||
found: set[str] = set()
|
||||
for table in symtable.symtable(module_source, "<udf>", "exec").get_children():
|
||||
visit(table, found)
|
||||
return found
|
||||
|
||||
|
||||
def _global_source(name: str, value: Any) -> str:
|
||||
"""One module-level line that rebinds `name` to `value` in the artifact:
|
||||
an import for modules and importable classes/functions, a literal otherwise."""
|
||||
if isinstance(value, types.ModuleType):
|
||||
if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
|
||||
raise ValueError(
|
||||
f"@udf cannot package dynamic namespace access: {value.__name__!r}"
|
||||
)
|
||||
try:
|
||||
imported = importlib.import_module(value.__name__)
|
||||
except ImportError:
|
||||
imported = None
|
||||
if imported is not value:
|
||||
raise TypeError(
|
||||
f"Function source references module {name!r} that does not import "
|
||||
f"as {value.__name__!r}"
|
||||
)
|
||||
return f"import {value.__name__} as {name}"
|
||||
module_name = getattr(value, "__module__", None)
|
||||
qualname = getattr(value, "__qualname__", None)
|
||||
if (
|
||||
isinstance(module_name, str)
|
||||
and isinstance(qualname, str)
|
||||
and module_name != "__main__"
|
||||
and "." not in qualname
|
||||
and "<" not in qualname
|
||||
):
|
||||
try:
|
||||
imported = getattr(importlib.import_module(module_name), qualname)
|
||||
except (ImportError, AttributeError):
|
||||
imported = None
|
||||
if imported is value:
|
||||
return f"from {module_name} import {qualname} as {name}"
|
||||
return f"{name} = {_literal_source(value)}"
|
||||
|
||||
|
||||
def _is_recursive_reference(function: Callable[..., Any], name: str) -> bool:
|
||||
"""`name` inside the body means the function itself unless the module has
|
||||
since bound it to something else."""
|
||||
if name != function.__name__:
|
||||
return False
|
||||
bound = function.__globals__.get(name, function)
|
||||
if bound is function:
|
||||
return True
|
||||
# The decorator's own result is the one wrapper known to call `function`
|
||||
# unchanged; any other binding may behave differently from a self-call.
|
||||
return type(bound) is UdfDefinition and bound._function is function
|
||||
|
||||
|
||||
def _package_source(function: Callable[..., Any]) -> bytes:
|
||||
if not inspect.isfunction(function) or inspect.iscoroutinefunction(function):
|
||||
raise TypeError("@udf requires a synchronous Python function")
|
||||
@@ -760,23 +849,46 @@ def _package_source(function: Callable[..., Any]) -> bytes:
|
||||
closure = inspect.getclosurevars(function)
|
||||
if closure.nonlocals:
|
||||
raise ValueError("@udf cannot package functions that capture closure values")
|
||||
if closure.unbound:
|
||||
raise ValueError(
|
||||
f"@udf source contains unresolved global names: {sorted(closure.unbound)!r}"
|
||||
)
|
||||
globals_source = []
|
||||
for name, value in sorted(closure.globals.items()):
|
||||
if isinstance(value, types.ModuleType):
|
||||
globals_source.append(f"import {value.__name__} as {name}")
|
||||
else:
|
||||
globals_source.append(f"{name} = {_literal_source(value)}")
|
||||
|
||||
function_source = ast.unparse(definition)
|
||||
parts = ["from __future__ import annotations"]
|
||||
module_header = "from __future__ import annotations"
|
||||
references = _module_references(f"{module_header}\n\n{function_source}\n")
|
||||
dynamic = _namespace_acquisition(definition, references)
|
||||
if dynamic:
|
||||
raise ValueError(f"@udf cannot package dynamic namespace access: {dynamic!r}")
|
||||
# Resolve every module-scope reference the way the interpreter would: the
|
||||
# function's own globals first (a module global may shadow a builtin, and
|
||||
# nested scopes are not visible to getclosurevars), then its builtins.
|
||||
# The artifact runs under the standard builtins; only the exact mapping is
|
||||
# provably equivalent (a subclass or copy can change lookups and hooks).
|
||||
if function.__builtins__ is not vars(builtins):
|
||||
raise ValueError("@udf cannot package a non-standard builtins environment")
|
||||
globals_source = []
|
||||
unresolved = []
|
||||
for name in sorted(references):
|
||||
if name == function.__name__:
|
||||
if not _is_recursive_reference(function, name):
|
||||
raise ValueError(
|
||||
f"@udf cannot package {name!r}: the module binds that name to "
|
||||
"another value, which the artifact's own definition would shadow"
|
||||
)
|
||||
continue
|
||||
if name in function.__globals__:
|
||||
globals_source.append(_global_source(name, function.__globals__[name]))
|
||||
elif hasattr(builtins, name):
|
||||
pass
|
||||
else:
|
||||
unresolved.append(name)
|
||||
if unresolved:
|
||||
raise ValueError(
|
||||
f"@udf source contains unresolved global names: {unresolved!r}"
|
||||
)
|
||||
|
||||
parts = [module_header]
|
||||
if globals_source:
|
||||
parts.extend(["", *globals_source])
|
||||
parts.extend(["", function_source, ""])
|
||||
return "\n".join(parts).encode("utf-8")
|
||||
packaged = "\n".join(parts)
|
||||
return packaged.encode("utf-8")
|
||||
|
||||
|
||||
class UdfDefinition:
|
||||
@@ -797,33 +909,32 @@ class UdfDefinition:
|
||||
output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
|
||||
pip: tuple[str, ...],
|
||||
env: Mapping[str, str],
|
||||
secrets: tuple[str, ...],
|
||||
python_version: Optional[str],
|
||||
conda: tuple[str, ...] = (),
|
||||
conda_channels: tuple[str, ...] = (),
|
||||
):
|
||||
function_name = name or function.__name__
|
||||
if not _FUNCTION_NAME.fullmatch(function_name):
|
||||
raise ValueError(f"invalid Function name: {function_name!r}")
|
||||
packages = tuple(sorted(set(pip)))
|
||||
if pip and conda:
|
||||
raise ValueError("a Function environment is pip or conda, not both")
|
||||
if conda_channels and not conda:
|
||||
raise ValueError("conda_channels requires conda packages")
|
||||
packages = tuple(sorted(set(conda if conda else pip)))
|
||||
if any(not package or package != package.strip() for package in packages):
|
||||
raise ValueError("pip requirements must be non-empty and trimmed")
|
||||
raise ValueError("package requirements must be non-empty and trimmed")
|
||||
if conda:
|
||||
environment_spec = PythonEnvironmentSpec(
|
||||
kind="conda", packages=packages, channels=tuple(conda_channels)
|
||||
)
|
||||
else:
|
||||
environment_spec = PythonEnvironmentSpec(kind="pip", packages=packages)
|
||||
environment = dict(env)
|
||||
if any(
|
||||
not isinstance(key, str) or not isinstance(value, str)
|
||||
for key, value in environment.items()
|
||||
):
|
||||
raise TypeError("Function env keys and values must be strings")
|
||||
required_secrets = tuple(sorted(set(secrets)))
|
||||
invalid_secrets = [
|
||||
secret for secret in required_secrets if not _SECRET_NAME.fullmatch(secret)
|
||||
]
|
||||
if invalid_secrets:
|
||||
raise ValueError(f"invalid Function secret names: {invalid_secrets!r}")
|
||||
overlap = set(environment) & set(required_secrets)
|
||||
if overlap:
|
||||
raise ValueError(
|
||||
f"Function env and secret names must be disjoint: {sorted(overlap)!r}"
|
||||
)
|
||||
|
||||
signature = _infer_signature(function, input_schema, output_schema)
|
||||
source = _package_source(function)
|
||||
digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
|
||||
@@ -831,7 +942,7 @@ class UdfDefinition:
|
||||
kind="python",
|
||||
python_version=python_version
|
||||
or f"{sys.version_info.major}.{sys.version_info.minor}",
|
||||
environment=PythonEnvironmentSpec(kind="pip", packages=packages),
|
||||
environment=environment_spec,
|
||||
env=environment,
|
||||
)
|
||||
self._function = function
|
||||
@@ -852,7 +963,6 @@ class UdfDefinition:
|
||||
),
|
||||
signature=signature,
|
||||
runtime=runtime,
|
||||
required_secrets=required_secrets,
|
||||
)
|
||||
functools.update_wrapper(self, function)
|
||||
|
||||
@@ -878,8 +988,9 @@ def udf(
|
||||
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
|
||||
pip: tuple[str, ...] | list[str] = (),
|
||||
env: Optional[Mapping[str, str]] = None,
|
||||
secrets: tuple[str, ...] | list[str] = (),
|
||||
python_version: Optional[str] = None,
|
||||
conda: tuple[str, ...] | list[str] = (),
|
||||
conda_channels: tuple[str, ...] | list[str] = (),
|
||||
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...
|
||||
|
||||
|
||||
@@ -891,8 +1002,9 @@ def udf(
|
||||
output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
|
||||
pip: tuple[str, ...] | list[str] = (),
|
||||
env: Optional[Mapping[str, str]] = None,
|
||||
secrets: tuple[str, ...] | list[str] = (),
|
||||
python_version: Optional[str] = None,
|
||||
conda: tuple[str, ...] | list[str] = (),
|
||||
conda_channels: tuple[str, ...] | list[str] = (),
|
||||
):
|
||||
"""Prepare a scalar Python callable for remote Function registration.
|
||||
|
||||
@@ -915,14 +1027,22 @@ def udf(
|
||||
provided together with ``input_schema``.
|
||||
pip : sequence of str, optional
|
||||
Pip requirements for the remote environment.
|
||||
conda : sequence of str, optional
|
||||
Conda packages for the remote environment, instead of ``pip``.
|
||||
conda_channels : sequence of str, optional
|
||||
Conda channels in priority order; requires ``conda``.
|
||||
env : mapping of str to str, optional
|
||||
Non-secret environment variables. Use ``secrets`` for credentials.
|
||||
secrets : sequence of str, optional
|
||||
Names of secrets resolved by the remote service. Secret values are not
|
||||
accepted by this API or included in the registration request.
|
||||
Environment variables included in the Function definition.
|
||||
python_version : str, optional
|
||||
Remote Python major/minor version. Defaults to the client version.
|
||||
|
||||
The packaged artifact is a snapshot: the function source plus exactly
|
||||
the module-level names it references (modules as imports, importable
|
||||
classes and functions as imports, literals inline). Code that reaches the
|
||||
module namespace another way -- ``globals()``/``eval``, ``sys.modules``,
|
||||
``builtins`` -- is rejected where it can be seen and otherwise
|
||||
unsupported; closures and a non-standard ``__builtins__`` are rejected.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UdfDefinition
|
||||
@@ -934,7 +1054,7 @@ def udf(
|
||||
Examples
|
||||
--------
|
||||
>>> from lancedb import udf
|
||||
>>> @udf(pip=["numpy==2.2.0"], secrets=["MODEL_TOKEN"])
|
||||
>>> @udf(pip=["numpy==2.2.0"])
|
||||
... def score(value: float) -> float:
|
||||
... return value * 2
|
||||
>>> score(1.5)
|
||||
@@ -949,8 +1069,9 @@ def udf(
|
||||
output_schema=output_schema,
|
||||
pip=tuple(pip),
|
||||
env={} if env is None else env,
|
||||
secrets=tuple(secrets),
|
||||
python_version=python_version,
|
||||
conda=tuple(conda),
|
||||
conda_channels=tuple(conda_channels),
|
||||
)
|
||||
|
||||
if function is None:
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing import List, Literal, Optional
|
||||
from ._lancedb import (
|
||||
IndexConfig,
|
||||
)
|
||||
from .query import DocumentGranularity
|
||||
from .types import BaseTokenizerType
|
||||
|
||||
lang_mapping = {
|
||||
@@ -121,6 +122,11 @@ class FTS:
|
||||
|
||||
>>> config = FTS(block_size=256)
|
||||
|
||||
Create an index that treats each deepest-list element as one document:
|
||||
|
||||
>>> from lancedb.query import DocumentGranularity
|
||||
>>> config = FTS(document_granularity=DocumentGranularity.LIST_ELEMENT)
|
||||
|
||||
Attributes
|
||||
----------
|
||||
with_position : bool, default False
|
||||
@@ -172,6 +178,11 @@ class FTS:
|
||||
roughly half of the available CPU cores. The effective value is
|
||||
limited by the available compute capacity. This build-only setting is
|
||||
not persisted with the index and does not apply to remote tables.
|
||||
document_granularity : DocumentGranularity, default ROW
|
||||
``ROW`` treats the selected text in one table row as one document.
|
||||
``LIST_ELEMENT`` treats each element of the deepest list on the indexed
|
||||
field path as one document and returns its physical coordinates in
|
||||
``_doc_index`` for matching queries.
|
||||
|
||||
Notes
|
||||
-----
|
||||
@@ -196,6 +207,7 @@ class FTS:
|
||||
custom_stop_words: Optional[List[str]] = None
|
||||
memory_limit: Optional[int] = None
|
||||
num_workers: Optional[int] = None
|
||||
document_granularity: DocumentGranularity = DocumentGranularity.ROW
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -391,6 +391,15 @@ def _table_to_pickle_state(table: Table) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _drop_base_version(permutation_data: pa.Table) -> pa.Table:
|
||||
"""Strip the recorded base version so the reader leaves the base table unpinned."""
|
||||
metadata = dict(permutation_data.schema.metadata or {})
|
||||
if metadata.pop(b"base_version", None) is None:
|
||||
return permutation_data
|
||||
metadata.pop(b"base_branch", None)
|
||||
return permutation_data.replace_schema_metadata(metadata)
|
||||
|
||||
|
||||
def _table_from_pickle_state(state: dict[str, Any]) -> Table:
|
||||
from . import connect
|
||||
|
||||
@@ -679,11 +688,15 @@ class Permutation:
|
||||
from . import connect
|
||||
|
||||
connection_factory = state["connection_factory"]
|
||||
rebuilt_base = False
|
||||
if connection_factory is not None:
|
||||
base_table = connection_factory(state["base_table_name"])
|
||||
elif "base_table_state" in state:
|
||||
base_table = _table_from_pickle_state(state["base_table_state"])
|
||||
base_state = state["base_table_state"]
|
||||
rebuilt_base = base_state["kind"] == "memory"
|
||||
base_table = _table_from_pickle_state(base_state)
|
||||
elif "base_table_data" in state:
|
||||
rebuilt_base = True
|
||||
# In-memory base table inlined into the pickle; rebuild the same
|
||||
# way we rebuild the in-memory permutation table.
|
||||
mem_db = connect("memory://")
|
||||
@@ -701,11 +714,14 @@ class Permutation:
|
||||
)
|
||||
|
||||
permutation_table: Optional[Table] = None
|
||||
if state["permutation_data"] is not None:
|
||||
permutation_data = state["permutation_data"]
|
||||
if permutation_data is not None:
|
||||
if rebuilt_base:
|
||||
# The base table was materialized from Arrow, so it is a fresh
|
||||
# single-version dataset and the recorded pin cannot resolve on it.
|
||||
permutation_data = _drop_base_version(permutation_data)
|
||||
mem_db = connect("memory://")
|
||||
permutation_table = mem_db.create_table(
|
||||
"permutation", state["permutation_data"]
|
||||
)
|
||||
permutation_table = mem_db.create_table("permutation", permutation_data)
|
||||
|
||||
self.base_table = base_table
|
||||
self.permutation_table = permutation_table
|
||||
|
||||
@@ -167,6 +167,12 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||||
return {"columns": projection}
|
||||
|
||||
|
||||
def _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
|
||||
if req.select_source_columns is not None:
|
||||
return req.select_source_columns
|
||||
return req.select
|
||||
|
||||
|
||||
def _scanner_kwargs_for_query(
|
||||
query: Query,
|
||||
blob_mode: BlobMode,
|
||||
@@ -375,6 +381,13 @@ class FullTextOperator(str, Enum):
|
||||
OR = "OR"
|
||||
|
||||
|
||||
class DocumentGranularity(str, Enum):
|
||||
"""The unit treated as one full-text-search document."""
|
||||
|
||||
ROW = "row"
|
||||
LIST_ELEMENT = "list_element"
|
||||
|
||||
|
||||
class Occur(str, Enum):
|
||||
SHOULD = "SHOULD"
|
||||
MUST = "MUST"
|
||||
@@ -478,6 +491,10 @@ class MatchQuery(FullTextQuery):
|
||||
prefix_length : int, optional
|
||||
The number of beginning characters being unchanged for fuzzy matching.
|
||||
This is useful to achieve prefix matching.
|
||||
document_granularity : DocumentGranularity, optional
|
||||
Explicitly select row or deepest-list-element documents. If omitted,
|
||||
the indexed granularity is inferred. When both granularities are indexed
|
||||
for the field, this must be specified. With no index, row granularity is used.
|
||||
"""
|
||||
|
||||
query: str
|
||||
@@ -487,6 +504,9 @@ class MatchQuery(FullTextQuery):
|
||||
max_expansions: int = pydantic.Field(50, kw_only=True)
|
||||
operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
|
||||
prefix_length: int = pydantic.Field(0, kw_only=True)
|
||||
document_granularity: Optional[DocumentGranularity] = pydantic.Field(
|
||||
None, kw_only=True
|
||||
)
|
||||
|
||||
def query_type(self) -> FullTextQueryType:
|
||||
return FullTextQueryType.MATCH
|
||||
@@ -503,11 +523,20 @@ class PhraseQuery(FullTextQuery):
|
||||
The query string to match against.
|
||||
column : str
|
||||
The name of the column to match against.
|
||||
slop : int, default 0
|
||||
The maximum number of intervening positions permitted in the phrase.
|
||||
document_granularity : DocumentGranularity, optional
|
||||
Explicitly select row or deepest-list-element documents. If omitted,
|
||||
the indexed granularity is inferred. When both granularities are indexed
|
||||
for the field, this must be specified. With no index, row granularity is used.
|
||||
"""
|
||||
|
||||
query: str
|
||||
column: str
|
||||
slop: int = pydantic.Field(0, kw_only=True)
|
||||
document_granularity: Optional[DocumentGranularity] = pydantic.Field(
|
||||
None, kw_only=True
|
||||
)
|
||||
|
||||
def query_type(self) -> FullTextQueryType:
|
||||
return FullTextQueryType.MATCH_PHRASE
|
||||
@@ -2776,15 +2805,16 @@ class AsyncQueryBase(object):
|
||||
|
||||
req = self._inner.to_query_request()
|
||||
schema = await self._table.schema()
|
||||
projection = _query_request_projection(req)
|
||||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
schema,
|
||||
req.select,
|
||||
projection,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if not self._blob_auto_row_id:
|
||||
self._blob_paths = ()
|
||||
return
|
||||
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
|
||||
self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
|
||||
self._inner.with_row_id()
|
||||
|
||||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||||
@@ -3378,9 +3408,10 @@ class AsyncQuery(AsyncStandardQuery):
|
||||
pass in multiple vectors. When multiple vectors are passed in, if the vector
|
||||
column is with multivector type, then the vectors will be treated as a single
|
||||
query. Or the vectors will be treated as multiple queries, this can be useful
|
||||
if you want to find the nearest vectors to multiple query vectors.
|
||||
This is not expected to be faster than making multiple queries concurrently;
|
||||
it is just a convenience method. If multiple vectors are passed in then
|
||||
if you want to find the nearest vectors to multiple query vectors. Flat
|
||||
searches share one table scan across the query vectors, avoiding the scan
|
||||
and memory amplification of making multiple queries concurrently. If
|
||||
multiple vectors are passed in then
|
||||
an additional column `query_index` will be added to the results. This column
|
||||
will contain the index of the query vector that the result is nearest to.
|
||||
"""
|
||||
@@ -3509,8 +3540,8 @@ class AsyncFTSQuery(AsyncStandardQuery):
|
||||
|
||||
Typically, a single vector is passed in as the query. However, you can also
|
||||
pass in multiple vectors. This can be useful if you want to find the nearest
|
||||
vectors to multiple query vectors. This is not expected to be faster than
|
||||
making multiple queries concurrently; it is just a convenience method.
|
||||
vectors to multiple query vectors. Flat searches share one table scan across
|
||||
the query vectors instead of issuing concurrent full scans.
|
||||
If multiple vectors are passed in then an additional column `query_index`
|
||||
will be added to the results. This column will contain the index of the
|
||||
query vector that the result is nearest to.
|
||||
@@ -3870,14 +3901,15 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||||
blob_paths: tuple[str, ...] = ()
|
||||
if self._table is not None:
|
||||
schema = await self._table.schema()
|
||||
projection = _query_request_projection(req)
|
||||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||||
schema,
|
||||
req.select,
|
||||
projection,
|
||||
with_row_id=self._with_row_id,
|
||||
)
|
||||
if blob_auto_row_id:
|
||||
blob_paths = tuple(
|
||||
blob_v2_projection_sources(schema, req.select).keys()
|
||||
blob_v2_projection_sources(schema, projection).keys()
|
||||
)
|
||||
self._blob_auto_row_id = blob_auto_row_id
|
||||
self._blob_paths = blob_paths
|
||||
|
||||
@@ -36,6 +36,7 @@ from lancedb._lancedb import (
|
||||
UpdateResult,
|
||||
)
|
||||
from lancedb.embeddings.base import EmbeddingFunctionConfig
|
||||
from lancedb.expr import Expr
|
||||
from lancedb.index import (
|
||||
FTS,
|
||||
BTree,
|
||||
@@ -61,11 +62,20 @@ from lancedb.table import _normalize_progress
|
||||
|
||||
from ..query import (
|
||||
AnalyzePlanDistributedMetrics,
|
||||
DocumentGranularity,
|
||||
LanceQueryBuilder,
|
||||
LanceTakeQueryBuilder,
|
||||
LanceVectorQueryBuilder,
|
||||
)
|
||||
from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
|
||||
from ..table import (
|
||||
AsyncTable,
|
||||
BlobMode,
|
||||
Branches,
|
||||
IndexStatistics,
|
||||
Query,
|
||||
Table,
|
||||
Tags,
|
||||
)
|
||||
from ..types import BaseTokenizerType
|
||||
|
||||
|
||||
@@ -349,6 +359,7 @@ class RemoteTable(Table):
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
block_size: int = 128,
|
||||
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
"""Create a full-text search index on a column.
|
||||
@@ -371,6 +382,7 @@ class RemoteTable(Table):
|
||||
ngram_max_length=ngram_max_length,
|
||||
prefix_only=prefix_only,
|
||||
block_size=block_size,
|
||||
document_granularity=document_granularity,
|
||||
)
|
||||
LOOP.run(
|
||||
self._table.create_index(
|
||||
@@ -610,6 +622,7 @@ class RemoteTable(Table):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -642,6 +655,8 @@ class RemoteTable(Table):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Not supported on LanceDB Cloud. Setting this raises.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -658,6 +673,7 @@ class RemoteTable(Table):
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -856,7 +872,7 @@ class RemoteTable(Table):
|
||||
|
||||
def update(
|
||||
self,
|
||||
where: Optional[str] = None,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -867,9 +883,11 @@ class RemoteTable(Table):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
|
||||
+101
-34
@@ -4,30 +4,34 @@
|
||||
|
||||
"""Schema helpers for Lance blob columns."""
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.ipc
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lance.blob import BlobType as BlobType
|
||||
|
||||
_BLOB_EXTENSION_NAME = "lance.blob.v2"
|
||||
_BLOB_V1_KEY = "lance-encoding:blob"
|
||||
_ARROW_EXT_NAME_KEY = "ARROW:extension:name"
|
||||
_BLOB_V2_STORAGE_TYPE = pa.struct(
|
||||
[
|
||||
pa.field("data", pa.large_binary(), nullable=True),
|
||||
pa.field("uri", pa.utf8(), nullable=True),
|
||||
pa.field("position", pa.uint64(), nullable=True),
|
||||
pa.field("size", pa.uint64(), nullable=True),
|
||||
]
|
||||
)
|
||||
_resolved_blob_type = None
|
||||
|
||||
|
||||
class BlobType(pa.ExtensionType):
|
||||
"""PyArrow extension type for a Lance blob v2 column.
|
||||
|
||||
Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files`
|
||||
for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes.
|
||||
"""
|
||||
class _FallbackBlobType(pa.ExtensionType):
|
||||
"""lance.blob.v2 extension type used when pylance is not installed."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
storage_type = pa.struct(
|
||||
[
|
||||
pa.field("data", pa.large_binary(), nullable=True),
|
||||
pa.field("uri", pa.utf8(), nullable=True),
|
||||
pa.field("position", pa.uint64(), nullable=True),
|
||||
pa.field("size", pa.uint64(), nullable=True),
|
||||
]
|
||||
)
|
||||
super().__init__(storage_type, _BLOB_EXTENSION_NAME)
|
||||
pa.ExtensionType.__init__(self, _BLOB_V2_STORAGE_TYPE, _BLOB_EXTENSION_NAME)
|
||||
|
||||
def __arrow_ext_serialize__(self) -> bytes:
|
||||
return b""
|
||||
@@ -35,23 +39,16 @@ class BlobType(pa.ExtensionType):
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(
|
||||
cls, storage_type: pa.DataType, serialized: bytes
|
||||
) -> "BlobType":
|
||||
) -> "_FallbackBlobType":
|
||||
return cls()
|
||||
|
||||
def __reduce__(self):
|
||||
# Ensure pickle round-trips on older pyarrow (apache/arrow#35599).
|
||||
return type(self).__arrow_ext_deserialize__, (
|
||||
self.storage_type,
|
||||
self.__arrow_ext_serialize__(),
|
||||
)
|
||||
|
||||
|
||||
try:
|
||||
pa.register_extension_type(BlobType()) # type: ignore[arg-type]
|
||||
except pa.ArrowKeyError:
|
||||
pass
|
||||
|
||||
|
||||
def _metadata_value(metadata: dict, key: str):
|
||||
return metadata.get(key.encode()) or metadata.get(key)
|
||||
|
||||
@@ -92,43 +89,105 @@ def is_blob_like_field(field: pa.Field) -> bool:
|
||||
return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {})
|
||||
|
||||
|
||||
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]:
|
||||
paths: list[str] = []
|
||||
def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[tuple[str, bool]]:
|
||||
"""Walk the schema and return (path, has_list_ancestor) for each blob field."""
|
||||
paths: list[tuple[str, bool]] = []
|
||||
|
||||
def walk(fields, prefix: str) -> None:
|
||||
def walk(fields, prefix: str, has_list_ancestor: bool) -> None:
|
||||
for field in fields:
|
||||
path = f"{prefix}.{field.name}" if prefix else field.name
|
||||
if is_blob(field):
|
||||
paths.append(path)
|
||||
paths.append((path, has_list_ancestor))
|
||||
elif pa.types.is_struct(field.type):
|
||||
walk(field.type, path)
|
||||
walk(field.type, path, has_list_ancestor)
|
||||
elif (
|
||||
pa.types.is_list(field.type)
|
||||
or pa.types.is_large_list(field.type)
|
||||
or pa.types.is_fixed_size_list(field.type)
|
||||
):
|
||||
walk([field.type.value_field], path)
|
||||
walk([field.type.value_field], path, True)
|
||||
|
||||
walk(schema, "")
|
||||
walk(schema, "", False)
|
||||
return paths
|
||||
|
||||
|
||||
def blob_column_paths(schema: pa.Schema) -> list[str]:
|
||||
"""Dotted paths of blob-like columns (v2 extension or legacy metadata)."""
|
||||
return _collect_blob_paths(schema, is_blob_like_field)
|
||||
return [path for path, _ in _collect_blob_paths(schema, is_blob_like_field)]
|
||||
|
||||
|
||||
def blob_v2_column_paths(schema: pa.Schema) -> list[str]:
|
||||
return _collect_blob_paths(schema, is_blob_v2_field)
|
||||
return [path for path, _ in _collect_blob_paths(schema, is_blob_v2_field)]
|
||||
|
||||
|
||||
def row_addressable_blob_v2_paths(schema: pa.Schema) -> list[str]:
|
||||
"""Blob v2 paths with one blob addressable by table row id.
|
||||
|
||||
``fetch_blobs`` and the descriptor row-id ride-along address one blob per
|
||||
row, so a blob inside a list container has no row-id slot and no fetch
|
||||
path. Those columns still store and query as raw descriptors.
|
||||
"""
|
||||
return [
|
||||
path
|
||||
for path, has_list_ancestor in _collect_blob_paths(schema, is_blob_v2_field)
|
||||
if not has_list_ancestor
|
||||
]
|
||||
|
||||
|
||||
def schema_has_blob_field(schema: pa.Schema) -> bool:
|
||||
return bool(blob_column_paths(schema))
|
||||
|
||||
|
||||
def _deserialize_registered_type(extension_type: pa.ExtensionType) -> pa.DataType:
|
||||
"""Return the type Arrow reconstructs for this extension name."""
|
||||
schema = pa.schema([pa.field("value", extension_type)])
|
||||
restored = pa.ipc.read_schema(schema.serialize())
|
||||
return restored.field("value").type
|
||||
|
||||
|
||||
def _resolve_blob_type():
|
||||
"""Return the BlobType class this process should use.
|
||||
|
||||
pylance's class when it owns the lance.blob.v2 registry entry,
|
||||
otherwise LanceDB's fallback. A different registered class is an error.
|
||||
"""
|
||||
global _resolved_blob_type
|
||||
if _resolved_blob_type is not None:
|
||||
return _resolved_blob_type
|
||||
try:
|
||||
blob_module = importlib.import_module("lance.blob")
|
||||
except ModuleNotFoundError as err:
|
||||
if err.name not in ("lance", "lance.blob"):
|
||||
raise
|
||||
else:
|
||||
blob_type = getattr(blob_module, "BlobType", None)
|
||||
if blob_type is not None:
|
||||
registered_type = _deserialize_registered_type(blob_type())
|
||||
if type(registered_type) is not blob_type:
|
||||
registered_cls = type(registered_type)
|
||||
raise ValueError(
|
||||
"lance.blob.v2 is already registered by "
|
||||
f"{registered_cls.__module__}.{registered_cls.__qualname__}"
|
||||
)
|
||||
_resolved_blob_type = blob_type
|
||||
return blob_type
|
||||
try:
|
||||
pa.register_extension_type(_FallbackBlobType()) # type: ignore[arg-type]
|
||||
except pa.ArrowKeyError as err:
|
||||
raise ValueError(
|
||||
"lance.blob.v2 is already registered by another extension class"
|
||||
) from err
|
||||
_resolved_blob_type = _FallbackBlobType
|
||||
return _resolved_blob_type
|
||||
|
||||
|
||||
def blob(name: str, nullable: bool = True) -> pa.Field:
|
||||
"""Create a Lance blob v2 column field."""
|
||||
return pa.field(name, BlobType(), nullable=nullable)
|
||||
"""Create a Lance blob v2 column field.
|
||||
|
||||
When pylance is installed this is ``lance.blob.BlobType``.
|
||||
"""
|
||||
blob_type = _resolve_blob_type()
|
||||
return pa.field(name, blob_type(), nullable=nullable)
|
||||
|
||||
|
||||
def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType:
|
||||
@@ -155,3 +214,11 @@ def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataTyp
|
||||
... ])
|
||||
"""
|
||||
return pa.list_(value_type, dimension)
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "BlobType":
|
||||
blob_type = _resolve_blob_type()
|
||||
globals()["BlobType"] = blob_type
|
||||
return blob_type
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
@@ -19,6 +19,7 @@ above.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import heapq
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
@@ -29,17 +30,18 @@ from collections import deque
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from copy import deepcopy
|
||||
from multiprocessing import RawArray
|
||||
from typing import Any, Callable, cast, Iterator, Literal, Optional, Union
|
||||
from typing import Any, Callable, cast, Iterator, Literal, NamedTuple, Optional, Union
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import torch
|
||||
from torch.utils.data import IterableDataset, get_worker_info
|
||||
from torch.utils.data import DataLoader, IterableDataset, get_worker_info
|
||||
|
||||
from .permutation import (
|
||||
Permutation,
|
||||
Transforms,
|
||||
permutation_builder,
|
||||
_drop_base_version,
|
||||
_table_from_pickle_state,
|
||||
_table_to_pickle_state,
|
||||
)
|
||||
@@ -55,6 +57,155 @@ DEFAULT_READ_BATCH_SIZE = 64
|
||||
DEFAULT_PREFETCH_BATCHES = 4
|
||||
|
||||
|
||||
class _WorkerSample(NamedTuple):
|
||||
data: Any
|
||||
dataset: "StreamingDataset"
|
||||
|
||||
|
||||
class _WorkerBatch(NamedTuple):
|
||||
data: Any
|
||||
state: dict
|
||||
|
||||
|
||||
class _ConsumerIteratorLease(NamedTuple):
|
||||
owner_token: int
|
||||
owner_thread: int
|
||||
|
||||
|
||||
class _CheckpointCollate:
|
||||
"""Attach the worker's post-fetch state to a collated batch."""
|
||||
|
||||
def __init__(self, collate_fn: Callable):
|
||||
self._collate_fn = collate_fn
|
||||
|
||||
def __call__(self, samples):
|
||||
try:
|
||||
if isinstance(samples, list):
|
||||
if not samples:
|
||||
return _WorkerBatch(self._collate_fn(samples), {})
|
||||
worker_samples = samples
|
||||
data = self._collate_fn([sample.data for sample in worker_samples])
|
||||
dataset = worker_samples[-1].dataset
|
||||
else:
|
||||
data = self._collate_fn(samples.data)
|
||||
dataset = samples.dataset
|
||||
except StopIteration as exc:
|
||||
raise RuntimeError(
|
||||
"collate_fn raised StopIteration before returning a batch"
|
||||
) from exc
|
||||
return _WorkerBatch(data, dataset._checkpoint_snapshot())
|
||||
|
||||
|
||||
class _StreamingDatasetAdapter(IterableDataset):
|
||||
"""Yield private sample wrappers for :class:`StreamingDataLoader`."""
|
||||
|
||||
def __init__(self, dataset: "StreamingDataset"):
|
||||
super().__init__()
|
||||
self.dataset = dataset
|
||||
|
||||
def __iter__(self):
|
||||
for sample in self.dataset._iter(consumer_checkpoint_transport=True):
|
||||
yield _WorkerSample(sample, self.dataset)
|
||||
|
||||
def __getattr__(self, name):
|
||||
dataset = self.__dict__.get("dataset")
|
||||
if dataset is None:
|
||||
raise AttributeError(name)
|
||||
return getattr(dataset, name)
|
||||
|
||||
|
||||
class _ConsumerCommitIterator:
|
||||
def __init__(
|
||||
self,
|
||||
iterator,
|
||||
dataset: "StreamingDataset",
|
||||
*,
|
||||
owner_token: int,
|
||||
require_uniform: bool,
|
||||
):
|
||||
self._iterator = iterator
|
||||
self._dataset = dataset
|
||||
self._owner_token = owner_token
|
||||
self._require_uniform = require_uniform
|
||||
self._released = False
|
||||
self._terminal = False
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self._terminal:
|
||||
raise StopIteration
|
||||
try:
|
||||
batch = next(self._iterator)
|
||||
except StopIteration:
|
||||
self._terminal = True
|
||||
self._release()
|
||||
raise
|
||||
except BaseException as exc:
|
||||
self._dataset._invalidate_checkpoint(
|
||||
f"a DataLoader batch failed before it was returned: {exc}"
|
||||
)
|
||||
raise
|
||||
try:
|
||||
if not isinstance(batch, _WorkerBatch):
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader did not receive worker checkpoint metadata"
|
||||
)
|
||||
self._dataset._commit_worker_state(
|
||||
batch.state, require_uniform=self._require_uniform
|
||||
)
|
||||
return batch.data
|
||||
except BaseException as exc:
|
||||
self._dataset._invalidate_checkpoint(
|
||||
f"a DataLoader batch failed before it was returned: {exc}"
|
||||
)
|
||||
raise
|
||||
|
||||
def _release(self) -> None:
|
||||
if self.__dict__.get("_released", True):
|
||||
return
|
||||
self._released = True
|
||||
dataset = self.__dict__.get("_dataset")
|
||||
if dataset is not None:
|
||||
dataset._release_consumer_iterator(self._owner_token)
|
||||
|
||||
def _shutdown_workers(self):
|
||||
if self.__dict__.get("_released", True):
|
||||
return None
|
||||
self._terminal = True
|
||||
iterator = self.__dict__.get("_iterator")
|
||||
shutdown = getattr(iterator, "_shutdown_workers", None)
|
||||
try:
|
||||
if shutdown is not None:
|
||||
shutdown()
|
||||
else:
|
||||
fetcher = getattr(iterator, "_dataset_fetcher", None)
|
||||
dataset_iterator = getattr(fetcher, "dataset_iter", None)
|
||||
close = getattr(dataset_iterator, "close", None)
|
||||
if close is None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader could not close its inner iterator"
|
||||
)
|
||||
close()
|
||||
except BaseException as exc:
|
||||
self._dataset._invalidate_checkpoint(
|
||||
f"a DataLoader iterator could not be shut down safely: {exc}"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
self._release()
|
||||
|
||||
def __del__(self):
|
||||
try:
|
||||
self._shutdown_workers()
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(self._iterator, name)
|
||||
|
||||
|
||||
class StreamingDataset(IterableDataset):
|
||||
"""An elastic, resumable PyTorch IterableDataset backed by a LanceDB table.
|
||||
|
||||
@@ -384,6 +535,22 @@ class StreamingDataset(IterableDataset):
|
||||
# rows_skipped]
|
||||
self._worker_stats: RawArray = RawArray(ctypes.c_int64, 8)
|
||||
|
||||
# A standard multi-process DataLoader cannot report which prefetched
|
||||
# batches were actually returned to its consumer. Workers set this
|
||||
# shared flag so state_dict() can reject a stale parent checkpoint
|
||||
# unless StreamingDataLoader installed the consumer-commit transport.
|
||||
self._untracked_worker_iteration: RawArray = RawArray(ctypes.c_int64, 1)
|
||||
|
||||
# Parent-side checkpoint lifecycle. A failed DataLoader task creates
|
||||
# a permanent hole in that iterator's delivery stream, while a
|
||||
# multi-worker checkpoint is safe to restore only after all splits
|
||||
# reach the same logical step boundary.
|
||||
self._checkpoint_invalid_reason: Optional[str] = None
|
||||
self._consumer_checkpoint_requires_uniform = False
|
||||
self._consumer_iterator_lock = threading.Lock()
|
||||
self._consumer_iterator_generation = 0
|
||||
self._consumer_iterator_lease: Optional[_ConsumerIteratorLease] = None
|
||||
|
||||
# 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.
|
||||
@@ -396,6 +563,10 @@ class StreamingDataset(IterableDataset):
|
||||
# step boundaries all splits have consumed this many samples, so a
|
||||
# single scalar captures the topology-independent checkpoint state.
|
||||
self._resume_offset: int = 0
|
||||
# Exact yielded-sample counts for splits this process has advanced.
|
||||
# Missing entries use _resume_offset, which remains the lower-bound
|
||||
# checkpoint inherited from an earlier uniform/global state.
|
||||
self._resume_samples: dict[int, int] = {}
|
||||
# Permutation position each split has consumed through, keyed by
|
||||
# global split index. Equal to _resume_offset for every split unless
|
||||
# on_transform_error skipped rows, in which case skipped positions
|
||||
@@ -521,11 +692,45 @@ class StreamingDataset(IterableDataset):
|
||||
return self._rank_splits[start : start + splits_per_worker]
|
||||
|
||||
def __iter__(self) -> Iterator[dict[str, Any]]:
|
||||
return self._iter()
|
||||
|
||||
def _iter(
|
||||
self, *, consumer_checkpoint_transport: bool = False
|
||||
) -> Iterator[dict[str, Any]]:
|
||||
owner_token = None
|
||||
previous_lease = self._consumer_iterator_lease
|
||||
if consumer_checkpoint_transport:
|
||||
if not self._consumer_iterator_active:
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader worker transport requires an active "
|
||||
"parent iterator reservation"
|
||||
)
|
||||
else:
|
||||
try:
|
||||
owner_token = self._acquire_consumer_iterator()
|
||||
except BaseException:
|
||||
self._release_consumer_iterator_after_failed_acquire(previous_lease)
|
||||
raise
|
||||
try:
|
||||
yield from self._iter_owned(
|
||||
consumer_checkpoint_transport=consumer_checkpoint_transport
|
||||
)
|
||||
finally:
|
||||
if owner_token is not None:
|
||||
self._release_consumer_iterator(owner_token)
|
||||
|
||||
def _iter_owned(
|
||||
self, *, consumer_checkpoint_transport: bool
|
||||
) -> 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."
|
||||
)
|
||||
real_worker = get_worker_info() is not None
|
||||
if real_worker and not consumer_checkpoint_transport:
|
||||
self._untracked_worker_iteration[0] = 1
|
||||
|
||||
my_splits = self._resolve_my_splits()
|
||||
if not my_splits:
|
||||
return
|
||||
@@ -533,6 +738,7 @@ class StreamingDataset(IterableDataset):
|
||||
# Set identity transform on each Permutation so __getitems__ returns
|
||||
# the raw RecordBatch. Stage 2 applies the real transform.
|
||||
permutations: list[Permutation] = []
|
||||
initial_samples: list[int] = []
|
||||
initial_positions: list[int] = []
|
||||
for split_idx in my_splits:
|
||||
perm = Permutation.from_tables(
|
||||
@@ -541,21 +747,22 @@ class StreamingDataset(IterableDataset):
|
||||
if self._columns is not None:
|
||||
perm = perm.select_columns(self._columns)
|
||||
perm = perm.with_transform(Transforms.arrow2arrow)
|
||||
sample_count = self._resume_samples.get(split_idx, self._resume_offset)
|
||||
# Both modes resume from absolute permutation positions. Packing
|
||||
# stores them separately because it also checkpoints partial blocks.
|
||||
start_pos = (
|
||||
self._pack_consumed[split_idx]
|
||||
if self._pack_sequences is not None
|
||||
else self._resume_positions.get(split_idx, self._resume_offset)
|
||||
else self._resume_positions.get(split_idx, sample_count)
|
||||
)
|
||||
if start_pos > 0:
|
||||
perm = perm.with_skip(start_pos)
|
||||
initial_samples.append(sample_count)
|
||||
initial_positions.append(start_pos)
|
||||
permutations.append(perm)
|
||||
|
||||
n = len(permutations)
|
||||
split_sizes = [perm.num_rows for perm in permutations]
|
||||
initial_offset = self._resume_offset
|
||||
local_consumed = [0] * n
|
||||
# Permutation position each split has consumed through (absolute,
|
||||
# i.e. counted from the start of the unskipped split). Runs ahead of
|
||||
@@ -853,6 +1060,27 @@ class StreamingDataset(IterableDataset):
|
||||
for i in range(n):
|
||||
_fill_io(i)
|
||||
|
||||
def _yield_row(i: int):
|
||||
pos, row = cooked[i].popleft()
|
||||
# Surface any completed prefetched failure before the
|
||||
# current row becomes durable checkpoint progress.
|
||||
_advance(i)
|
||||
local_consumed[i] += 1
|
||||
pos_consumed[i] = pos + 1
|
||||
split_idx = my_splits[i]
|
||||
self._resume_samples[split_idx] = (
|
||||
initial_samples[i] + local_consumed[i]
|
||||
)
|
||||
self._resume_positions[split_idx] = pos_consumed[i]
|
||||
return row
|
||||
|
||||
def _update_progress_stats() -> None:
|
||||
if not real_worker:
|
||||
self._resume_offset = min(
|
||||
initial_samples[j] + local_consumed[j] for j in range(n)
|
||||
)
|
||||
_update_stats()
|
||||
|
||||
if self._pack_sequences is not None:
|
||||
first_count = pack_blocks_emitted[my_splits[0]]
|
||||
if any(
|
||||
@@ -878,12 +1106,38 @@ class StreamingDataset(IterableDataset):
|
||||
tokens.extend([pad_id] * (pack_len - len(tokens)))
|
||||
block = _emit_block(i)
|
||||
pack_blocks_emitted[my_splits[i]] += 1
|
||||
# Checkpoint state must advance before yielding so
|
||||
# StreamingDataLoader can attach the exact state to
|
||||
# the batch it transports to the parent process.
|
||||
_commit_pack_state()
|
||||
if i == n - 1:
|
||||
_commit_pack_state()
|
||||
_update_stats()
|
||||
yield block
|
||||
return
|
||||
|
||||
# A checkpoint taken between round-robin split turns has
|
||||
# non-uniform counts. Resume lagging splits first so the
|
||||
# exact canonical sequence continues without replaying
|
||||
# already-consumed rows.
|
||||
if len(set(initial_samples)) > 1:
|
||||
catch_up_to = max(initial_samples)
|
||||
pending = [
|
||||
(initial_samples[i], my_splits[i], i)
|
||||
for i in range(n)
|
||||
if initial_samples[i] < catch_up_to
|
||||
]
|
||||
heapq.heapify(pending)
|
||||
while pending:
|
||||
consumed, _, i = heapq.heappop(pending)
|
||||
_ensure_cooked(i)
|
||||
if not cooked[i]:
|
||||
return
|
||||
row = _yield_row(i)
|
||||
if consumed + 1 < catch_up_to:
|
||||
heapq.heappush(pending, (consumed + 1, my_splits[i], i))
|
||||
_update_progress_stats()
|
||||
yield row
|
||||
|
||||
while True:
|
||||
# A cycle only runs if every split can still produce a
|
||||
# row. Without skips all splits exhaust simultaneously
|
||||
@@ -904,20 +1158,14 @@ class StreamingDataset(IterableDataset):
|
||||
break
|
||||
|
||||
for i in range(n):
|
||||
pos, row = cooked[i].popleft()
|
||||
local_consumed[i] += 1
|
||||
pos_consumed[i] = pos + 1
|
||||
_advance(i)
|
||||
row = _yield_row(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]
|
||||
for j, split_idx in enumerate(my_splits):
|
||||
self._resume_positions[split_idx] = pos_consumed[j]
|
||||
_update_stats()
|
||||
_update_progress_stats()
|
||||
|
||||
yield row
|
||||
finally:
|
||||
@@ -1064,6 +1312,7 @@ class StreamingDataset(IterableDataset):
|
||||
"_local_consumed_ref",
|
||||
):
|
||||
state[key] = None
|
||||
state["_consumer_iterator_lock"] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
@@ -1074,19 +1323,31 @@ class StreamingDataset(IterableDataset):
|
||||
table_state = state.pop("_table")
|
||||
perm_name, perm_data = state.pop("_perm_table")
|
||||
self.__dict__.update(state)
|
||||
self._consumer_iterator_lock = threading.Lock()
|
||||
if self._connection_factory is not None:
|
||||
self._table = self._connection_factory(table_name)
|
||||
else:
|
||||
self._table = _table_from_pickle_state(table_state)
|
||||
if table_state["kind"] == "memory":
|
||||
# Rebuilt from Arrow, so the recorded pin cannot resolve on it.
|
||||
perm_data = _drop_base_version(perm_data)
|
||||
self._perm_table = _connect("memory://").create_table(perm_name, perm_data)
|
||||
|
||||
def state_dict(self) -> dict:
|
||||
"""Snapshot the dataset's consumption state.
|
||||
|
||||
When using DataLoader workers, construct a
|
||||
[StreamingDataLoader][lancedb.streaming.StreamingDataLoader]. It
|
||||
commits worker state only when a prefetched batch is returned to the
|
||||
trainer. A standard multi-process ``DataLoader`` cannot expose that
|
||||
boundary, so calling this method after one has started raises
|
||||
``RuntimeError`` instead of returning stale producer state.
|
||||
|
||||
In row mode, the returned dict is topology-independent at global step
|
||||
boundaries. ``positions_consumed_per_split`` records how far each
|
||||
split's permutation has advanced, which can differ from the sample
|
||||
count when ``on_transform_error`` skips rows. Combine state dicts from
|
||||
count when ``on_transform_error`` skips rows. ``StreamingDataLoader``
|
||||
combines worker state in its parent process. Combine state dicts from
|
||||
every rank with
|
||||
[merge_state_dicts][lancedb.streaming.StreamingDataset.merge_state_dicts]
|
||||
before resuming on a different topology.
|
||||
@@ -1095,6 +1356,43 @@ class StreamingDataset(IterableDataset):
|
||||
for every logical split. When packing is sharded, merge every rank
|
||||
state with ``merge_state_dicts`` before loading it.
|
||||
"""
|
||||
if self._untracked_worker_iteration[0] and get_worker_info() is None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataset cannot checkpoint a standard DataLoader with "
|
||||
"num_workers > 0 because prefetched worker progress is not "
|
||||
"consumer-committed. Use StreamingDataLoader instead."
|
||||
)
|
||||
if self._checkpoint_invalid_reason is not None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataset checkpointing is invalid because "
|
||||
f"{self._checkpoint_invalid_reason}. Load the last valid "
|
||||
"checkpoint into a fresh dataset before continuing."
|
||||
)
|
||||
state = self._checkpoint_snapshot()
|
||||
if self._pack_sequences is not None:
|
||||
rank_blocks = [
|
||||
state["blocks_emitted_per_split"][split] for split in self._rank_splits
|
||||
]
|
||||
if len(set(rank_blocks)) > 1:
|
||||
raise RuntimeError(
|
||||
"Packed StreamingDataset checkpointing is only safe at a "
|
||||
"complete logical step boundary, when every split assigned "
|
||||
"to this rank has emitted the same block count. Consume more "
|
||||
"batches before calling state_dict()."
|
||||
)
|
||||
elif self._consumer_checkpoint_requires_uniform:
|
||||
samples = state["samples_consumed_per_split"]
|
||||
rank_samples = [samples[split] for split in self._rank_splits]
|
||||
if len(set(rank_samples)) > 1:
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader checkpointing with multiple workers is "
|
||||
"only safe at a complete logical step boundary, when every "
|
||||
"split assigned to this rank has the same consumed-sample "
|
||||
"count. Consume more batches before calling state_dict()."
|
||||
)
|
||||
return state
|
||||
|
||||
def _checkpoint_snapshot(self) -> dict:
|
||||
if self._pack_sequences is not None:
|
||||
return {
|
||||
"shuffle_seed": self._shuffle_seed,
|
||||
@@ -1108,18 +1406,141 @@ class StreamingDataset(IterableDataset):
|
||||
"blocks_emitted_per_split": list(self._pack_blocks_emitted),
|
||||
"pack_buffers": deepcopy(self._pack_buffers),
|
||||
}
|
||||
samples = [
|
||||
self._resume_samples.get(split, self._resume_offset)
|
||||
for split in range(self._num_splits)
|
||||
]
|
||||
positions = [
|
||||
self._resume_positions.get(split, self._resume_offset)
|
||||
self._resume_positions.get(split, samples[split])
|
||||
for split in range(self._num_splits)
|
||||
]
|
||||
return {
|
||||
"shuffle_seed": self._shuffle_seed,
|
||||
"num_splits": self._num_splits,
|
||||
"epoch": self._epoch,
|
||||
"samples_consumed_per_split": [self._resume_offset] * self._num_splits,
|
||||
"samples_consumed_per_split": samples,
|
||||
"positions_consumed_per_split": positions,
|
||||
}
|
||||
|
||||
def _invalidate_checkpoint(self, reason: str) -> None:
|
||||
if self._checkpoint_invalid_reason is None:
|
||||
self._checkpoint_invalid_reason = reason
|
||||
|
||||
@property
|
||||
def _consumer_iterator_active(self) -> bool:
|
||||
return self._consumer_iterator_lease is not None
|
||||
|
||||
@property
|
||||
def _consumer_iterator_owner(self) -> Optional[int]:
|
||||
lease = self._consumer_iterator_lease
|
||||
return lease.owner_token if lease is not None else None
|
||||
|
||||
@property
|
||||
def _consumer_iterator_owner_thread(self) -> Optional[int]:
|
||||
lease = self._consumer_iterator_lease
|
||||
return lease.owner_thread if lease is not None else None
|
||||
|
||||
def _acquire_consumer_iterator(self) -> int:
|
||||
"""Reserve this parent dataset for one checkpoint-aware iterator."""
|
||||
with self._consumer_iterator_lock:
|
||||
if self._consumer_iterator_active or self._raw_batches_ref is not None:
|
||||
raise RuntimeError(
|
||||
"StreamingDataset does not support concurrent iteration. "
|
||||
"Only one active iterator per dataset instance is allowed."
|
||||
)
|
||||
owner_thread = threading.get_ident()
|
||||
owner_token = self._consumer_iterator_generation + 1
|
||||
lease = _ConsumerIteratorLease(owner_token, owner_thread)
|
||||
self._consumer_iterator_generation = owner_token
|
||||
self._consumer_iterator_lease = lease
|
||||
return owner_token
|
||||
|
||||
def _release_consumer_iterator(self, owner_token: int) -> None:
|
||||
with self._consumer_iterator_lock:
|
||||
lease = self._consumer_iterator_lease
|
||||
if lease is not None and lease.owner_token == owner_token:
|
||||
self._consumer_iterator_lease = None
|
||||
|
||||
def _release_consumer_iterator_after_failed_acquire(
|
||||
self, previous_lease: Optional[_ConsumerIteratorLease]
|
||||
) -> None:
|
||||
"""Clean up when an interrupted acquire set a lease but did not return it."""
|
||||
owner_thread = threading.current_thread().ident
|
||||
with self._consumer_iterator_lock:
|
||||
lease = self._consumer_iterator_lease
|
||||
if (
|
||||
lease is not None
|
||||
and lease is not previous_lease
|
||||
and lease.owner_thread == owner_thread
|
||||
):
|
||||
self._consumer_iterator_lease = None
|
||||
|
||||
def _commit_worker_state(self, state: dict, *, require_uniform: bool) -> None:
|
||||
"""Merge one trainer-consumed worker batch into parent state."""
|
||||
for key, expected in (
|
||||
("shuffle_seed", self._shuffle_seed),
|
||||
("num_splits", self._num_splits),
|
||||
("epoch", self._epoch),
|
||||
):
|
||||
if state.get(key) != expected:
|
||||
raise ValueError(
|
||||
f"{key} mismatch in worker checkpoint: "
|
||||
f"{state.get(key)} != {expected}"
|
||||
)
|
||||
packed = "pack_buffers" in state
|
||||
if packed != (self._pack_sequences is not None):
|
||||
raise ValueError("worker checkpoint mode does not match the dataset")
|
||||
if packed:
|
||||
for key in ("pack_sequences", "eos_id", "pad_id", "blocks_per_epoch"):
|
||||
expected = getattr(self, f"_{key}")
|
||||
if state.get(key) != expected:
|
||||
raise ValueError(
|
||||
f"{key} mismatch in worker checkpoint: "
|
||||
f"{state.get(key)} != {expected}"
|
||||
)
|
||||
samples = state["samples_consumed_per_split"]
|
||||
emitted = state["blocks_emitted_per_split"]
|
||||
if len(samples) != self._num_splits or len(emitted) != self._num_splits:
|
||||
raise ValueError(
|
||||
"packed worker checkpoint must contain one entry per split"
|
||||
)
|
||||
buffers = state["pack_buffers"]
|
||||
for split, (count, blocks) in enumerate(zip(samples, emitted)):
|
||||
incoming = (int(blocks), int(count))
|
||||
current = (
|
||||
self._pack_blocks_emitted[split],
|
||||
self._pack_consumed[split],
|
||||
)
|
||||
if incoming > current:
|
||||
self._pack_blocks_emitted[split] = incoming[0]
|
||||
self._pack_consumed[split] = incoming[1]
|
||||
buffer = buffers.get(split, buffers.get(str(split)))
|
||||
if buffer is None:
|
||||
self._pack_buffers.pop(split, None)
|
||||
else:
|
||||
self._pack_buffers[split] = {
|
||||
"tokens": list(buffer["tokens"]),
|
||||
"starts": list(buffer["starts"]),
|
||||
}
|
||||
self._consumer_checkpoint_requires_uniform |= require_uniform
|
||||
return
|
||||
|
||||
samples = state["samples_consumed_per_split"]
|
||||
positions = state.get("positions_consumed_per_split", samples)
|
||||
for split, count in enumerate(samples):
|
||||
current = self._resume_samples.get(split, self._resume_offset)
|
||||
self._resume_samples[split] = max(current, int(count))
|
||||
for split, position in enumerate(positions):
|
||||
current = self._resume_positions.get(
|
||||
split, self._resume_samples.get(split, self._resume_offset)
|
||||
)
|
||||
self._resume_positions[split] = max(current, int(position))
|
||||
self._resume_offset = min(
|
||||
self._resume_samples.get(split, self._resume_offset)
|
||||
for split in range(self._num_splits)
|
||||
)
|
||||
self._consumer_checkpoint_requires_uniform |= require_uniform
|
||||
|
||||
def load_state_dict(self, state: dict) -> None:
|
||||
"""Resume from a previously snapshotted state.
|
||||
|
||||
@@ -1139,6 +1560,7 @@ class StreamingDataset(IterableDataset):
|
||||
f"shuffle_seed mismatch: checkpoint has {state['shuffle_seed']}, "
|
||||
f"current dataset has {self._shuffle_seed}"
|
||||
)
|
||||
self._consumer_checkpoint_requires_uniform = False
|
||||
|
||||
if "pack_buffers" in state or self._pack_sequences is not None:
|
||||
for key in (
|
||||
@@ -1165,14 +1587,17 @@ class StreamingDataset(IterableDataset):
|
||||
return
|
||||
|
||||
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
|
||||
self._resume_offset = min(consumed) if consumed else 0
|
||||
self._resume_samples = {
|
||||
split: int(count) for split, count in enumerate(consumed)
|
||||
}
|
||||
else:
|
||||
self._resume_offset = int(consumed)
|
||||
self._resume_samples = {}
|
||||
# Older checkpoints predate positions_consumed_per_split; without
|
||||
# skipped rows positions equal sample counts, so falling back to
|
||||
# _resume_offset (the .get default in __iter__) is exact.
|
||||
# the per-split sample count (the .get default in __iter__) is exact.
|
||||
positions = state.get("positions_consumed_per_split")
|
||||
if positions is None:
|
||||
self._resume_positions = {}
|
||||
@@ -1185,10 +1610,11 @@ class StreamingDataset(IterableDataset):
|
||||
def merge_state_dicts(states: list[dict]) -> dict:
|
||||
"""Merge state dicts saved by different ranks into one exact state.
|
||||
|
||||
For row mode, the elementwise maximum of permutation positions recovers
|
||||
splits advanced by different ranks after transform failures. For packed
|
||||
mode, the state that emitted the most blocks for each logical split
|
||||
supplies that split's permutation position and partial token buffer. Packed
|
||||
In row mode, each rank records exact consumer-committed progress for
|
||||
its own splits and lower bounds for the rest, so elementwise maxima
|
||||
recover both sample counts and permutation positions. In packed mode,
|
||||
the state that emitted the most blocks for each logical split supplies
|
||||
that split's permutation position and partial token buffer. Packed
|
||||
states must cover every rank at the same global step.
|
||||
|
||||
Raises ``ValueError`` if the states are empty, were not produced by
|
||||
@@ -1299,17 +1725,13 @@ class StreamingDataset(IterableDataset):
|
||||
merged["pack_buffers"] = merged_buffers
|
||||
return merged
|
||||
|
||||
for state in states[1:]:
|
||||
if (
|
||||
state["samples_consumed_per_split"]
|
||||
!= first["samples_consumed_per_split"]
|
||||
):
|
||||
raise ValueError(
|
||||
"samples_consumed_per_split mismatch across state dicts; "
|
||||
"state_dict() must be called at the same global step "
|
||||
"boundary on every rank"
|
||||
)
|
||||
merged = dict(first)
|
||||
merged["samples_consumed_per_split"] = [
|
||||
max(per_split)
|
||||
for per_split in zip(
|
||||
*(state["samples_consumed_per_split"] for state in states)
|
||||
)
|
||||
]
|
||||
all_positions = [
|
||||
state.get(
|
||||
"positions_consumed_per_split", state["samples_consumed_per_split"]
|
||||
@@ -1320,3 +1742,113 @@ class StreamingDataset(IterableDataset):
|
||||
max(per_split) for per_split in zip(*all_positions)
|
||||
]
|
||||
return merged
|
||||
|
||||
|
||||
class StreamingDataLoader(DataLoader):
|
||||
"""A PyTorch DataLoader with consumer-committed dataset checkpoints.
|
||||
|
||||
PyTorch workers prefetch batches ahead of the trainer, so worker-local
|
||||
producer progress is not a safe checkpoint. This loader carries a state
|
||||
snapshot alongside every internal batch and applies it to the parent
|
||||
[StreamingDataset][lancedb.streaming.StreamingDataset] only when that batch
|
||||
is returned by ``next()``.
|
||||
The trainer receives the same collated batch it would receive from a
|
||||
standard ``torch.utils.data.DataLoader``.
|
||||
|
||||
With more than one worker, row-mode ``state_dict()`` is available only at
|
||||
complete logical step boundaries, when every split assigned to the rank has
|
||||
the same consumed-sample count. Packed checkpoints require equal emitted-block
|
||||
counts across the rank's splits for any worker count. ``persistent_workers=True``
|
||||
is not supported because prefetched worker copies cannot be restored from
|
||||
parent-committed state. If batch collation raises, checkpointing remains
|
||||
invalid for that dataset instance; restore the last valid checkpoint into a
|
||||
fresh dataset before continuing.
|
||||
Only one active iterator may own a dataset at a time, including when worker
|
||||
processes are used. Exhausting or explicitly shutting down the iterator
|
||||
releases that ownership. ``drop_last=True`` is not supported because worker
|
||||
replicas discard incomplete tails independently, which cannot produce a
|
||||
topology-independent checkpoint.
|
||||
|
||||
Parameters are the same as ``torch.utils.data.DataLoader`` except that
|
||||
``dataset`` must be a
|
||||
[StreamingDataset][lancedb.streaming.StreamingDataset].
|
||||
Subclasses that override ``StreamingDataset.__iter__`` are not supported
|
||||
because the custom iterator cannot provide the exact per-yield checkpoint
|
||||
snapshots required by this loader.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> # dataset = StreamingDataset(table, num_splits=2)
|
||||
>>> # loader = StreamingDataLoader(dataset, batch_size=8, num_workers=2)
|
||||
>>> # batch = next(iter(loader))
|
||||
>>> # checkpoint = dataset.state_dict()
|
||||
"""
|
||||
|
||||
def __init__(self, dataset: StreamingDataset, *args, **kwargs):
|
||||
if not isinstance(dataset, StreamingDataset):
|
||||
raise TypeError("StreamingDataLoader requires a StreamingDataset")
|
||||
if type(dataset).__iter__ is not StreamingDataset.__iter__:
|
||||
raise TypeError(
|
||||
"StreamingDataLoader does not support StreamingDataset subclasses "
|
||||
"that override __iter__ because they cannot provide exact "
|
||||
"per-yield checkpoint state"
|
||||
)
|
||||
if kwargs.get("in_order", True) is False:
|
||||
raise ValueError(
|
||||
"StreamingDataLoader requires in_order=True for deterministic "
|
||||
"consumer checkpoints"
|
||||
)
|
||||
if kwargs.get("persistent_workers", False):
|
||||
raise ValueError(
|
||||
"StreamingDataLoader does not support persistent_workers=True "
|
||||
"because worker prefetch state cannot be reset from a checkpoint"
|
||||
)
|
||||
self._streaming_dataset = dataset
|
||||
super().__init__(_StreamingDatasetAdapter(dataset), *args, **kwargs)
|
||||
if self.drop_last:
|
||||
raise ValueError(
|
||||
"StreamingDataLoader does not support drop_last=True because "
|
||||
"discarded worker tails cannot be checkpointed "
|
||||
"topology-independently"
|
||||
)
|
||||
self.collate_fn = _CheckpointCollate(self.collate_fn)
|
||||
|
||||
def __iter__(self):
|
||||
dataset = self._streaming_dataset
|
||||
previous_lease = dataset._consumer_iterator_lease
|
||||
owner_token = None
|
||||
try:
|
||||
owner_token = dataset._acquire_consumer_iterator()
|
||||
state = dataset._checkpoint_snapshot()
|
||||
packed = dataset._pack_sequences is not None
|
||||
if packed:
|
||||
blocks = state["blocks_emitted_per_split"]
|
||||
rank_blocks = [blocks[split] for split in dataset._rank_splits]
|
||||
if len(set(rank_blocks)) > 1:
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader cannot start from a partial packed "
|
||||
"logical step; resume from a checkpoint whose splits "
|
||||
"assigned to this rank have equal emitted-block counts"
|
||||
)
|
||||
elif self.num_workers > 1:
|
||||
samples = state["samples_consumed_per_split"]
|
||||
rank_samples = [samples[split] for split in dataset._rank_splits]
|
||||
if len(set(rank_samples)) > 1:
|
||||
raise RuntimeError(
|
||||
"StreamingDataLoader cannot start multiple workers from a "
|
||||
"partial logical step; resume from a checkpoint whose "
|
||||
"splits assigned to this rank have equal consumed-sample "
|
||||
"counts"
|
||||
)
|
||||
return _ConsumerCommitIterator(
|
||||
super().__iter__(),
|
||||
dataset,
|
||||
owner_token=owner_token,
|
||||
require_uniform=self.num_workers > 1 or packed,
|
||||
)
|
||||
except BaseException:
|
||||
if owner_token is not None:
|
||||
dataset._release_consumer_iterator(owner_token)
|
||||
else:
|
||||
dataset._release_consumer_iterator_after_failed_acquire(previous_lease)
|
||||
raise
|
||||
|
||||
+314
-83
@@ -85,6 +85,7 @@ from .query import (
|
||||
AsyncQuery,
|
||||
AsyncTakeQuery,
|
||||
AsyncVectorQuery,
|
||||
DocumentGranularity,
|
||||
FullTextQuery,
|
||||
LanceEmptyQueryBuilder,
|
||||
LanceFtsQueryBuilder,
|
||||
@@ -103,7 +104,12 @@ from .util import (
|
||||
value_to_sql,
|
||||
)
|
||||
from .index import lang_mapping
|
||||
from .schema import blob_v2_column_paths, schema_has_blob_field
|
||||
from .schema import (
|
||||
blob_v2_column_paths,
|
||||
is_blob_v2_field,
|
||||
row_addressable_blob_v2_paths,
|
||||
schema_has_blob_field,
|
||||
)
|
||||
|
||||
|
||||
def _should_push_down_query_table(
|
||||
@@ -425,6 +431,7 @@ def _cast_to_target_schema(
|
||||
|
||||
def gen():
|
||||
for batch in reader:
|
||||
batch = _coerce_blob_write_columns(batch, reordered_schema)
|
||||
# Table but not RecordBatch has cast.
|
||||
cast_batches = (
|
||||
pa.Table.from_batches([batch]).cast(reordered_schema).to_batches()
|
||||
@@ -437,6 +444,166 @@ def _cast_to_target_schema(
|
||||
return pa.RecordBatchReader.from_batches(reordered_schema, gen())
|
||||
|
||||
|
||||
def _coerce_blob_write_columns(
|
||||
batch: pa.RecordBatch, target_schema: pa.Schema
|
||||
) -> pa.RecordBatch:
|
||||
"""Materialize blob storage structs before the stream leaves Python.
|
||||
|
||||
merge_insert requires its source reader to already match the table's
|
||||
physical schema. Unlike add and insert, it does not pass through
|
||||
LanceDB's Rust blob coercion, so preserving binary input here would
|
||||
reach Lance as binary and fail the schema check.
|
||||
"""
|
||||
columns = []
|
||||
fields = []
|
||||
changed = False
|
||||
for field, column in zip(batch.schema, batch.columns):
|
||||
target_field = target_schema.field(field.name)
|
||||
coerced = _coerce_blob_value(column, target_field)
|
||||
if coerced is not column:
|
||||
column = coerced
|
||||
field = pa.field(
|
||||
field.name,
|
||||
coerced.type,
|
||||
field.nullable,
|
||||
target_field.metadata,
|
||||
)
|
||||
changed = True
|
||||
columns.append(column)
|
||||
fields.append(field)
|
||||
if not changed:
|
||||
return batch
|
||||
return pa.RecordBatch.from_arrays(
|
||||
columns, schema=pa.schema(fields, metadata=batch.schema.metadata)
|
||||
)
|
||||
|
||||
|
||||
def _coerce_blob_value(column: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
if is_blob_v2_field(target_field) and _can_coerce_to_blob(column.type):
|
||||
return _coerce_value_to_blob(column, target_field)
|
||||
|
||||
target_type = target_field.type
|
||||
if pa.types.is_struct(target_type) and pa.types.is_struct(column.type):
|
||||
children = []
|
||||
fields = []
|
||||
changed = False
|
||||
for source_field in column.type:
|
||||
source_column = column.field(source_field.name)
|
||||
nested_target = next(
|
||||
(field for field in target_type if field.name == source_field.name),
|
||||
None,
|
||||
)
|
||||
if nested_target is None:
|
||||
children.append(source_column)
|
||||
fields.append(source_field)
|
||||
continue
|
||||
coerced = _coerce_blob_value(source_column, nested_target)
|
||||
if coerced is not source_column:
|
||||
changed = True
|
||||
child_array, child_type = _physical_array_and_type(coerced)
|
||||
children.append(child_array)
|
||||
fields.append(
|
||||
pa.field(
|
||||
source_field.name,
|
||||
child_type,
|
||||
source_field.nullable,
|
||||
nested_target.metadata,
|
||||
)
|
||||
)
|
||||
if not changed:
|
||||
return column
|
||||
return pa.StructArray.from_arrays(
|
||||
children,
|
||||
fields=fields,
|
||||
mask=column.is_null() if column.null_count else None,
|
||||
)
|
||||
|
||||
if _is_list_like(target_type) and _is_list_like(column.type):
|
||||
return _coerce_blob_list_values(column, target_type.value_field)
|
||||
|
||||
return column
|
||||
|
||||
|
||||
def _coerce_blob_list_values(
|
||||
column: pa.Array, target_value_field: pa.Field
|
||||
) -> pa.Array:
|
||||
"""Coerce blob values inside a list column, preserving offsets and nulls.
|
||||
|
||||
Works on the raw child values window instead of ``pc.list_flatten`` because
|
||||
flatten drops values spanned by null slots, which would misalign offsets.
|
||||
"""
|
||||
mask = column.is_null() if column.null_count else None
|
||||
if pa.types.is_fixed_size_list(column.type):
|
||||
list_size = column.type.list_size
|
||||
values = column.values.slice(column.offset * list_size, len(column) * list_size)
|
||||
coerced = _coerce_blob_value(values, target_value_field)
|
||||
if coerced is values:
|
||||
return column
|
||||
physical_values, _ = _physical_array_and_type(coerced)
|
||||
return pa.FixedSizeListArray.from_arrays(physical_values, list_size, mask=mask)
|
||||
offsets = column.offsets
|
||||
first_offset = offsets[0].as_py()
|
||||
values = column.values.slice(
|
||||
first_offset,
|
||||
offsets[-1].as_py() - first_offset,
|
||||
)
|
||||
coerced = _coerce_blob_value(values, target_value_field)
|
||||
if coerced is values:
|
||||
return column
|
||||
physical_values, _ = _physical_array_and_type(coerced)
|
||||
if first_offset:
|
||||
offsets = pc.subtract(offsets, pa.scalar(first_offset, offsets.type))
|
||||
if pa.types.is_large_list(column.type):
|
||||
return pa.LargeListArray.from_arrays(offsets, physical_values, mask=mask)
|
||||
return pa.ListArray.from_arrays(offsets, physical_values, mask=mask)
|
||||
|
||||
|
||||
def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array:
|
||||
if pa.types.is_null(values.type):
|
||||
data = pa.nulls(len(values), type=pa.large_binary())
|
||||
elif pa.types.is_large_binary(values.type):
|
||||
data = values
|
||||
else:
|
||||
data = values.cast(pa.large_binary())
|
||||
length = len(values)
|
||||
storage_type = target_field.type
|
||||
if isinstance(storage_type, pa.ExtensionType):
|
||||
storage_type = storage_type.storage_type
|
||||
storage_fields = list(storage_type)
|
||||
children = []
|
||||
for storage_field in storage_fields:
|
||||
if storage_field.name == "data":
|
||||
children.append(data)
|
||||
else:
|
||||
children.append(pa.nulls(length, type=storage_field.type))
|
||||
storage = pa.StructArray.from_arrays(
|
||||
children,
|
||||
fields=storage_fields,
|
||||
mask=values.is_null() if values.null_count else None,
|
||||
)
|
||||
if isinstance(target_field.type, pa.ExtensionType):
|
||||
return pa.ExtensionArray.from_storage(target_field.type, storage)
|
||||
return storage
|
||||
|
||||
|
||||
def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]:
|
||||
if isinstance(array.type, pa.ExtensionType):
|
||||
return array.storage, array.type.storage_type
|
||||
return array, array.type
|
||||
|
||||
|
||||
def _can_coerce_to_blob(data_type: pa.DataType) -> bool:
|
||||
return _is_binary_like(data_type) or pa.types.is_null(data_type)
|
||||
|
||||
|
||||
def _is_binary_like(data_type: pa.DataType) -> bool:
|
||||
return (
|
||||
pa.types.is_binary(data_type)
|
||||
or pa.types.is_large_binary(data_type)
|
||||
or pa.types.is_binary_view(data_type)
|
||||
)
|
||||
|
||||
|
||||
def _field_extension_name(field: pa.Field) -> Optional[str]:
|
||||
extension_name = getattr(field.type, "extension_name", None)
|
||||
if extension_name is not None:
|
||||
@@ -463,63 +630,71 @@ def _align_field_types(
|
||||
target_field = next((f for f in target_fields if f.name == field.name), None)
|
||||
if target_field is None:
|
||||
raise ValueError(f"Field '{field.name}' not found in target schema")
|
||||
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
|
||||
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
|
||||
# input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
if (
|
||||
_field_extension_name(field) == "arrow.json"
|
||||
and _field_extension_name(target_field) == "lance.json"
|
||||
):
|
||||
new_fields.append(field)
|
||||
continue
|
||||
if pa.types.is_struct(target_field.type):
|
||||
if pa.types.is_struct(field.type):
|
||||
new_type = pa.struct(
|
||||
_align_field_types(
|
||||
field.type.fields,
|
||||
target_field.type.fields,
|
||||
)
|
||||
new_fields.append(_align_field(field, target_field))
|
||||
return new_fields
|
||||
|
||||
|
||||
def _align_list_value_field(
|
||||
value_field: pa.Field, target_value_field: pa.Field
|
||||
) -> pa.Field:
|
||||
# A list has exactly one child, so the inferred child name ("item") aligns
|
||||
# positionally and adopts the table's child name; pa.Table.cast renames it.
|
||||
return _align_field(value_field, target_value_field).with_name(
|
||||
target_value_field.name
|
||||
)
|
||||
|
||||
|
||||
def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field:
|
||||
# Preserve arrow.json input until it reaches Lance. LanceDB exposes stored
|
||||
# JSON columns as lance.json (JSONB-backed LargeBinary), but casting the
|
||||
# input to that storage type here merely relabels the raw JSON bytes as
|
||||
# JSONB. Lance must see arrow.json so it can perform the JSONB encoding.
|
||||
if (
|
||||
_field_extension_name(field) == "arrow.json"
|
||||
and _field_extension_name(target_field) == "lance.json"
|
||||
):
|
||||
return field
|
||||
if pa.types.is_struct(target_field.type):
|
||||
if pa.types.is_struct(field.type):
|
||||
new_type = pa.struct(
|
||||
_align_field_types(
|
||||
field.type.fields,
|
||||
target_field.type.fields,
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_field_types(
|
||||
[field.type.value_field],
|
||||
[target_field.type.value_field],
|
||||
)[0]
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_large_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.large_list(
|
||||
_align_field_types(
|
||||
[field.type.value_field],
|
||||
[target_field.type.value_field],
|
||||
)[0]
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_fixed_size_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_field_types(
|
||||
[field.type.value_field],
|
||||
[target_field.type.value_field],
|
||||
)[0],
|
||||
target_field.type.list_size,
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
new_fields.append(
|
||||
pa.field(field.name, new_type, field.nullable, target_field.metadata)
|
||||
)
|
||||
return new_fields
|
||||
elif pa.types.is_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_list_value_field(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
)
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_large_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.large_list(
|
||||
_align_list_value_field(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
)
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
elif pa.types.is_fixed_size_list(target_field.type):
|
||||
if _is_list_like(field.type):
|
||||
new_type = pa.list_(
|
||||
_align_list_value_field(
|
||||
field.type.value_field, target_field.type.value_field
|
||||
),
|
||||
target_field.type.list_size,
|
||||
)
|
||||
else:
|
||||
new_type = target_field.type
|
||||
else:
|
||||
new_type = target_field.type
|
||||
return pa.field(field.name, new_type, field.nullable, target_field.metadata)
|
||||
|
||||
|
||||
def _infer_subschema(
|
||||
@@ -588,7 +763,7 @@ def sanitize_create_table(
|
||||
schema = data.schema
|
||||
else:
|
||||
if schema is not None:
|
||||
data = pa.Table.from_pylist([], schema)
|
||||
data = pa.Table.from_batches([], schema=schema)
|
||||
if schema is None:
|
||||
if data is None:
|
||||
raise ValueError("Either data or schema must be provided")
|
||||
@@ -1168,6 +1343,7 @@ class Table(ABC):
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
block_size: int = 128,
|
||||
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
|
||||
wait_timeout: Optional[timedelta] = None,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
@@ -1246,6 +1422,11 @@ class Table(ABC):
|
||||
The number of documents per compressed posting block. Must be 128
|
||||
or 256. A value of 256 uses the experimental FTS V3 format and
|
||||
may introduce breaking changes.
|
||||
document_granularity: DocumentGranularity, default ROW
|
||||
``ROW`` treats the selected text in one table row as one document.
|
||||
``LIST_ELEMENT`` treats each element of the deepest list on the field
|
||||
path as one document and returns its physical coordinates in
|
||||
``_doc_index`` for matching queries.
|
||||
wait_timeout: timedelta, optional
|
||||
The timeout to wait if indexing is asynchronous.
|
||||
name: str, optional
|
||||
@@ -1269,6 +1450,7 @@ class Table(ABC):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [Table][lancedb.table.Table].
|
||||
|
||||
@@ -1320,6 +1502,10 @@ class Table(ABC):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Store blob URIs that sit outside registered blob bases. The row
|
||||
keeps a reference, so the object has to stay readable. Local
|
||||
tables only.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -1732,7 +1918,7 @@ class Table(ABC):
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
where: Optional[str] = None,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -1747,9 +1933,11 @@ class Table(ABC):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
@@ -1767,6 +1955,7 @@ class Table(ABC):
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
|
||||
>>> db = lancedb.connect("./.lancedb")
|
||||
@@ -1776,7 +1965,7 @@ class Table(ABC):
|
||||
0 1 [1.0, 2.0]
|
||||
1 2 [3.0, 4.0]
|
||||
2 3 [5.0, 6.0]
|
||||
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
|
||||
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
|
||||
UpdateResult(rows_updated=1, version=2)
|
||||
>>> table.to_pandas()
|
||||
x vector
|
||||
@@ -1972,13 +2161,15 @@ class Table(ABC):
|
||||
A mapping with one ``FunctionApplication`` value keeps its scalar
|
||||
or named-struct result in the named table column. A bare
|
||||
named-struct application expands its ordered result fields as one
|
||||
atomic sibling group; aliases come from ``rename(columns=...)``.
|
||||
atomic binding; aliases come from ``rename(columns=...)``.
|
||||
Function columns are supported only on LanceDB Cloud and
|
||||
Enterprise.
|
||||
computed: Dict[str, str], optional
|
||||
A map of column name to a SQL expression defining the column. The
|
||||
column's type and inputs are derived from the expression, so no
|
||||
data type is supplied.
|
||||
A mapping from output column names to SQL expressions derives each
|
||||
output field from its expression. A direct projection of a Blob v2
|
||||
field inherits Blob v2 semantics; other expressions derive their
|
||||
ordinary Arrow type. Mapping order is declaration and dependency
|
||||
order.
|
||||
|
||||
Unlike ``transforms``, the expression is stored rather than
|
||||
evaluated now: the column is committed with no values, and rows get
|
||||
@@ -2115,12 +2306,25 @@ class Table(ABC):
|
||||
----------
|
||||
updates : dict
|
||||
One or more dicts, each with:
|
||||
|
||||
- "path": str — dot-path to the field (e.g. "embedding" or "a.b.c").
|
||||
- "metadata": dict[str, str | None] — keys to set; a value of ``None``
|
||||
deletes that key.
|
||||
- "replace": bool, optional — replace the field's whole metadata map
|
||||
instead of merging (default False).
|
||||
|
||||
The following keys are treated specially, by convention, and should
|
||||
be used when appropriate:
|
||||
|
||||
- "lancedb:description": for a human-readable description of a field.
|
||||
- ``"lancedb:tag:<name>"`` for a user-defined key-value tag, where the
|
||||
suffix names the tag category; e.g. "lancedb:tag:model": "clip".
|
||||
- "lancedb:logical-column" for a column grouping; e.g. "feature_v1"
|
||||
and "feature_v2" might be in the same logical column.
|
||||
- "lancedb:status" for status options ("production", "candidate",
|
||||
"deprecated", "archived") to designate the current life cycle
|
||||
state of this column.
|
||||
|
||||
Returns
|
||||
-------
|
||||
UpdateFieldMetadataResult
|
||||
@@ -2670,7 +2874,7 @@ class LanceTable(Table):
|
||||
arrow_tbl = self.to_arrow()
|
||||
if blob_mode == "descriptions":
|
||||
arrow_tbl = strip_auto_row_ids(
|
||||
arrow_tbl, blob_v2_column_paths(self.schema)
|
||||
arrow_tbl, row_addressable_blob_v2_paths(self.schema)
|
||||
)
|
||||
return arrow_tbl.to_pandas(**kwargs)
|
||||
|
||||
@@ -3268,6 +3472,7 @@ class LanceTable(Table):
|
||||
ngram_max_length: int = 3,
|
||||
prefix_only: bool = False,
|
||||
block_size: int = 128,
|
||||
document_granularity: DocumentGranularity = DocumentGranularity.ROW,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
"""Create a full-text search index on a column.
|
||||
@@ -3319,7 +3524,11 @@ class LanceTable(Table):
|
||||
tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name)
|
||||
tokenizer_configs["custom_stop_words"] = custom_stop_words
|
||||
|
||||
config = FTS(block_size=block_size, **tokenizer_configs)
|
||||
config = FTS(
|
||||
block_size=block_size,
|
||||
document_granularity=document_granularity,
|
||||
**tokenizer_configs,
|
||||
)
|
||||
|
||||
try:
|
||||
LOOP.run(
|
||||
@@ -3409,6 +3618,7 @@ class LanceTable(Table):
|
||||
fill_value: float = 0.0,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add data to the table.
|
||||
If vector columns are missing and the table
|
||||
@@ -3436,6 +3646,9 @@ class LanceTable(Table):
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
Local tables only.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -3452,6 +3665,7 @@ class LanceTable(Table):
|
||||
fill_value=fill_value,
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
@@ -3806,7 +4020,7 @@ class LanceTable(Table):
|
||||
|
||||
def update(
|
||||
self,
|
||||
where: Optional[str] = None,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
values: Optional[dict] = None,
|
||||
*,
|
||||
values_sql: Optional[Dict[str, str]] = None,
|
||||
@@ -3817,9 +4031,11 @@ class LanceTable(Table):
|
||||
|
||||
Parameters
|
||||
----------
|
||||
where: str, optional
|
||||
The SQL where clause to use when updating rows. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error.
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. The filter must not be empty, or it will
|
||||
error.
|
||||
values: dict, optional
|
||||
The values to update. The keys are the column names and the values
|
||||
are the values to set.
|
||||
@@ -3837,6 +4053,7 @@ class LanceTable(Table):
|
||||
Examples
|
||||
--------
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]})
|
||||
>>> db = lancedb.connect("./.lancedb")
|
||||
@@ -3846,7 +4063,7 @@ class LanceTable(Table):
|
||||
0 1 [1.0, 2.0]
|
||||
1 2 [3.0, 4.0]
|
||||
2 3 [5.0, 6.0]
|
||||
>>> table.update(where="x = 2", values={"vector": [10.0, 10]})
|
||||
>>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]})
|
||||
UpdateResult(rows_updated=1, version=2)
|
||||
>>> table.to_pandas()
|
||||
x vector
|
||||
@@ -5061,7 +5278,9 @@ class AsyncTable:
|
||||
if blob_mode == "descriptions" or not schema_has_blob_field(schema):
|
||||
arrow_tbl = await self.to_arrow()
|
||||
if blob_mode == "descriptions":
|
||||
arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema))
|
||||
arrow_tbl = strip_auto_row_ids(
|
||||
arrow_tbl, row_addressable_blob_v2_paths(schema)
|
||||
)
|
||||
return arrow_tbl.to_pandas(**kwargs)
|
||||
|
||||
if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory":
|
||||
@@ -5365,6 +5584,7 @@ class AsyncTable:
|
||||
fill_value: Optional[float] = None,
|
||||
progress: Optional[Union[bool, Callable, Any]] = None,
|
||||
write_parallelism: Optional[int] = None,
|
||||
allow_external_blob_outside_bases: bool = False,
|
||||
) -> AddResult:
|
||||
"""Add more data to the [AsyncTable][lancedb.table.AsyncTable].
|
||||
|
||||
@@ -5395,6 +5615,9 @@ class AsyncTable:
|
||||
data in flight. Defaults to an estimate based on the data size,
|
||||
capped at the number of CPU cores. Lower this if bulk ingestion is
|
||||
using too much memory.
|
||||
allow_external_blob_outside_bases: bool, default False
|
||||
Allow blob URIs outside registered bases. See :meth:`Table.add`.
|
||||
Local tables only.
|
||||
|
||||
"""
|
||||
schema = await self.schema()
|
||||
@@ -5431,6 +5654,7 @@ class AsyncTable:
|
||||
mode or "append",
|
||||
progress=progress,
|
||||
write_parallelism=write_parallelism,
|
||||
allow_external_blob_outside_bases=allow_external_blob_outside_bases,
|
||||
)
|
||||
except RuntimeError as e:
|
||||
if "Cast error" in str(e):
|
||||
@@ -5955,7 +6179,7 @@ class AsyncTable:
|
||||
self,
|
||||
updates: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
where: Optional[str] = None,
|
||||
where: Optional[Union[str, Expr]] = None,
|
||||
updates_sql: Optional[Dict[str, str]] = None,
|
||||
) -> UpdateResult:
|
||||
"""
|
||||
@@ -5970,9 +6194,11 @@ class AsyncTable:
|
||||
The updates to apply. The keys should be the name of the column to
|
||||
update. The values should be the new values to assign. This is
|
||||
required unless updates_sql is supplied.
|
||||
where: str, optional
|
||||
An SQL filter that controls which rows are updated. For example, 'x = 2'
|
||||
or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated.
|
||||
where: str or [Expr][lancedb.expr.Expr], optional
|
||||
The filter condition. Can be a SQL string or a type-safe
|
||||
[Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and
|
||||
[lit][lancedb.expr.lit]. Only rows that satisfy this filter will
|
||||
be updated.
|
||||
updates_sql: dict, optional
|
||||
The updates to apply, expressed as SQL expression strings. The keys should
|
||||
be column names. The values should be SQL expressions. These can be SQL
|
||||
@@ -5990,13 +6216,14 @@ class AsyncTable:
|
||||
--------
|
||||
>>> import asyncio
|
||||
>>> import lancedb
|
||||
>>> from lancedb.expr import col
|
||||
>>> import pandas as pd
|
||||
>>> async def demo_update():
|
||||
... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]})
|
||||
... db = await lancedb.connect_async("./.lancedb")
|
||||
... table = await db.create_table("my_table", data)
|
||||
... # x is [1, 2], vector is [[1, 2], [3, 4]]
|
||||
... await table.update({"vector": [10, 10]}, where="x = 2")
|
||||
... await table.update({"vector": [10, 10]}, where=col("x") == 2)
|
||||
... # x is [1, 2], vector is [[1, 2], [10, 10]]
|
||||
... await table.update(updates_sql={"x": "x + 1"})
|
||||
... # x is [2, 3], vector is [[1, 2], [10, 10]]
|
||||
@@ -6010,7 +6237,8 @@ class AsyncTable:
|
||||
if updates is not None:
|
||||
updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
|
||||
|
||||
return await self._inner.update(updates_sql, where)
|
||||
predicate = where.to_sql() if isinstance(where, Expr) else where
|
||||
return await self._inner.update(updates_sql, predicate)
|
||||
|
||||
async def add_columns(
|
||||
self,
|
||||
@@ -6038,12 +6266,15 @@ class AsyncTable:
|
||||
A mapping with one ``FunctionApplication`` value keeps its scalar
|
||||
or named-struct result in the named table column. A bare
|
||||
named-struct application expands its ordered result fields as one
|
||||
atomic sibling group; aliases come from ``rename(columns=...)``.
|
||||
atomic binding; aliases come from ``rename(columns=...)``.
|
||||
Function columns are supported only on LanceDB Cloud and
|
||||
Enterprise.
|
||||
computed: Dict[str, str], optional
|
||||
A map of column name to a SQL expression defining the column. The
|
||||
column's type and inputs are derived from the expression.
|
||||
A mapping from output column names to SQL expressions derives each
|
||||
output field from its expression. A direct projection of a Blob v2
|
||||
field inherits Blob v2 semantics; other expressions derive their
|
||||
ordinary Arrow type. Mapping order is declaration and dependency
|
||||
order.
|
||||
|
||||
Unlike ``transforms``, the expression is stored rather than
|
||||
evaluated now: the column is committed with no values, and rows get
|
||||
@@ -6075,7 +6306,7 @@ class AsyncTable:
|
||||
isinstance(value, FunctionApplication) for value in transforms.values()
|
||||
):
|
||||
raise ValueError(
|
||||
"one add_columns call declares exactly one Function sibling group"
|
||||
"one add_columns call declares exactly one Function binding"
|
||||
)
|
||||
function_output_name, function_application = next(iter(transforms.items()))
|
||||
|
||||
|
||||
@@ -105,7 +105,7 @@ def test_quickstart(tmp_path):
|
||||
tbl.create_index(num_sub_vectors=1)
|
||||
# --8<-- [end:create_index]
|
||||
# --8<-- [start:delete_rows]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows]
|
||||
# --8<-- [start:drop_table]
|
||||
db.drop_table("my_table")
|
||||
@@ -201,7 +201,7 @@ async def test_quickstart_async(tmp_path):
|
||||
await tbl.create_index("vector")
|
||||
# --8<-- [end:create_index_async]
|
||||
# --8<-- [start:delete_rows_async]
|
||||
await tbl.delete('item = "fizz"')
|
||||
await tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_rows_async]
|
||||
# --8<-- [start:drop_table_async]
|
||||
await db.drop_table("my_table_async")
|
||||
|
||||
@@ -266,7 +266,7 @@ def test_table():
|
||||
tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_from_pydantic]
|
||||
# --8<-- [start:delete_row]
|
||||
tbl.delete('item = "fizz"')
|
||||
tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row]
|
||||
# --8<-- [start:delete_specific_row]
|
||||
data = [
|
||||
@@ -538,7 +538,7 @@ async def test_table_async():
|
||||
await async_tbl.add(pydantic_model_items)
|
||||
# --8<-- [end:add_table_async_from_pydantic]
|
||||
# --8<-- [start:delete_row_async]
|
||||
await async_tbl.delete('item = "fizz"')
|
||||
await async_tbl.delete("item = 'fizz'")
|
||||
# --8<-- [end:delete_row_async]
|
||||
# --8<-- [start:delete_specific_row_async]
|
||||
data = [
|
||||
|
||||
@@ -2,17 +2,41 @@
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import io
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
import lance
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
import pytest
|
||||
from lance.blob import BlobType as LanceBlobType
|
||||
|
||||
import lancedb
|
||||
from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids
|
||||
from lancedb._blob import (
|
||||
blob_v2_projection_sources,
|
||||
read_row_ids_from_hits,
|
||||
stash_auto_row_ids,
|
||||
)
|
||||
from lancedb.expr import col
|
||||
from lancedb.index import FTS
|
||||
from lancedb.schema import blob_column_paths, blob_v2_column_paths
|
||||
|
||||
|
||||
_HIDE_LANCE_BLOB = """\
|
||||
import importlib.abc
|
||||
import sys
|
||||
|
||||
class _MissingLanceBlob(importlib.abc.MetaPathFinder):
|
||||
def find_spec(self, fullname, path, target=None):
|
||||
if fullname == "lance.blob" or fullname.startswith("lance.blob."):
|
||||
raise ModuleNotFoundError(fullname, name="lance.blob")
|
||||
|
||||
sys.modules.pop("lance.blob", None)
|
||||
sys.meta_path.insert(0, _MissingLanceBlob())
|
||||
"""
|
||||
|
||||
|
||||
def _blob_table(name, rows):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
@@ -46,6 +70,181 @@ def test_blob_factory_declares_v2_field():
|
||||
field = lancedb.blob("image")
|
||||
assert isinstance(field.type, pa.ExtensionType)
|
||||
assert field.type.extension_name == "lance.blob.v2"
|
||||
assert lancedb.BlobType is LanceBlobType
|
||||
assert type(field.type) is LanceBlobType
|
||||
|
||||
|
||||
def test_blob_type_works_without_pylance():
|
||||
script = _HIDE_LANCE_BLOB + textwrap.dedent(
|
||||
"""\
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
|
||||
field = lancedb.blob("image")
|
||||
if not isinstance(field.type, pa.ExtensionType):
|
||||
raise SystemExit("expected an extension type")
|
||||
if field.type.extension_name != "lance.blob.v2":
|
||||
raise SystemExit(field.type.extension_name)
|
||||
if lancedb.BlobType is not type(field.type):
|
||||
raise SystemExit("BlobType is not the field type class")
|
||||
if lancedb.BlobType.__module__ != "lancedb.schema":
|
||||
raise SystemExit(lancedb.BlobType.__module__)
|
||||
|
||||
db = lancedb.connect("memory:///")
|
||||
table = db.create_table(
|
||||
"images",
|
||||
schema=pa.schema([pa.field("id", pa.int64()), field]),
|
||||
)
|
||||
table.add([{"id": 1, "image": b"hello"}])
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
|
||||
)
|
||||
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
|
||||
raise SystemExit(
|
||||
f"merge_insert rows updated={result.num_updated_rows} "
|
||||
f"inserted={result.num_inserted_rows}"
|
||||
)
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_blob_resolves_pylance_type_without_eager_import():
|
||||
script = textwrap.dedent(
|
||||
"""\
|
||||
import sys
|
||||
import lancedb
|
||||
|
||||
if "lance.blob" in sys.modules:
|
||||
raise SystemExit("import lancedb imported lance.blob")
|
||||
field = lancedb.blob("image")
|
||||
from lance.blob import BlobType
|
||||
|
||||
if type(field.type) is not BlobType:
|
||||
raise SystemExit(f"{type(field.type)} is not {BlobType}")
|
||||
import lance
|
||||
|
||||
image = lance.blob_array([b"x"])
|
||||
if type(image.type) is not BlobType:
|
||||
raise SystemExit("blob_array used a different class")
|
||||
if type(image.type) is not type(field.type):
|
||||
raise SystemExit("field and array classes differ")
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_blob_fallback_fails_if_name_already_registered():
|
||||
script = _HIDE_LANCE_BLOB + textwrap.dedent(
|
||||
"""\
|
||||
import pyarrow as pa
|
||||
|
||||
class OtherBlobType(pa.ExtensionType):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
pa.struct([pa.field("data", pa.large_binary())]),
|
||||
"lance.blob.v2",
|
||||
)
|
||||
|
||||
def __arrow_ext_serialize__(self):
|
||||
return b""
|
||||
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(cls, storage_type, serialized):
|
||||
return cls()
|
||||
|
||||
pa.register_extension_type(OtherBlobType())
|
||||
import lancedb
|
||||
|
||||
try:
|
||||
lancedb.blob("image")
|
||||
except ValueError as err:
|
||||
if "already registered" not in str(err):
|
||||
raise SystemExit(err)
|
||||
else:
|
||||
raise SystemExit("expected ValueError")
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_blob_type_rejects_competing_registration_with_pylance():
|
||||
script = textwrap.dedent(
|
||||
"""\
|
||||
import pyarrow as pa
|
||||
import pyarrow.ipc
|
||||
|
||||
class OtherBlobType(pa.ExtensionType):
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("data", pa.large_binary()),
|
||||
pa.field("uri", pa.utf8()),
|
||||
pa.field("position", pa.uint64()),
|
||||
pa.field("size", pa.uint64()),
|
||||
]
|
||||
),
|
||||
"lance.blob.v2",
|
||||
)
|
||||
|
||||
def __arrow_ext_serialize__(self):
|
||||
return b""
|
||||
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(cls, storage_type, serialized):
|
||||
return cls()
|
||||
|
||||
pa.register_extension_type(OtherBlobType())
|
||||
|
||||
from lance.blob import BlobType
|
||||
|
||||
if BlobType is OtherBlobType:
|
||||
raise SystemExit("pylance BlobType was replaced")
|
||||
schema = pa.schema([pa.field("value", BlobType())])
|
||||
restored = pa.ipc.read_schema(schema.serialize())
|
||||
if type(restored.field("value").type) is not OtherBlobType:
|
||||
raise SystemExit(type(restored.field("value").type))
|
||||
|
||||
import lancedb
|
||||
|
||||
try:
|
||||
lancedb.blob("image")
|
||||
except ValueError as err:
|
||||
if "__main__.OtherBlobType" not in str(err):
|
||||
raise SystemExit(err)
|
||||
else:
|
||||
raise SystemExit("expected ValueError")
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_blob_v2_column_paths_include_list_children():
|
||||
@@ -70,6 +269,14 @@ def test_blob_v2_column_paths_include_list_children():
|
||||
]
|
||||
|
||||
|
||||
def test_blob_v2_projection_sources_use_typed_column_name():
|
||||
schema = pa.schema([lancedb.blob("blob")])
|
||||
|
||||
assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == {
|
||||
"blob_alias": "blob"
|
||||
}
|
||||
|
||||
|
||||
def _legacy_v1_table(name):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema(
|
||||
@@ -166,6 +373,20 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
|
||||
assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_typed_blob_projection_preserves_source_column():
|
||||
db = await lancedb.connect_async("memory:///typed_blob_projection")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
|
||||
table = await db.create_table("typed_blob_projection", schema=schema)
|
||||
await table.add([{"id": 1, "blob": b"alpha"}])
|
||||
|
||||
hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
|
||||
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = await table.fetch_blobs("blob", hits)
|
||||
assert blobs.to_pylist() == [b"alpha"]
|
||||
|
||||
|
||||
def test_fetch_blobs_round_trip():
|
||||
table = _blob_table(
|
||||
"round_trip",
|
||||
@@ -176,6 +397,292 @@ def test_fetch_blobs_round_trip():
|
||||
assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"]
|
||||
|
||||
|
||||
def test_merge_insert_writes_python_bytes():
|
||||
table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}])
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}])
|
||||
)
|
||||
assert result.num_updated_rows == 1
|
||||
assert result.num_inserted_rows == 1
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
assert blobs.to_pylist() == [b"updated", b"inserted"]
|
||||
|
||||
|
||||
def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("images", schema=schema)
|
||||
table.add([{"id": 1, "image": b"hello"}])
|
||||
|
||||
script = textwrap.dedent(
|
||||
f"""\
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect({str(tmp_path)!r})
|
||||
table = db.open_table("images")
|
||||
image_type = table.schema.field("image").type
|
||||
if type(image_type).__name__ != "StructType":
|
||||
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(
|
||||
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
|
||||
)
|
||||
)
|
||||
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
|
||||
raise SystemExit(
|
||||
f"rows updated={{result.num_updated_rows}} "
|
||||
f"inserted={{result.num_inserted_rows}}"
|
||||
)
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
if blobs.to_pylist() != [b"updated", b"inserted"]:
|
||||
raise SystemExit(blobs.to_pylist())
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("images", schema=schema)
|
||||
table.add([{"id": 1, "image": b"hello"}])
|
||||
|
||||
script = _HIDE_LANCE_BLOB + textwrap.dedent(
|
||||
f"""\
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect({str(tmp_path)!r})
|
||||
table = db.open_table("images")
|
||||
image_type = table.schema.field("image").type
|
||||
if type(image_type).__name__ != "StructType":
|
||||
raise SystemExit(f"expected StructType, got {{type(image_type)}}")
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(
|
||||
[{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}]
|
||||
)
|
||||
)
|
||||
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
|
||||
raise SystemExit(
|
||||
f"rows updated={{result.num_updated_rows}} "
|
||||
f"inserted={{result.num_inserted_rows}}"
|
||||
)
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
if blobs.to_pylist() != [b"updated", b"inserted"]:
|
||||
raise SystemExit(blobs.to_pylist())
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("images", schema=schema)
|
||||
table.add([{"id": 1, "image": b"before"}])
|
||||
|
||||
script = textwrap.dedent(
|
||||
f"""\
|
||||
import pyarrow as pa
|
||||
import lancedb
|
||||
|
||||
db = lancedb.connect({str(tmp_path)!r})
|
||||
table = db.open_table("images")
|
||||
image_type = table.schema.field("image").type
|
||||
if type(image_type).__name__ != "StructType":
|
||||
raise SystemExit(
|
||||
f"expected StructType before lance import, got {{type(image_type)}}"
|
||||
)
|
||||
|
||||
import lance
|
||||
|
||||
updates = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array([1, 2], type=pa.int64()),
|
||||
lance.blob_array([b"updated", b"inserted"]),
|
||||
],
|
||||
names=["id", "image"],
|
||||
)
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(updates)
|
||||
)
|
||||
if result.num_updated_rows != 1 or result.num_inserted_rows != 1:
|
||||
raise SystemExit(
|
||||
f"rows updated={{result.num_updated_rows}} "
|
||||
f"inserted={{result.num_inserted_rows}}"
|
||||
)
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
if blobs.to_pylist() != [b"updated", b"inserted"]:
|
||||
raise SystemExit(blobs.to_pylist())
|
||||
"""
|
||||
)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_add_all_null_blob_column():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("all_null", schema=schema)
|
||||
table.add([{"id": 1, "image": None}, {"id": 2, "image": None}])
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
assert blobs.to_pylist() == [None, None]
|
||||
|
||||
|
||||
def test_create_table_nested_blob_schema_without_rows():
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("info", pa.struct([lancedb.blob("blob")])),
|
||||
pa.field("images", pa.list_(lancedb.blob("image"))),
|
||||
]
|
||||
)
|
||||
table = db.create_table("nested_empty", schema=schema)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_merge_insert_nested_blob_dicts():
|
||||
db = lancedb.connect("memory:///")
|
||||
info = pa.StructArray.from_arrays(
|
||||
[
|
||||
pa.array(["first"], type=pa.string()),
|
||||
_blob_array("blob", [b"before"]),
|
||||
],
|
||||
names=["name", "blob"],
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1], type=pa.int64()), info],
|
||||
names=["id", "info"],
|
||||
)
|
||||
table = db.create_table("nested_merge", data=data)
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.execute([{"id": 1, "info": {"name": "first", "blob": b"after"}}])
|
||||
)
|
||||
assert result.num_updated_rows == 1
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("info.blob", [by_id[1]])
|
||||
assert blobs.to_pylist() == [b"after"]
|
||||
|
||||
|
||||
def _list_blob_table(name):
|
||||
db = lancedb.connect("memory:///")
|
||||
blob_field = lancedb.blob("image")
|
||||
images = pa.ListArray.from_arrays(
|
||||
pa.array([0, 1], type=pa.int32()), _blob_array("image", [b"before"])
|
||||
)
|
||||
data = pa.Table.from_arrays(
|
||||
[pa.array([1], type=pa.int64()), images],
|
||||
schema=pa.schema(
|
||||
[pa.field("id", pa.int64()), pa.field("images", pa.list_(blob_field))]
|
||||
),
|
||||
)
|
||||
return db.create_table(name, data=data)
|
||||
|
||||
|
||||
def test_merge_insert_list_blob_dicts():
|
||||
table = _list_blob_table("list_merge")
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute([{"id": 1, "images": [b"one", b"two"]}, {"id": 2, "images": None}])
|
||||
)
|
||||
assert result.num_updated_rows == 1
|
||||
assert result.num_inserted_rows == 1
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
sizes = {
|
||||
row["id"]: None if row["images"] is None else [d["size"] for d in row["images"]]
|
||||
for row in hits.to_pylist()
|
||||
}
|
||||
assert sizes == {1: [3, 3], 2: None}
|
||||
|
||||
|
||||
def test_list_blob_column_queries_as_raw_descriptors():
|
||||
table = _list_blob_table("list_query")
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
element = hits.schema.field("images").type.value_type
|
||||
assert pa.types.is_struct(element)
|
||||
assert "_lance_row_id" not in element.names
|
||||
with pytest.raises(ValueError, match="expected struct before segment"):
|
||||
table.fetch_blobs("images.image", [0])
|
||||
|
||||
|
||||
def test_row_addressable_paths_exclude_list_children():
|
||||
from lancedb.schema import row_addressable_blob_v2_paths
|
||||
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("info", pa.struct([lancedb.blob("blob")])),
|
||||
pa.field("images", pa.list_(lancedb.blob("image"))),
|
||||
]
|
||||
)
|
||||
assert blob_v2_column_paths(schema) == ["info.blob", "images.image"]
|
||||
assert row_addressable_blob_v2_paths(schema) == ["info.blob"]
|
||||
|
||||
|
||||
def test_merge_insert_writes_pylance_blob_array():
|
||||
table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}])
|
||||
image = lance.blob_array([b"updated", b"inserted"])
|
||||
assert type(image.type) is LanceBlobType
|
||||
assert type(image.type) is type(lancedb.BlobType())
|
||||
updates = pa.Table.from_arrays(
|
||||
[pa.array([1, 2], type=pa.int64()), image], names=["id", "image"]
|
||||
)
|
||||
|
||||
result = (
|
||||
table.merge_insert("id")
|
||||
.when_matched_update_all()
|
||||
.when_not_matched_insert_all()
|
||||
.execute(updates)
|
||||
)
|
||||
|
||||
assert result.num_updated_rows == 1
|
||||
assert result.num_inserted_rows == 1
|
||||
by_id = _row_ids_by_id(table)
|
||||
blobs = table.fetch_blobs("image", [by_id[1], by_id[2]])
|
||||
assert blobs.to_pylist() == [b"updated", b"inserted"]
|
||||
|
||||
|
||||
def test_fetch_blobs_accepts_query_result():
|
||||
table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}])
|
||||
hits = table.search().limit(10).to_arrow()
|
||||
@@ -403,6 +910,50 @@ async def test_blob_v2_hybrid_fetch_blobs_async():
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_hybrid_typed_blob_projection_preserves_source_column():
|
||||
db = await lancedb.connect_async("memory:///hybrid_typed_blob")
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("text", pa.utf8()),
|
||||
pa.field("vector", pa.list_(pa.float32(), list_size=2)),
|
||||
lancedb.blob("blob"),
|
||||
]
|
||||
)
|
||||
table = await db.create_table("hybrid_typed_blob", schema=schema)
|
||||
await table.add(
|
||||
[
|
||||
{
|
||||
"id": 1,
|
||||
"text": "hello alpha",
|
||||
"vector": [1.0, 0.0],
|
||||
"blob": b"alpha",
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"text": "hello beta",
|
||||
"vector": [0.9, 0.1],
|
||||
"blob": b"beta",
|
||||
},
|
||||
]
|
||||
)
|
||||
await table.create_index("text", config=FTS(with_position=False))
|
||||
|
||||
hits = await (
|
||||
table.query()
|
||||
.nearest_to([1.0, 0.0])
|
||||
.nearest_to_text("hello")
|
||||
.select({"blob_alias": col("blob")})
|
||||
.limit(2)
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
|
||||
blobs = await table.fetch_blobs("blob", hits)
|
||||
assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"}
|
||||
|
||||
|
||||
def test_blob_file_seek_read_and_read_range():
|
||||
payload = _identifiable_payload(1024)
|
||||
table = _blob_table("seek_read", [{"id": 1, "image": payload}])
|
||||
@@ -617,3 +1168,71 @@ def test_fetch_blobs_nested_path_survives_sort_after_query():
|
||||
def _identifiable_payload(size: int) -> bytes:
|
||||
block = 256
|
||||
return b"".join(bytes([i % 256]) * block for i in range(size // block))
|
||||
|
||||
|
||||
def _external_uri_blob_array(uris):
|
||||
blob_type = lancedb.blob("image").type
|
||||
storage_type = blob_type.storage_type
|
||||
child_names = [field.name for field in storage_type]
|
||||
assert "uri" in child_names, "blob layout no longer has a uri child"
|
||||
children = [
|
||||
pa.array(uris if field.name == "uri" else [None] * len(uris), type=field.type)
|
||||
for field in storage_type
|
||||
]
|
||||
storage = pa.StructArray.from_arrays(children, fields=list(storage_type))
|
||||
return pa.ExtensionArray.from_storage(blob_type, storage)
|
||||
|
||||
|
||||
def _external_uri_table_and_rows(name, uris):
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table(name, schema=schema)
|
||||
rows = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array(range(len(uris)), type=pa.int64()),
|
||||
_external_uri_blob_array(uris),
|
||||
],
|
||||
schema=schema,
|
||||
)
|
||||
return table, rows
|
||||
|
||||
|
||||
def test_add_external_uri_struct_round_trips_with_flag(tmp_path):
|
||||
payload = b"external-uri-bytes"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
|
||||
table, rows = _external_uri_table_and_rows("external_struct", [blob_path.as_uri()])
|
||||
table.add(rows, allow_external_blob_outside_bases=True)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert blobs[0].as_py() == payload
|
||||
|
||||
|
||||
def test_add_external_uri_without_flag_raises(tmp_path):
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(b"unreachable")
|
||||
|
||||
table, rows = _external_uri_table_and_rows("external_no_flag", [blob_path.as_uri()])
|
||||
with pytest.raises(ValueError, match="allow_external_blob_outside_bases"):
|
||||
table.add(rows)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_add_external_uri_string_round_trips_with_flag(tmp_path):
|
||||
payload = b"external-uri-bytes"
|
||||
blob_path = tmp_path / "payload.bin"
|
||||
blob_path.write_bytes(payload)
|
||||
|
||||
db = lancedb.connect("memory:///")
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
table = db.create_table("external_string", schema=schema)
|
||||
table.add(
|
||||
[{"id": 1, "image": blob_path.as_uri()}],
|
||||
allow_external_blob_outside_bases=True,
|
||||
)
|
||||
|
||||
hits = table.search().to_arrow()
|
||||
blobs = table.fetch_blobs("image", hits)
|
||||
assert blobs[0].as_py() == payload
|
||||
|
||||
@@ -32,6 +32,7 @@ Parameters used throughout:
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
import lancedb
|
||||
@@ -46,6 +47,7 @@ from utils import (
|
||||
torch = pytest.importorskip("torch")
|
||||
streaming = pytest.importorskip("lancedb.streaming")
|
||||
StreamingDataset = streaming.StreamingDataset
|
||||
StreamingDataLoader = streaming.StreamingDataLoader
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset parameters
|
||||
@@ -92,6 +94,27 @@ class FakeWorkerInfo:
|
||||
num_workers: int
|
||||
|
||||
|
||||
def _collate_with_first_batch_error(samples):
|
||||
ids = [sample["id"] for sample in samples]
|
||||
if ids == [0, 1]:
|
||||
raise ValueError("first batch fails")
|
||||
return ids
|
||||
|
||||
|
||||
def _collate_with_first_batch_stop(samples):
|
||||
ids = [sample["id"] for sample in samples]
|
||||
if ids == [0, 1]:
|
||||
raise StopIteration("first batch stopped")
|
||||
return ids
|
||||
|
||||
|
||||
def _collate_with_first_batch_interrupt(samples):
|
||||
ids = [sample["id"] for sample in samples]
|
||||
if ids == [0, 1]:
|
||||
raise KeyboardInterrupt("first batch interrupted")
|
||||
return ids
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1008,6 +1031,565 @@ def test_multi_worker_elastic_det_across_worker_counts(lance_table):
|
||||
# ── Resumability with num_workers ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_streaming_dataloader_commits_only_consumed_worker_batches(tmp_path):
|
||||
"""Prefetched worker state is committed only as the trainer receives it."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"worker_commit", pa.table({"id": [1, 2, 3, 4, 10, 20, 30, 40]})
|
||||
)
|
||||
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
prefetch_factor=4,
|
||||
)
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
first = next(iterator)["id"].tolist()
|
||||
|
||||
assert first == [1, 2]
|
||||
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2, 0]
|
||||
with pytest.raises(RuntimeError, match="complete logical step boundary"):
|
||||
dataset.state_dict()
|
||||
|
||||
second = next(iterator)["id"].tolist()
|
||||
assert second == [10, 20]
|
||||
checkpoint = dataset.state_dict()
|
||||
assert checkpoint["samples_consumed_per_split"] == [2, 2]
|
||||
uninterrupted = [batch["id"].tolist() for batch in iterator]
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
resumed.load_state_dict(checkpoint)
|
||||
resumed_loader = StreamingDataLoader(
|
||||
resumed,
|
||||
batch_size=2,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
prefetch_factor=4,
|
||||
)
|
||||
resumed_iterator = iter(resumed_loader)
|
||||
try:
|
||||
remaining = [batch["id"].tolist() for batch in resumed_iterator]
|
||||
finally:
|
||||
resumed_iterator._shutdown_workers()
|
||||
assert remaining == uninterrupted == [[3, 4], [30, 40]]
|
||||
|
||||
|
||||
def test_distributed_checkpoint_uses_rank_local_worker_boundary(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("rank_boundary", pa.table({"id": list(range(8))}))
|
||||
dataset = StreamingDataset(
|
||||
table,
|
||||
num_splits=4,
|
||||
shuffle=False,
|
||||
rank=0,
|
||||
world_size=2,
|
||||
)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=1,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
assert next(iterator)["id"].tolist() == [0]
|
||||
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [
|
||||
1,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
]
|
||||
with pytest.raises(RuntimeError, match="complete logical step boundary"):
|
||||
dataset.state_dict()
|
||||
|
||||
assert next(iterator)["id"].tolist() == [2]
|
||||
checkpoint = dataset.state_dict()
|
||||
remaining = [batch["id"].tolist() for batch in iterator]
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
assert checkpoint["samples_consumed_per_split"] == [1, 1, 0, 0]
|
||||
assert remaining == [[1], [3]]
|
||||
|
||||
|
||||
def test_standard_dataloader_rejects_stale_parent_checkpoint(tmp_path):
|
||||
"""A standard DataLoader must not expose prefetched producer progress."""
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("untracked_workers", pa.table({"id": [1, 2, 10, 20]}))
|
||||
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
# Merely constructing the checkpoint-aware loader must not authorize a
|
||||
# later plain DataLoader's worker progress.
|
||||
StreamingDataLoader(dataset, batch_size=2, num_workers=0)
|
||||
loader = torch.utils.data.DataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
assert next(iterator)["id"].tolist() == [1, 2]
|
||||
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
|
||||
dataset.state_dict()
|
||||
list(iterator)
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
|
||||
def test_streaming_dataloader_rejects_persistent_workers(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("persistent_workers", pa.table({"id": [1, 2]}))
|
||||
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
|
||||
with pytest.raises(ValueError, match="persistent_workers=True"):
|
||||
StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=1,
|
||||
num_workers=2,
|
||||
persistent_workers=True,
|
||||
)
|
||||
|
||||
|
||||
def test_collate_failure_invalidates_consumer_checkpoint(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"collate_failure", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
|
||||
)
|
||||
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
collate_fn=_collate_with_first_batch_error,
|
||||
prefetch_factor=2,
|
||||
)
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
with pytest.raises(ValueError, match="first batch fails"):
|
||||
next(iterator)
|
||||
assert next(iterator) == [100, 101]
|
||||
assert next(iterator) == [2, 3]
|
||||
with pytest.raises(RuntimeError, match="failed before it was returned"):
|
||||
dataset.state_dict()
|
||||
list(iterator)
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
|
||||
def test_collate_stop_iteration_invalidates_consumer_checkpoint(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("collate_stop", pa.table({"id": list(range(6))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=0,
|
||||
collate_fn=_collate_with_first_batch_stop,
|
||||
)
|
||||
iterator = iter(loader)
|
||||
|
||||
with pytest.raises(RuntimeError, match="collate_fn raised StopIteration"):
|
||||
next(iterator)
|
||||
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
|
||||
with pytest.raises(RuntimeError, match="failed before it was returned"):
|
||||
dataset.state_dict()
|
||||
assert list(iterator) == [[2, 3], [4, 5]]
|
||||
|
||||
|
||||
def test_batch_base_exception_invalidates_consumer_checkpoint(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("collate_interrupt", pa.table({"id": list(range(6))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=0,
|
||||
collate_fn=_collate_with_first_batch_interrupt,
|
||||
)
|
||||
iterator = iter(loader)
|
||||
|
||||
with pytest.raises(KeyboardInterrupt, match="first batch interrupted"):
|
||||
next(iterator)
|
||||
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
|
||||
with pytest.raises(RuntimeError, match="failed before it was returned"):
|
||||
dataset.state_dict()
|
||||
assert list(iterator) == [[2, 3], [4, 5]]
|
||||
|
||||
|
||||
def test_parent_commit_base_exception_invalidates_consumer_checkpoint(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("commit_interrupt", pa.table({"id": list(range(4))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
|
||||
iterator = iter(loader)
|
||||
real_commit = dataset._commit_worker_state
|
||||
|
||||
def interrupt_after_commit(state, *, require_uniform):
|
||||
real_commit(state, require_uniform=require_uniform)
|
||||
raise KeyboardInterrupt("after parent commit")
|
||||
|
||||
with patch.object(
|
||||
dataset, "_commit_worker_state", side_effect=interrupt_after_commit
|
||||
):
|
||||
with pytest.raises(KeyboardInterrupt, match="after parent commit"):
|
||||
next(iterator)
|
||||
|
||||
assert dataset._checkpoint_snapshot()["samples_consumed_per_split"] == [2]
|
||||
with pytest.raises(RuntimeError, match="failed before it was returned"):
|
||||
dataset.state_dict()
|
||||
|
||||
|
||||
def test_direct_iteration_surfaces_prefetch_failure_before_committing_row(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("prefetch_failure", pa.table({"id": list(range(4))}))
|
||||
release = threading.Event()
|
||||
failed = threading.Event()
|
||||
real_getitems = streaming.Permutation.__getitems__
|
||||
|
||||
def controlled_getitems(permutation, indices):
|
||||
if indices and indices[0] >= 2:
|
||||
assert release.wait(timeout=5)
|
||||
failed.set()
|
||||
raise RuntimeError("later prefetched I/O failed")
|
||||
return real_getitems(permutation, indices)
|
||||
|
||||
class SignalDict(dict):
|
||||
def __setitem__(self, key, value):
|
||||
super().__setitem__(key, value)
|
||||
release.set()
|
||||
assert failed.wait(timeout=5)
|
||||
|
||||
monkeypatch.setattr(streaming.Permutation, "__getitems__", controlled_getitems)
|
||||
dataset = StreamingDataset(
|
||||
table,
|
||||
num_splits=1,
|
||||
shuffle=False,
|
||||
read_batch_size=2,
|
||||
io_queue_depth=2,
|
||||
)
|
||||
dataset._resume_positions = SignalDict()
|
||||
iterator = iter(dataset)
|
||||
|
||||
assert next(iterator)["id"] == 0
|
||||
with pytest.raises(RuntimeError, match="later prefetched I/O failed"):
|
||||
next(iterator)
|
||||
|
||||
checkpoint = dataset.state_dict()
|
||||
assert checkpoint["samples_consumed_per_split"] == [1]
|
||||
assert checkpoint["positions_consumed_per_split"] == [1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("workers", [0, 1, 2])
|
||||
def test_streaming_dataloader_rejects_drop_last(tmp_path, workers):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("drop_last", pa.table({"id": [0, 1, 2]}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
worker_options = {"multiprocessing_context": "spawn"} if workers else {}
|
||||
|
||||
with pytest.raises(ValueError, match="drop_last=True"):
|
||||
StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=workers,
|
||||
drop_last=True,
|
||||
**worker_options,
|
||||
)
|
||||
|
||||
|
||||
def test_streaming_dataloader_owns_one_iterator_until_teardown(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("iterator_owner", pa.table({"id": list(range(4))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=1,
|
||||
multiprocessing_context="spawn",
|
||||
)
|
||||
|
||||
first = iter(loader)
|
||||
try:
|
||||
assert next(first)["id"].tolist() == [0, 1]
|
||||
with pytest.raises(RuntimeError, match="concurrent iteration"):
|
||||
iter(loader)
|
||||
finally:
|
||||
first._shutdown_workers()
|
||||
|
||||
second = iter(loader)
|
||||
try:
|
||||
assert [batch["id"].tolist() for batch in second] == [[2, 3]]
|
||||
except BaseException:
|
||||
second._shutdown_workers()
|
||||
raise
|
||||
|
||||
# Natural exhaustion releases ownership too.
|
||||
third = iter(loader)
|
||||
try:
|
||||
assert list(third) == []
|
||||
finally:
|
||||
third._shutdown_workers()
|
||||
|
||||
|
||||
def test_zero_worker_shutdown_closes_inner_iterator_before_release(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("zero_worker_shutdown", pa.table({"id": list(range(6))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
|
||||
|
||||
first = iter(loader)
|
||||
assert next(first)["id"].tolist() == [0, 1]
|
||||
first._shutdown_workers()
|
||||
|
||||
assert dataset._consumer_iterator_active is False
|
||||
assert dataset._raw_batches_ref is None
|
||||
second = iter(loader)
|
||||
try:
|
||||
with pytest.raises(StopIteration):
|
||||
next(first)
|
||||
assert next(second)["id"].tolist() == [2, 3]
|
||||
finally:
|
||||
second._shutdown_workers()
|
||||
|
||||
|
||||
def test_direct_and_loader_admission_share_one_atomic_lease(tmp_path, monkeypatch):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("direct_loader_lease", pa.table({"id": list(range(4))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
direct_result = []
|
||||
direct_error = []
|
||||
contender = []
|
||||
real_resolve = dataset._resolve_my_splits
|
||||
|
||||
def controlled_resolve():
|
||||
if threading.current_thread().name == "direct-start":
|
||||
entered.set()
|
||||
assert release.wait(timeout=5)
|
||||
return real_resolve()
|
||||
|
||||
def advance_direct(iterator):
|
||||
try:
|
||||
direct_result.append(next(iterator)["id"])
|
||||
except BaseException as exc:
|
||||
direct_error.append(exc)
|
||||
|
||||
monkeypatch.setattr(dataset, "_resolve_my_splits", controlled_resolve)
|
||||
direct = iter(dataset)
|
||||
thread = threading.Thread(
|
||||
target=advance_direct, args=(direct,), name="direct-start"
|
||||
)
|
||||
thread.start()
|
||||
assert entered.wait(timeout=5)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="concurrent iteration"):
|
||||
contender.append(iter(loader))
|
||||
finally:
|
||||
release.set()
|
||||
thread.join(timeout=5)
|
||||
if contender:
|
||||
contender[0]._shutdown_workers()
|
||||
direct.close()
|
||||
|
||||
assert not thread.is_alive()
|
||||
assert direct_error == []
|
||||
assert direct_result == [0]
|
||||
|
||||
|
||||
def test_loader_acquires_before_snapshot_and_cleans_interrupted_acquire(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("lease_snapshot", pa.table({"id": list(range(4))}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(dataset, batch_size=2, num_workers=0)
|
||||
first = iter(loader)
|
||||
assert next(first)["id"].tolist() == [0, 1]
|
||||
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
pending = []
|
||||
pending_errors = []
|
||||
observed_snapshots = []
|
||||
real_acquire = dataset._acquire_consumer_iterator
|
||||
real_snapshot = dataset._checkpoint_snapshot
|
||||
|
||||
def controlled_acquire():
|
||||
if threading.current_thread().name == "stale-start":
|
||||
entered.set()
|
||||
assert release.wait(timeout=5)
|
||||
return real_acquire()
|
||||
|
||||
def recording_snapshot():
|
||||
state = real_snapshot()
|
||||
if threading.current_thread().name == "stale-start":
|
||||
observed_snapshots.append(state["samples_consumed_per_split"])
|
||||
return state
|
||||
|
||||
def create_pending_iterator():
|
||||
try:
|
||||
pending.append(iter(loader))
|
||||
except BaseException as exc:
|
||||
pending_errors.append(exc)
|
||||
|
||||
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", controlled_acquire)
|
||||
monkeypatch.setattr(dataset, "_checkpoint_snapshot", recording_snapshot)
|
||||
thread = threading.Thread(target=create_pending_iterator, name="stale-start")
|
||||
thread.start()
|
||||
assert entered.wait(timeout=5)
|
||||
assert next(first)["id"].tolist() == [2, 3]
|
||||
with pytest.raises(StopIteration):
|
||||
next(first)
|
||||
release.set()
|
||||
thread.join(timeout=5)
|
||||
|
||||
assert not thread.is_alive()
|
||||
assert pending_errors == []
|
||||
assert observed_snapshots == [[4]]
|
||||
assert len(pending) == 1
|
||||
assert list(pending[0]) == []
|
||||
assert dataset.state_dict()["samples_consumed_per_split"] == [4]
|
||||
|
||||
def interrupted_acquire():
|
||||
real_acquire()
|
||||
raise KeyboardInterrupt("after acquire")
|
||||
|
||||
monkeypatch.setattr(dataset, "_acquire_consumer_iterator", interrupted_acquire)
|
||||
with pytest.raises(KeyboardInterrupt, match="after acquire"):
|
||||
iter(loader)
|
||||
assert dataset._consumer_iterator_active is False
|
||||
|
||||
|
||||
def test_consumer_iterator_lease_publication_is_atomic(tmp_path, monkeypatch):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("atomic_lease", pa.table({"id": [0, 1]}))
|
||||
dataset = StreamingDataset(table, num_splits=1, shuffle=False)
|
||||
loader = StreamingDataLoader(dataset, batch_size=1, num_workers=0)
|
||||
real_get_ident = streaming.threading.get_ident
|
||||
calls = 0
|
||||
|
||||
def interrupt_during_publication():
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
raise KeyboardInterrupt("during lease mutation")
|
||||
return real_get_ident()
|
||||
|
||||
monkeypatch.setattr(streaming.threading, "get_ident", interrupt_during_publication)
|
||||
with pytest.raises(KeyboardInterrupt, match="during lease mutation"):
|
||||
iter(loader)
|
||||
monkeypatch.setattr(streaming.threading, "get_ident", real_get_ident)
|
||||
|
||||
assert dataset._consumer_iterator_active is False
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
assert next(iterator)["id"].tolist() == [0]
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
|
||||
def test_streaming_dataloader_rejects_dataset_iter_override(tmp_path):
|
||||
class CustomizedDataset(StreamingDataset):
|
||||
def __iter__(self):
|
||||
return iter([1000, 1001])
|
||||
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("custom_iteration", pa.table({"id": [0, 1, 2]}))
|
||||
dataset = CustomizedDataset(table, num_splits=1, shuffle=False)
|
||||
|
||||
assert list(dataset) == [1000, 1001]
|
||||
with pytest.raises(TypeError, match="override __iter__"):
|
||||
StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=2,
|
||||
num_workers=0,
|
||||
collate_fn=list,
|
||||
)
|
||||
|
||||
|
||||
def test_interleaved_adapters_do_not_authorize_plain_iteration(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table_a = db.create_table("adapter_a", pa.table({"id": [0, 1]}))
|
||||
table_b = db.create_table("adapter_b", pa.table({"id": [10, 11]}))
|
||||
dataset_a = StreamingDataset(table_a, num_splits=1, shuffle=False)
|
||||
dataset_b = StreamingDataset(table_b, num_splits=1, shuffle=False)
|
||||
initial_state = dataset_a.state_dict()
|
||||
|
||||
owner_a = dataset_a._acquire_consumer_iterator()
|
||||
owner_b = dataset_b._acquire_consumer_iterator()
|
||||
try:
|
||||
iterator_a = iter(streaming._StreamingDatasetAdapter(dataset_a))
|
||||
iterator_b = iter(streaming._StreamingDatasetAdapter(dataset_b))
|
||||
assert next(iterator_a).data["id"] == 0
|
||||
assert next(iterator_b).data["id"] == 10
|
||||
assert [sample.data["id"] for sample in iterator_a] == [1]
|
||||
assert [sample.data["id"] for sample in iterator_b] == [11]
|
||||
finally:
|
||||
dataset_a._release_consumer_iterator(owner_a)
|
||||
dataset_b._release_consumer_iterator(owner_b)
|
||||
|
||||
dataset_a.load_state_dict(initial_state)
|
||||
with patch(
|
||||
"lancedb.streaming.get_worker_info",
|
||||
return_value=FakeWorkerInfo(id=0, num_workers=1),
|
||||
):
|
||||
plain_iterator = iter(dataset_a)
|
||||
assert next(plain_iterator)["id"] == 0
|
||||
plain_iterator.close()
|
||||
|
||||
assert dataset_a._untracked_worker_iteration[0] == 1
|
||||
with pytest.raises(RuntimeError, match="Use StreamingDataLoader"):
|
||||
dataset_a.state_dict()
|
||||
|
||||
|
||||
def test_resume_from_partial_split_cycle_preserves_remaining_order(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("partial_cycle", pa.table({"id": [1, 2, 10, 20]}))
|
||||
dataset = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
iterator = iter(dataset)
|
||||
|
||||
assert next(iterator)["id"] == 1
|
||||
checkpoint = dataset.state_dict()
|
||||
iterator.close()
|
||||
assert checkpoint["samples_consumed_per_split"] == [1, 0]
|
||||
|
||||
resumed = StreamingDataset(table, num_splits=2, shuffle=False)
|
||||
resumed.load_state_dict(checkpoint)
|
||||
assert [row["id"] for row in resumed] == [10, 2, 20]
|
||||
|
||||
|
||||
def test_partial_cycle_resume_preserves_skip_truncation(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table(
|
||||
"partial_skip", pa.table({"id": [0, 1, 2, 3, 100, 101, 102, 103]})
|
||||
)
|
||||
kwargs = dict(
|
||||
num_splits=2,
|
||||
shuffle=False,
|
||||
transform=_failing_transform({1, 2, 3}),
|
||||
on_transform_error="skip",
|
||||
)
|
||||
dataset = StreamingDataset(table, **kwargs)
|
||||
iterator = iter(dataset)
|
||||
|
||||
assert next(iterator)["id"] == 0
|
||||
checkpoint = dataset.state_dict()
|
||||
uninterrupted = [row["id"] for row in iterator]
|
||||
|
||||
resumed = StreamingDataset(table, **kwargs)
|
||||
resumed.load_state_dict(checkpoint)
|
||||
assert [row["id"] for row in resumed] == uninterrupted == [100]
|
||||
|
||||
|
||||
def test_multi_worker_resumability_same_topology(lance_table):
|
||||
"""Checkpoint with num_workers=2, resume with num_workers=2: exact continuation."""
|
||||
world_size = 1
|
||||
@@ -2018,6 +2600,23 @@ def test_merge_state_dicts_validates_consistency(lance_table):
|
||||
StreamingDataset.merge_state_dicts([])
|
||||
|
||||
|
||||
def test_merge_state_dicts_combines_nonuniform_consumer_progress(lance_table):
|
||||
dataset = StreamingDataset(
|
||||
lance_table, num_splits=2, shuffle=False, shuffle_seed=SHUFFLE_SEED
|
||||
)
|
||||
rank0 = dataset.state_dict()
|
||||
rank0["samples_consumed_per_split"] = [2, 0]
|
||||
rank0["positions_consumed_per_split"] = [2, 0]
|
||||
rank1 = dataset.state_dict()
|
||||
rank1["samples_consumed_per_split"] = [0, 2]
|
||||
rank1["positions_consumed_per_split"] = [0, 2]
|
||||
|
||||
merged = StreamingDataset.merge_state_dicts([rank0, rank1])
|
||||
|
||||
assert merged["samples_consumed_per_split"] == [2, 2]
|
||||
assert merged["positions_consumed_per_split"] == [2, 2]
|
||||
|
||||
|
||||
def test_load_state_dict_without_positions_key(lance_table):
|
||||
"""Checkpoints from before positions_consumed_per_split existed still
|
||||
resume exactly (positions equal sample counts when nothing is skipped)."""
|
||||
@@ -2254,6 +2853,65 @@ def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):
|
||||
]
|
||||
|
||||
|
||||
def test_packed_checkpoint_requires_complete_split_cycle(tmp_path):
|
||||
table = _create_token_table(tmp_path, [[1], [2], [10], [20]])
|
||||
dataset = _packed_dataset(table, pack_sequences=3, blocks_per_epoch=4, num_splits=2)
|
||||
iterator = iter(dataset)
|
||||
|
||||
next(iterator)
|
||||
with pytest.raises(RuntimeError, match="complete logical step boundary"):
|
||||
dataset.state_dict()
|
||||
|
||||
next(iterator)
|
||||
assert dataset.state_dict()["blocks_emitted_per_split"] == [1, 1]
|
||||
iterator.close()
|
||||
|
||||
|
||||
def test_streaming_dataloader_commits_consumed_packed_batches(tmp_path):
|
||||
table = _create_token_table(
|
||||
tmp_path,
|
||||
[[1], [2], [3], [4], [10], [20], [30], [40]],
|
||||
)
|
||||
kwargs = dict(pack_sequences=4, blocks_per_epoch=4, num_splits=2)
|
||||
dataset = _packed_dataset(table, **kwargs)
|
||||
loader = StreamingDataLoader(
|
||||
dataset,
|
||||
batch_size=1,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
prefetch_factor=2,
|
||||
)
|
||||
iterator = iter(loader)
|
||||
try:
|
||||
next(iterator)
|
||||
with pytest.raises(RuntimeError, match="complete logical step boundary"):
|
||||
dataset.state_dict()
|
||||
|
||||
next(iterator)
|
||||
checkpoint = dataset.state_dict()
|
||||
uninterrupted = [batch["input_ids"].tolist() for batch in iterator]
|
||||
finally:
|
||||
iterator._shutdown_workers()
|
||||
|
||||
resumed = _packed_dataset(table, **kwargs)
|
||||
resumed.load_state_dict(checkpoint)
|
||||
resumed_loader = StreamingDataLoader(
|
||||
resumed,
|
||||
batch_size=1,
|
||||
num_workers=2,
|
||||
multiprocessing_context="spawn",
|
||||
prefetch_factor=2,
|
||||
)
|
||||
resumed_iterator = iter(resumed_loader)
|
||||
try:
|
||||
remaining = [batch["input_ids"].tolist() for batch in resumed_iterator]
|
||||
finally:
|
||||
resumed_iterator._shutdown_workers()
|
||||
|
||||
assert checkpoint["blocks_emitted_per_split"] == [1, 1]
|
||||
assert remaining == uninterrupted
|
||||
|
||||
|
||||
def test_pack_sequences_validates_configuration_and_tokens(tmp_path):
|
||||
table = _create_token_table(tmp_path, [[1, 2]])
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ class TestExprConstruction:
|
||||
def test_func(self):
|
||||
e = func("lower", col("name"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "lower(name)"
|
||||
assert e.to_sql() == "lower(`name`)"
|
||||
|
||||
def test_func_unknown_raises(self):
|
||||
with pytest.raises(Exception):
|
||||
@@ -115,7 +115,7 @@ class TestExprOperators:
|
||||
def test_and_operator(self):
|
||||
e = (col("age") > lit(18)) & (col("status") == lit("active"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "((age > 18) AND (status = 'active'))"
|
||||
assert e.to_sql() == "((age > 18) AND (`status` = 'active'))"
|
||||
|
||||
def test_or_operator(self):
|
||||
e = (col("a") == lit(1)) | (col("b") == lit(2))
|
||||
@@ -166,7 +166,7 @@ class TestExprOperators:
|
||||
def test_coerce_plain_str(self):
|
||||
e = col("name") == "alice"
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(name = 'alice')"
|
||||
assert e.to_sql() == "(`name` = 'alice')"
|
||||
|
||||
def test_reflexive_comparisons(self):
|
||||
# 10 < col("age") swaps to col("age") > 10
|
||||
@@ -198,85 +198,85 @@ class TestExprBytesLiteral:
|
||||
|
||||
def test_bytes_equality_expr_sql(self):
|
||||
e = col("data") == lit(b"\xca\xfe")
|
||||
assert e.to_sql() == "(data = X'CAFE')"
|
||||
assert e.to_sql() == "(`data` = X'CAFE')"
|
||||
|
||||
def test_bytes_ne_expr_sql(self):
|
||||
e = col("data") != lit(b"\xff")
|
||||
assert e.to_sql() == "(data <> X'FF')"
|
||||
assert e.to_sql() == "(`data` <> X'FF')"
|
||||
|
||||
def test_bytes_compound_expr_sql(self):
|
||||
e = (col("data") == lit(b"\x01")) & (col("id") > lit(5))
|
||||
assert e.to_sql() == "((data = X'01') AND (id > 5))"
|
||||
assert e.to_sql() == "((`data` = X'01') AND (id > 5))"
|
||||
|
||||
def test_bytes_in_function_call(self):
|
||||
# Regression test: binary literals inside scalar function calls
|
||||
# used to fail because DataFusion's unparser does not support Binary
|
||||
# scalars. Now handled via a placeholder-substitution rewrite.
|
||||
e = func("contains", col("data"), lit(b"\xff"))
|
||||
assert e.to_sql() == "contains(data, X'FF')"
|
||||
assert e.to_sql() == "contains(`data`, X'FF')"
|
||||
|
||||
def test_bytes_in_not(self):
|
||||
e = ~(col("data") == lit(b"\xff"))
|
||||
assert e.to_sql() == "NOT (data = X'FF')"
|
||||
assert e.to_sql() == "NOT (`data` = X'FF')"
|
||||
|
||||
|
||||
class TestExprStringMethods:
|
||||
def test_lower(self):
|
||||
e = col("name").lower()
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "lower(name)"
|
||||
assert e.to_sql() == "lower(`name`)"
|
||||
|
||||
def test_upper(self):
|
||||
e = col("name").upper()
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "upper(name)"
|
||||
assert e.to_sql() == "upper(`name`)"
|
||||
|
||||
def test_contains(self):
|
||||
e = col("text").contains(lit("hello"))
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "contains(text, 'hello')"
|
||||
assert e.to_sql() == "contains(`text`, 'hello')"
|
||||
|
||||
def test_contains_with_str_coerce(self):
|
||||
e = col("text").contains("hello")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "contains(text, 'hello')"
|
||||
assert e.to_sql() == "contains(`text`, 'hello')"
|
||||
|
||||
def test_chained_lower_eq(self):
|
||||
e = col("name").lower() == lit("alice")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "(lower(name) = 'alice')"
|
||||
assert e.to_sql() == "(lower(`name`) = 'alice')"
|
||||
|
||||
|
||||
class TestExprCast:
|
||||
def test_cast_string(self):
|
||||
e = col("id").cast("string")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(id AS VARCHAR)"
|
||||
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
|
||||
|
||||
def test_cast_int32(self):
|
||||
e = col("score").cast("int32")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(score AS INTEGER)"
|
||||
assert e.to_sql() == "arrow_cast(score, 'Int32')"
|
||||
|
||||
def test_cast_float64(self):
|
||||
e = col("val").cast("float64")
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(val AS DOUBLE)"
|
||||
assert e.to_sql() == "arrow_cast(val, 'Float64')"
|
||||
|
||||
def test_cast_pyarrow_type(self):
|
||||
e = col("score").cast(pa.int32())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(score AS INTEGER)"
|
||||
assert e.to_sql() == "arrow_cast(score, 'Int32')"
|
||||
|
||||
def test_cast_pyarrow_float64(self):
|
||||
e = col("val").cast(pa.float64())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(val AS DOUBLE)"
|
||||
assert e.to_sql() == "arrow_cast(val, 'Float64')"
|
||||
|
||||
def test_cast_pyarrow_string(self):
|
||||
e = col("id").cast(pa.string())
|
||||
assert isinstance(e, Expr)
|
||||
assert e.to_sql() == "CAST(id AS VARCHAR)"
|
||||
assert e.to_sql() == "arrow_cast(id, 'Utf8')"
|
||||
|
||||
def test_cast_pyarrow_and_string_equivalent(self):
|
||||
# pa.int32() and "int32" should produce equivalent SQL
|
||||
@@ -597,14 +597,14 @@ class TestExprIsin:
|
||||
def test_isin_strs(self):
|
||||
assert (
|
||||
col("status").isin(["active", "pending"]).to_sql()
|
||||
== "status IN ('active', 'pending')"
|
||||
== "`status` IN ('active', 'pending')"
|
||||
)
|
||||
|
||||
def test_isin_coerces_and_mixes(self):
|
||||
assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)"
|
||||
|
||||
def test_isin_empty(self):
|
||||
assert col("id").isin([]).to_sql() == "id IN ()"
|
||||
assert col("id").isin([]).to_sql() == "false"
|
||||
|
||||
def test_isin_filter(self, simple_table):
|
||||
result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow()
|
||||
|
||||
@@ -37,21 +37,6 @@ def job_result(name: str) -> dict:
|
||||
return json.loads(fixture(name))["result"]
|
||||
|
||||
|
||||
def assert_no_secret_values(value):
|
||||
if isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
assert key not in {
|
||||
"secret_value",
|
||||
"secret_values",
|
||||
"resolved_secret",
|
||||
"resolved_secrets",
|
||||
}
|
||||
assert_no_secret_values(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
assert_no_secret_values(child)
|
||||
|
||||
|
||||
def test_public_function_values_are_in_api_reference():
|
||||
docs = Path(__file__).parents[3] / "docs" / "src" / "python" / "python.md"
|
||||
rendered = docs.read_text()
|
||||
@@ -109,7 +94,6 @@ def test_function_version_identity_is_immutable_and_exact():
|
||||
version = FunctionVersion.from_json(json.dumps(value))
|
||||
assert version.name == "embed"
|
||||
assert version.version == "fv_01K3EXACT"
|
||||
assert version.required_secrets == ("HF_TOKEN",)
|
||||
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
version.version = "fv_changed"
|
||||
@@ -121,7 +105,7 @@ def test_function_version_identity_is_immutable_and_exact():
|
||||
assert FunctionVersion(**changed) != version
|
||||
|
||||
|
||||
def test_function_version_binds_named_columns_as_one_immutable_group():
|
||||
def test_function_version_binds_named_columns_as_one_immutable_application():
|
||||
version = FunctionVersion.from_json(
|
||||
json.dumps(job_result("remote_function_job.json"))
|
||||
)
|
||||
@@ -131,13 +115,10 @@ def test_function_version_binds_named_columns_as_one_immutable_group():
|
||||
assert application.function.name == version.name
|
||||
assert application.function.version == version.version
|
||||
assert application.output is version.signature.output
|
||||
assert application.group_id.startswith("fg_")
|
||||
assert [
|
||||
(value.parameter, value.kind, value.value["path"])
|
||||
for value in application.inputs
|
||||
] == [("text", "column", "documents.body")]
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
application.group_id = "fg_changed"
|
||||
|
||||
|
||||
def test_function_version_binding_validates_names_and_direct_columns():
|
||||
@@ -156,7 +137,7 @@ def test_function_version_binding_validates_names_and_direct_columns():
|
||||
def test_function_version_keeps_named_struct_outputs_in_one_application():
|
||||
value = job_result("remote_function_job.json")
|
||||
value["name"] = "text_features"
|
||||
value["version"] = "fv_grouped"
|
||||
value["version"] = "fv_multi_output"
|
||||
value["signature"] = {
|
||||
"inputs": [
|
||||
{"name": "title", "arrow_type": "utf8", "nullable": True},
|
||||
@@ -221,7 +202,6 @@ def test_function_application_uses_rename_columns_only():
|
||||
assert application.columns["normalized_text"] == "search_text"
|
||||
assert renamed.columns["normalized_text"] == "body_normalized"
|
||||
assert renamed.function == application.function
|
||||
assert renamed.group_id == application.group_id
|
||||
assert not hasattr(application, "rename_outputs")
|
||||
with pytest.raises(TypeError, match="immutable"):
|
||||
renamed.columns["normalized_text"] = "changed"
|
||||
@@ -242,7 +222,6 @@ def test_function_application_uses_rename_columns_only():
|
||||
|
||||
def test_binding_and_refresh_result_keep_stable_remote_fields():
|
||||
binding = FunctionBinding.from_json(fixture("remote_function_binding.json"))
|
||||
assert binding.revision == 3
|
||||
assert binding.function.version == "fv_01K3TEXT"
|
||||
assert [output.output_ordinal for output in binding.outputs] == [0, 1]
|
||||
assert binding.input_schema is not None
|
||||
@@ -297,15 +276,6 @@ def test_refresh_result_rejects_non_u64_values(field):
|
||||
RefreshColumnResult.from_json(json.dumps(value))
|
||||
|
||||
|
||||
def test_canonical_client_values_contain_secret_names_only():
|
||||
version = FunctionVersion.from_json(
|
||||
json.dumps(job_result("remote_function_job.json"))
|
||||
)
|
||||
canonical = json.loads(version.to_canonical_json())
|
||||
assert canonical["required_secrets"] == ["HF_TOKEN"]
|
||||
assert_no_secret_values(canonical)
|
||||
|
||||
|
||||
class _FunctionDeclarationInner:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
@@ -322,7 +292,7 @@ def known_application() -> FunctionApplication:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically():
|
||||
async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomically():
|
||||
inner = _FunctionDeclarationInner()
|
||||
table = AsyncTable(inner)
|
||||
application = known_application()
|
||||
@@ -343,12 +313,12 @@ async def test_add_columns_routes_struct_as_one_and_grouped_expansion_atomically
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_columns_rejects_mixed_groups_and_unknown_newer_application():
|
||||
async def test_add_columns_rejects_multiple_bindings_and_unknown_newer_application():
|
||||
inner = _FunctionDeclarationInner()
|
||||
table = AsyncTable(inner)
|
||||
application = known_application()
|
||||
|
||||
with pytest.raises(ValueError, match="exactly one Function sibling group"):
|
||||
with pytest.raises(ValueError, match="exactly one Function binding"):
|
||||
await table.add_columns({"a": application, "b": application})
|
||||
|
||||
future = json.loads(fixture("remote_function_application.json"))
|
||||
@@ -376,7 +346,6 @@ def test_rename_requires_named_struct_and_keeps_partial_mapping_immutable():
|
||||
"arrow_type": "list<float32>",
|
||||
"nullable": False,
|
||||
},
|
||||
"group_id": "fg_scalar",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
@@ -3,7 +3,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import contextlib
|
||||
import functools
|
||||
import importlib.util
|
||||
import types
|
||||
from datetime import date
|
||||
import http.server
|
||||
import json
|
||||
from pathlib import Path
|
||||
@@ -16,6 +21,9 @@ import pytest
|
||||
import lancedb
|
||||
from lancedb.functions import UdfDefinition, udf
|
||||
|
||||
THRESHOLD = 20
|
||||
_CACHE = None
|
||||
|
||||
|
||||
FIXTURES = (
|
||||
Path(__file__).parents[3]
|
||||
@@ -31,28 +39,12 @@ FIXTURES = (
|
||||
@udf(
|
||||
pip=["numpy>=2"],
|
||||
env={"MODE": "test"},
|
||||
secrets=["API_TOKEN"],
|
||||
python_version="3.12",
|
||||
)
|
||||
def normalize_score(value: float) -> float:
|
||||
return value / 100.0
|
||||
|
||||
|
||||
def _assert_no_secret_values(value):
|
||||
if isinstance(value, dict):
|
||||
for key, child in value.items():
|
||||
assert key not in {
|
||||
"secret_value",
|
||||
"secret_values",
|
||||
"resolved_secret",
|
||||
"resolved_secrets",
|
||||
}
|
||||
_assert_no_secret_values(child)
|
||||
elif isinstance(value, list):
|
||||
for child in value:
|
||||
_assert_no_secret_values(child)
|
||||
|
||||
|
||||
def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
assert isinstance(normalize_score, UdfDefinition)
|
||||
assert normalize_score(25.0) == 0.25
|
||||
@@ -67,13 +59,417 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable():
|
||||
"kind": "scalar_to_arrow_batch",
|
||||
"version": 1,
|
||||
}
|
||||
assert request["required_secrets"] == ["API_TOKEN"]
|
||||
_assert_no_secret_values(request)
|
||||
|
||||
|
||||
def _run_packaged(definition, *args):
|
||||
"""Execute the shipped artifact in a fresh namespace, as a worker would."""
|
||||
source = base64.b64decode(definition.registration_request.artifact.content.data)
|
||||
namespace: dict = {}
|
||||
exec(compile(source, "<udf>", "exec"), namespace)
|
||||
return namespace[definition.registration_request.artifact.entrypoint](*args)
|
||||
|
||||
|
||||
def test_udf_conda_environment():
|
||||
@udf(conda=["scipy", "numpy"], conda_channels=["conda-forge", "defaults"])
|
||||
def halve(value: float) -> float:
|
||||
return value / 2
|
||||
|
||||
request = json.loads(halve.registration_request.to_canonical_json())
|
||||
assert request["runtime"]["environment"] == {
|
||||
"kind": "conda",
|
||||
"packages": ["numpy", "scipy"],
|
||||
"channels": ["conda-forge", "defaults"],
|
||||
}
|
||||
pip_request = json.loads(normalize_score.registration_request.to_canonical_json())
|
||||
assert "channels" not in pip_request["runtime"]["environment"]
|
||||
|
||||
with pytest.raises(ValueError, match="not both"):
|
||||
udf(name="both", pip=["numpy"], conda=["numpy"])(lambda value: value)
|
||||
with pytest.raises(ValueError, match="requires conda"):
|
||||
udf(name="channels", conda_channels=["conda-forge"])(lambda value: value)
|
||||
|
||||
|
||||
def test_udf_packages_attribute_access_and_body_imports():
|
||||
@udf
|
||||
def word_norm(body: str) -> float:
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
words = body.split()
|
||||
except AttributeError as error:
|
||||
raise ValueError(str(error)) from error
|
||||
return float(np.linalg.norm([len(w) for w in words]))
|
||||
|
||||
assert _run_packaged(word_norm, "aa bb") == pytest.approx(8**0.5)
|
||||
|
||||
|
||||
def test_udf_packages_module_globals_and_global_caches():
|
||||
@udf
|
||||
def label(value: int) -> str:
|
||||
return "big" if value >= THRESHOLD else "small"
|
||||
|
||||
assert _run_packaged(label, 21) == "big"
|
||||
|
||||
@udf
|
||||
def cached(value: int) -> int:
|
||||
global _CACHE
|
||||
if _CACHE is None:
|
||||
_CACHE = 40
|
||||
return _CACHE + value
|
||||
|
||||
assert _run_packaged(cached, 2) == 42
|
||||
|
||||
|
||||
def test_udf_annotations_are_not_runtime_names():
|
||||
@udf
|
||||
def identity(value: date) -> date:
|
||||
return value
|
||||
|
||||
assert _run_packaged(identity, date(2026, 8, 25)) == date(2026, 8, 25)
|
||||
|
||||
|
||||
def test_udf_nested_scopes_resolve_lexically():
|
||||
@udf
|
||||
def score(value: int) -> int:
|
||||
offset = 2
|
||||
|
||||
def add_offset() -> int:
|
||||
return value + offset
|
||||
|
||||
return add_offset() + sum(v for v in [0])
|
||||
|
||||
assert _run_packaged(score, 3) == 5
|
||||
|
||||
|
||||
def test_udf_resolves_module_globals_before_builtins(tmp_path):
|
||||
module_path = tmp_path / "shadowing_udfs.py"
|
||||
module_path.write_text(
|
||||
"max = 7\n"
|
||||
"len = lambda _: 99\n"
|
||||
"\n"
|
||||
"def uses_literal_shadow(value: int) -> int:\n"
|
||||
" def nested() -> int:\n"
|
||||
" return max\n"
|
||||
" return nested() + value\n"
|
||||
"\n"
|
||||
"def uses_callable_shadow(value: int) -> int:\n"
|
||||
" def nested() -> int:\n"
|
||||
" return len([1])\n"
|
||||
" return nested() + value\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("shadowing_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# The module's `max = 7` is what the interpreter would use, so it ships.
|
||||
assert _run_packaged(udf(module.uses_literal_shadow), 1) == 8
|
||||
# A callable global cannot ship; it must not be silently swapped for the builtin.
|
||||
with pytest.raises(TypeError, match="unsupported global value of type function"):
|
||||
udf(module.uses_callable_shadow)
|
||||
|
||||
|
||||
def test_canonical_arrow_type_is_exactly_the_grammar():
|
||||
from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
|
||||
|
||||
golden = json.loads(
|
||||
(
|
||||
Path(__file__).parents[3]
|
||||
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
|
||||
).read_text()
|
||||
)
|
||||
primitives = [
|
||||
case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
|
||||
]
|
||||
assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
|
||||
for outside in [
|
||||
pa.timestamp("us"),
|
||||
pa.decimal128(10, 2),
|
||||
pa.large_string(),
|
||||
pa.large_binary(),
|
||||
pa.binary(4),
|
||||
pa.duration("s"),
|
||||
pa.struct([pa.field("a", pa.int32())]),
|
||||
pa.list_(pa.float32(), 0),
|
||||
pa.list_(pa.timestamp("us")),
|
||||
]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(outside)
|
||||
|
||||
|
||||
def test_udf_nested_annotations_are_postponed_in_the_artifact():
|
||||
@udf
|
||||
def score(value: int) -> int:
|
||||
def identity(item: date) -> date:
|
||||
return item
|
||||
|
||||
identity(date(2026, 8, 25))
|
||||
return value
|
||||
|
||||
assert _run_packaged(score, 3) == 3
|
||||
|
||||
|
||||
def test_udf_ships_globals_the_body_deletes():
|
||||
@udf
|
||||
def clear(value: int) -> int:
|
||||
global _CACHE
|
||||
del _CACHE
|
||||
return value
|
||||
|
||||
assert _run_packaged(clear, 3) == 3
|
||||
|
||||
|
||||
def test_udf_rejects_a_module_global_that_does_not_import_as_itself(tmp_path):
|
||||
module_path = tmp_path / "fake_module_udfs.py"
|
||||
module_path.write_text(
|
||||
"import types\n"
|
||||
"np = types.ModuleType('numpy')\n"
|
||||
"np.sqrt = lambda x: 0\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return int(np.sqrt(value))\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("fake_module_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
with pytest.raises(TypeError, match="does not import as 'numpy'"):
|
||||
udf(module.score)
|
||||
|
||||
|
||||
def test_udf_rejects_a_module_level_namespace_alias(tmp_path):
|
||||
module_path = tmp_path / "aliasing_udfs.py"
|
||||
module_path.write_text(
|
||||
"import builtins as b\n"
|
||||
"THRESHOLD = 5\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return value + b.vars(b.__import__('aliasing_udfs'))['THRESHOLD']\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("aliasing_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
with pytest.raises(ValueError, match="dynamic namespace access"):
|
||||
udf(module.score)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"access",
|
||||
[
|
||||
"globals()['THRESHOLD']",
|
||||
"eval('THRESHOLD')",
|
||||
"(lambda g: g()['THRESHOLD'])(globals)",
|
||||
"__import__('sys').modules[__name__].THRESHOLD",
|
||||
"sys.modules[__name__].THRESHOLD",
|
||||
],
|
||||
)
|
||||
def test_udf_rejects_dynamic_namespace_access(access):
|
||||
namespace: dict = {}
|
||||
exec(
|
||||
f"def score(value: int) -> int:\n return value + {access}\n",
|
||||
{"THRESHOLD": 5},
|
||||
namespace,
|
||||
)
|
||||
with pytest.raises(ValueError, match="dynamic namespace access"):
|
||||
_package_from_text(
|
||||
"def score(value: int) -> int:\n"
|
||||
" import sys\n"
|
||||
f" return value + {access}\n"
|
||||
)
|
||||
|
||||
|
||||
def _package_from_text(source: str, module_globals: dict | None = None):
|
||||
"""Load `source` as a real module file so the packager can inspect it."""
|
||||
import tempfile
|
||||
|
||||
directory = tempfile.mkdtemp()
|
||||
path = Path(directory) / "generated_udf_module.py"
|
||||
path.write_text(source)
|
||||
spec = importlib.util.spec_from_file_location(f"generated_udf_{id(source)}", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
if module_globals:
|
||||
module.__dict__.update(module_globals)
|
||||
spec.loader.exec_module(module)
|
||||
functions = [
|
||||
value
|
||||
for value in vars(module).values()
|
||||
if callable(value) and getattr(value, "__module__", None) == module.__name__
|
||||
]
|
||||
return udf(functions[0])
|
||||
|
||||
|
||||
def test_udf_rejects_a_non_standard_builtins_environment():
|
||||
def score(value: int) -> int:
|
||||
return len([1]) + value
|
||||
|
||||
score.__globals__ # noqa: B018 -- real function, real globals
|
||||
import builtins
|
||||
|
||||
patched = types.FunctionType(
|
||||
score.__code__,
|
||||
{"__builtins__": {**vars(builtins), "len": lambda _: 99}},
|
||||
"score",
|
||||
)
|
||||
patched.__annotations__ = score.__annotations__
|
||||
assert patched(3) == 102
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(patched)
|
||||
|
||||
class ReportingDict(dict): # reports standard entries, resolves differently
|
||||
def __missing__(self, key):
|
||||
return vars(builtins)[key]
|
||||
|
||||
disguised = types.FunctionType(
|
||||
score.__code__, {"__builtins__": ReportingDict(len=lambda _: 99)}, "score"
|
||||
)
|
||||
disguised.__annotations__ = score.__annotations__
|
||||
assert disguised(3) == 102
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(disguised)
|
||||
|
||||
hooked = types.FunctionType(
|
||||
score.__code__,
|
||||
{"__builtins__": {**vars(builtins), "__import__": lambda *a, **k: None}},
|
||||
"score",
|
||||
)
|
||||
hooked.__annotations__ = score.__annotations__
|
||||
with pytest.raises(ValueError, match="non-standard builtins environment"):
|
||||
udf(hooked)
|
||||
|
||||
|
||||
def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
|
||||
module_path = tmp_path / "rebound_udfs.py"
|
||||
module_path.write_text(
|
||||
"def fact(value: int) -> int:\n"
|
||||
" return 1 if value <= 1 else value * fact(value - 1)\n"
|
||||
"\n"
|
||||
"def score(value: int) -> int:\n"
|
||||
" return score + value\n"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("rebound_udfs", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
assert _run_packaged(udf(module.fact), 5) == 120
|
||||
raw = module.score
|
||||
module.score = 10
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw)
|
||||
# A wrapper that merely exposes __wrapped__ is not the function.
|
||||
module.score = functools.wraps(raw)(lambda value: 41)
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw)
|
||||
# The decorator's own result is; a subclass of it is not.
|
||||
module.fact = udf(module.fact)
|
||||
assert _run_packaged(module.fact, 4) == 24
|
||||
|
||||
class Twisted(UdfDefinition):
|
||||
def __call__(self, *args, **kwargs):
|
||||
return 41
|
||||
|
||||
raw_fact = module.fact._function
|
||||
module.fact = Twisted(
|
||||
raw_fact,
|
||||
name=None,
|
||||
input_schema=None,
|
||||
output_schema=None,
|
||||
pip=(),
|
||||
env={},
|
||||
python_version=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match="binds that name to another value"):
|
||||
udf(raw_fact)
|
||||
|
||||
|
||||
def test_canonical_arrow_type_rejects_unrepresentable_list_children():
|
||||
from lancedb.functions import _canonical_arrow_type
|
||||
|
||||
for outside in [
|
||||
pa.list_(pa.float32()), # pyarrow default: nullable child
|
||||
pa.list_(pa.field("custom", pa.float32(), nullable=False)),
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
|
||||
]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(outside)
|
||||
assert (
|
||||
_canonical_arrow_type(
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 3)
|
||||
)
|
||||
== "fixed_size_list<float32, 3>"
|
||||
)
|
||||
|
||||
|
||||
def _calls_missing(value: int) -> int:
|
||||
return missing(value) # noqa: F821
|
||||
|
||||
|
||||
def _shadows_missing_in_a_comprehension(value: int) -> int:
|
||||
return missing(value) + sum(missing for missing in ()) # noqa: F821
|
||||
|
||||
|
||||
def _shadows_missing_in_a_lambda(value: int) -> int:
|
||||
return (lambda missing: missing)(value) + missing # noqa: F821
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"function",
|
||||
[_calls_missing, _shadows_missing_in_a_comprehension, _shadows_missing_in_a_lambda],
|
||||
)
|
||||
def test_udf_rejects_a_truly_unresolved_global(function):
|
||||
with pytest.raises(ValueError, match=r"unresolved global names: \['missing'\]"):
|
||||
udf(function)
|
||||
|
||||
|
||||
def _arrow_type_from_golden(spec: dict) -> pa.DataType:
|
||||
kind = spec["type"]
|
||||
if kind in ("list", "large_list", "fixed_size_list"):
|
||||
item = _arrow_type_from_golden(spec["fields"][0]["type"])
|
||||
field = pa.field("item", item, nullable=False)
|
||||
if kind == "list":
|
||||
return pa.list_(field)
|
||||
if kind == "large_list":
|
||||
return pa.large_list(field)
|
||||
return pa.list_(field, spec["length"])
|
||||
return {
|
||||
"null": pa.null(),
|
||||
"bool": pa.bool_(),
|
||||
"utf8": pa.string(),
|
||||
"binary": pa.binary(),
|
||||
"float16": pa.float16(),
|
||||
"float32": pa.float32(),
|
||||
"float64": pa.float64(),
|
||||
"date32": pa.date32(),
|
||||
"date64": pa.date64(),
|
||||
}.get(kind) or getattr(pa, kind)()
|
||||
|
||||
|
||||
def test_arrow_type_grammar_matches_the_shared_golden():
|
||||
golden = json.loads(
|
||||
(
|
||||
Path(__file__).parents[3]
|
||||
/ "rust/lancedb/tests/fixtures/first_class_functions/v1/arrow_types.json"
|
||||
).read_text()
|
||||
)
|
||||
from lancedb.functions import _canonical_arrow_type
|
||||
|
||||
emitted = {
|
||||
case["arrow_type"]: _canonical_arrow_type(_arrow_type_from_golden(case["json"]))
|
||||
for case in golden["valid"]
|
||||
}
|
||||
assert emitted == {
|
||||
case["arrow_type"]: case["arrow_type"] for case in golden["valid"]
|
||||
}
|
||||
assert not set(emitted) & set(golden["invalid"])
|
||||
for case in golden["server_only"]:
|
||||
with pytest.raises(TypeError, match="unsupported Arrow type"):
|
||||
_canonical_arrow_type(_arrow_type_from_golden(case["json"]))
|
||||
|
||||
|
||||
def test_explicit_arrow_schema_is_deterministic():
|
||||
input_schema = pa.schema([pa.field("value", pa.float32(), nullable=True)])
|
||||
output_schema = pa.field("embedding", pa.list_(pa.float32(), 3), nullable=False)
|
||||
output_schema = pa.field(
|
||||
"embedding",
|
||||
pa.list_(pa.field("item", pa.float32(), nullable=False), 3),
|
||||
nullable=False,
|
||||
)
|
||||
|
||||
@udf(input_schema=input_schema, output_schema=output_schema)
|
||||
def explicit(value):
|
||||
@@ -82,7 +478,7 @@ def test_explicit_arrow_schema_is_deterministic():
|
||||
signature = explicit.registration_request.signature
|
||||
assert signature.inputs[0].arrow_type == "float32"
|
||||
assert signature.inputs[0].nullable is True
|
||||
assert signature.output.arrow_type == "fixed_size_list<float32>[3]"
|
||||
assert signature.output.arrow_type == "fixed_size_list<float32, 3>"
|
||||
assert signature.output.nullable is False
|
||||
|
||||
|
||||
@@ -130,14 +526,6 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
|
||||
return value
|
||||
|
||||
|
||||
def test_environment_rejects_secret_value_overlap():
|
||||
with pytest.raises(ValueError, match="must be disjoint"):
|
||||
|
||||
@udf(env={"TOKEN": "plaintext"}, secrets=["TOKEN"])
|
||||
def overlapping(value: int) -> int:
|
||||
return value
|
||||
|
||||
|
||||
def test_local_function_catalog_operations_are_not_supported(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
message = "Function catalog operations are not supported by this database"
|
||||
@@ -174,7 +562,6 @@ def _mock_remote_function_catalog():
|
||||
"runtime": body["runtime"],
|
||||
"runtime_digest": "sha256:runtime",
|
||||
"environment_digest": "sha256:environment",
|
||||
"required_secrets": body.get("required_secrets", []),
|
||||
"created_at": "2026-08-21T00:00:00Z",
|
||||
}
|
||||
response = {"job_id": "job-register"}
|
||||
@@ -233,7 +620,6 @@ def test_remote_registration_job_and_exact_version_reopen_round_trip():
|
||||
assert create_request == json.loads(
|
||||
normalize_score.registration_request.to_canonical_json()
|
||||
)
|
||||
_assert_no_secret_values(create_request)
|
||||
|
||||
|
||||
def test_blocking_remote_registration_returns_function_version():
|
||||
|
||||
@@ -25,6 +25,7 @@ from lancedb.db import DBConnection
|
||||
from lancedb.index import FTS
|
||||
from lancedb.query import (
|
||||
BoostQuery,
|
||||
DocumentGranularity,
|
||||
MatchQuery,
|
||||
MultiMatchQuery,
|
||||
PhraseQuery,
|
||||
@@ -245,6 +246,55 @@ def test_create_inverted_index_rejects_invalid_block_size(table):
|
||||
table.create_index("text", config=FTS(block_size=129))
|
||||
|
||||
|
||||
def test_list_element_document_granularity(tmp_path):
|
||||
docs_type = pa.list_(pa.struct([pa.field("content", pa.string())]))
|
||||
docs = pa.array(
|
||||
[
|
||||
[
|
||||
{"content": "alpha beta"},
|
||||
None,
|
||||
{"content": ""},
|
||||
{"content": "the and"},
|
||||
{"content": "alpha beta"},
|
||||
]
|
||||
],
|
||||
type=docs_type,
|
||||
)
|
||||
table = ldb.connect(tmp_path).create_table(
|
||||
"list_element_docs", pa.table({"id": [0], "docs": docs})
|
||||
)
|
||||
row_table = ldb.connect(tmp_path).create_table(
|
||||
"row_docs", pa.table({"id": [0], "docs": docs})
|
||||
)
|
||||
row_table.create_index("docs.content", config=FTS())
|
||||
row_result = row_table.search(MatchQuery("alpha", "docs.content")).to_arrow()
|
||||
assert row_result.num_rows == 1
|
||||
assert "_doc_index" not in row_result.column_names
|
||||
|
||||
granularity = DocumentGranularity.LIST_ELEMENT
|
||||
table.create_index(
|
||||
"docs.content",
|
||||
config=FTS(with_position=True, document_granularity=granularity),
|
||||
)
|
||||
assert table.list_indices()[0].columns == ["docs.content"]
|
||||
|
||||
def coordinates(query):
|
||||
result = table.search(query).limit(10).to_arrow()
|
||||
doc_index_type = result.schema.field("_doc_index").type
|
||||
assert pa.types.is_list(doc_index_type)
|
||||
assert doc_index_type.value_type == pa.uint32()
|
||||
return sorted(result["_doc_index"].to_pylist())
|
||||
|
||||
assert coordinates(
|
||||
MatchQuery("alpha", "docs.content", document_granularity=granularity)
|
||||
) == [[0], [4]]
|
||||
assert coordinates(
|
||||
PhraseQuery("alpha beta", "docs.content", document_granularity=granularity)
|
||||
) == [[0], [4]]
|
||||
assert coordinates(MatchQuery("alpha", "docs.content")) == [[0], [4]]
|
||||
assert FTS().document_granularity is DocumentGranularity.ROW
|
||||
|
||||
|
||||
def test_create_inverted_index_respects_build_memory_limit(table):
|
||||
with pytest.raises(ValueError, match="exceeds worker memory limit"):
|
||||
table.create_index(
|
||||
@@ -1089,6 +1139,20 @@ def test_fts_query_to_json():
|
||||
)
|
||||
assert json_str == expected
|
||||
|
||||
# Test MatchQuery with list-element document granularity
|
||||
match_query = MatchQuery(
|
||||
"hello world",
|
||||
"text",
|
||||
document_granularity=DocumentGranularity.LIST_ELEMENT,
|
||||
)
|
||||
json_str = match_query.to_json()
|
||||
expected = (
|
||||
'{"match":{"column":"text","terms":"hello world","boost":1.0,'
|
||||
'"fuzziness":0,"max_expansions":50,"operator":"Or","prefix_length":0,'
|
||||
'"document_granularity":"list_element"}}'
|
||||
)
|
||||
assert json_str == expected
|
||||
|
||||
# Test MatchQuery with options
|
||||
match_query = MatchQuery("puppy", "text", fuzziness=2, boost=1.5, prefix_length=3)
|
||||
json_str = match_query.to_json()
|
||||
@@ -1098,6 +1162,19 @@ def test_fts_query_to_json():
|
||||
)
|
||||
assert json_str == expected
|
||||
|
||||
# Test PhraseQuery with list-element document granularity
|
||||
phrase_query = PhraseQuery(
|
||||
"quick brown fox",
|
||||
"title",
|
||||
document_granularity=DocumentGranularity.LIST_ELEMENT,
|
||||
)
|
||||
json_str = phrase_query.to_json()
|
||||
expected = (
|
||||
'{"phrase":{"column":"title","terms":"quick brown fox","slop":0,'
|
||||
'"document_granularity":"list_element"}}'
|
||||
)
|
||||
assert json_str == expected
|
||||
|
||||
# Test PhraseQuery
|
||||
phrase_query = PhraseQuery("quick brown fox", "title")
|
||||
json_str = phrase_query.to_json()
|
||||
|
||||
@@ -56,6 +56,31 @@ def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
|
||||
assert permutation_tbl._conn.read_consistency_interval is None
|
||||
|
||||
|
||||
def test_pickled_permutation_reads_pinned_version(tmp_path):
|
||||
"""An unpickled copy must still read the pinned version, which also covers the
|
||||
version surviving the ``to_arrow()`` round trip in ``__getstate__``."""
|
||||
import pickle
|
||||
|
||||
db = connect(tmp_path)
|
||||
tbl = db.create_table("base", pa.table({"idx": range(20)}))
|
||||
permutation_tbl = permutation_builder(tbl).execute()
|
||||
perm = Permutation.from_tables(tbl, permutation_tbl)
|
||||
|
||||
payload = pickle.dumps(perm)
|
||||
|
||||
# Compact so the stored row addresses no longer describe these rows at latest.
|
||||
tbl.delete("true")
|
||||
tbl.optimize()
|
||||
assert tbl.count_rows() == 0
|
||||
|
||||
# Unpickle after the mutation: __setstate__ reopens at latest, so this only
|
||||
# passes if the recorded version is applied on reopen.
|
||||
restored = pickle.loads(payload)
|
||||
assert len(restored) == 20
|
||||
rows = restored.__getitems__(list(range(20)))
|
||||
assert sorted(row["idx"] for row in rows) == list(range(20))
|
||||
|
||||
|
||||
def test_split_random_counts(mem_db):
|
||||
"""Test random splitting with absolute counts."""
|
||||
tbl = mem_db.create_table(
|
||||
|
||||
@@ -675,6 +675,21 @@ def test_distance_range(table: lancedb.table.Table):
|
||||
assert res["_distance"].to_pylist() == [min_dist, max_dist]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expression", ["1 - _distance", "1.0 - _distance"])
|
||||
def test_select_arithmetic_with_distance(table, expression):
|
||||
result = (
|
||||
table.search([10, 10])
|
||||
.select({"similarity": expression, "_distance": "_distance"})
|
||||
.distance_type("cosine")
|
||||
.to_arrow()
|
||||
)
|
||||
|
||||
assert result.schema.field("similarity").type == pa.float32()
|
||||
assert result["similarity"].to_pylist() == pytest.approx(
|
||||
[1 - distance for distance in result["_distance"].to_pylist()]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_distance_range_async(table_async: AsyncTable):
|
||||
q = [0, 0]
|
||||
@@ -897,6 +912,23 @@ def test_query_builder_batches(table):
|
||||
assert rs_list["id"][1] == 2
|
||||
|
||||
|
||||
def test_batch_vector_query_shares_filtered_flat_scan(table):
|
||||
query = (
|
||||
table.search([[1.0, 2.0], [3.0, 4.0]])
|
||||
.where("id > 0", prefilter=True)
|
||||
.limit(1)
|
||||
.select(["id"])
|
||||
)
|
||||
|
||||
plan = query.explain_plan(verbose=True)
|
||||
assert "KNNVectorDistance: queries=2" in plan
|
||||
assert "UnionExec" not in plan
|
||||
|
||||
results = query.to_arrow()
|
||||
assert len(results) == 2
|
||||
assert results["query_index"].to_pylist() == [0, 1]
|
||||
|
||||
|
||||
def test_dynamic_projection(table):
|
||||
rs = (
|
||||
LanceVectorQueryBuilder(table, [0, 0], "vector")
|
||||
|
||||
@@ -1618,6 +1618,49 @@ def test_query_sync_fts():
|
||||
)
|
||||
|
||||
|
||||
def test_query_sync_fts_document_granularity():
|
||||
from lancedb.query import DocumentGranularity, MatchQuery
|
||||
|
||||
def handler(body):
|
||||
assert body == {
|
||||
"full_text_query": {
|
||||
"query": {
|
||||
"match": {
|
||||
"column": "docs.content",
|
||||
"terms": "alpha",
|
||||
"boost": 1.0,
|
||||
"fuzziness": 0,
|
||||
"max_expansions": 50,
|
||||
"operator": "Or",
|
||||
"prefix_length": 0,
|
||||
"document_granularity": "list_element",
|
||||
}
|
||||
}
|
||||
},
|
||||
"k": 10,
|
||||
"prefilter": True,
|
||||
"vector": [],
|
||||
"version": None,
|
||||
}
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 1],
|
||||
"_doc_index": pa.array([[0], [4]], type=pa.list_(pa.uint32())),
|
||||
}
|
||||
)
|
||||
|
||||
with query_test_table(handler, server_version=Version("0.6.0")) as table:
|
||||
result = table.search(
|
||||
MatchQuery(
|
||||
"alpha",
|
||||
"docs.content",
|
||||
document_granularity=DocumentGranularity.LIST_ELEMENT,
|
||||
)
|
||||
).to_arrow()
|
||||
|
||||
assert result["_doc_index"].to_pylist() == [[0], [4]]
|
||||
|
||||
|
||||
def test_query_sync_hybrid():
|
||||
def handler(body):
|
||||
if "full_text_query" in body:
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import timedelta
|
||||
import threading
|
||||
|
||||
@@ -86,6 +87,25 @@ def test_s3_lifecycle(s3_bucket: str):
|
||||
asyncio.run(test())
|
||||
|
||||
|
||||
@pytest.mark.s3_test
|
||||
def test_concurrent_open_table(s3_bucket: str):
|
||||
uri = f"s3://{s3_bucket}/test_concurrent_open_table"
|
||||
db = lancedb.connect(uri, storage_options=copy.copy(CONFIG))
|
||||
db.create_table("test", pa.table({"x": [1, 2, 3]}))
|
||||
|
||||
num_workers = 32
|
||||
barrier = threading.Barrier(num_workers)
|
||||
|
||||
def open_and_count(_):
|
||||
barrier.wait()
|
||||
return db.open_table("test").count_rows()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as pool:
|
||||
row_counts = list(pool.map(open_and_count, range(num_workers)))
|
||||
|
||||
assert row_counts == [3] * num_workers
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def kms_key():
|
||||
kms = get_boto3_client("kms", endpoint_url=CONFIG["aws_endpoint"])
|
||||
|
||||
@@ -11,6 +11,7 @@ import warnings
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from decimal import Decimal
|
||||
from time import sleep
|
||||
from typing import List
|
||||
from unittest.mock import patch
|
||||
@@ -336,6 +337,21 @@ async def test_update_async(mem_db_async: AsyncConnection):
|
||||
assert await table.count_rows("id == 10") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection):
|
||||
values = ["5", "4.66e-84", "it's"]
|
||||
table = await mem_db_async.create_table(
|
||||
"update_expr_literals",
|
||||
data=[{"field": value, "result": "original"} for value in values],
|
||||
)
|
||||
|
||||
for value in values:
|
||||
update_res = await table.update({"result": value}, where=col("field") == value)
|
||||
assert update_res.rows_updated == 1
|
||||
|
||||
assert (await table.to_arrow())["result"].to_pylist() == values
|
||||
|
||||
|
||||
def test_create_table(mem_db: DBConnection):
|
||||
schema = pa.schema(
|
||||
{
|
||||
@@ -2343,6 +2359,148 @@ def test_update(mem_db: DBConnection):
|
||||
assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]]))
|
||||
|
||||
|
||||
def test_update_expr_filter_literals(mem_db: DBConnection):
|
||||
values = ["5", "4.66e-84", "it's"]
|
||||
table = mem_db.create_table(
|
||||
"update_expr_literals",
|
||||
data=[{"field": value, "result": "original"} for value in values],
|
||||
)
|
||||
|
||||
for value in values:
|
||||
update_res = table.update(where=col("field") == value, values={"result": value})
|
||||
assert update_res.rows_updated == 1
|
||||
|
||||
assert table.to_arrow()["result"].to_pylist() == values
|
||||
|
||||
|
||||
def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
|
||||
low = Decimal("1.234567890123456789")
|
||||
high = Decimal("1.234567890123456790")
|
||||
decimal_schema = pa.schema(
|
||||
[("val", pa.decimal128(19, 18)), ("result", pa.string())]
|
||||
)
|
||||
decimal_table = mem_db.create_table(
|
||||
"update_expr_decimal",
|
||||
pa.table(
|
||||
{"val": [low, high], "result": ["old", "old"]},
|
||||
schema=decimal_schema,
|
||||
),
|
||||
)
|
||||
predicate = col("val") < lit(high)
|
||||
assert decimal_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = decimal_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
keyword_table = mem_db.create_table(
|
||||
"update_expr_keyword", [{"null": 1, "result": "old"}]
|
||||
)
|
||||
predicate = col("null") == 1
|
||||
assert keyword_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = keyword_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
empty_in_table = mem_db.create_table(
|
||||
"update_expr_empty_in", [{"id": 1, "result": "old"}]
|
||||
)
|
||||
predicate = col("id").isin([])
|
||||
assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0
|
||||
result = empty_in_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 0
|
||||
|
||||
marker = "__lancedb_binary_placeholder_0__"
|
||||
binary_schema = pa.schema(
|
||||
[("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())]
|
||||
)
|
||||
binary_table = mem_db.create_table(
|
||||
"update_expr_binary",
|
||||
pa.table(
|
||||
{
|
||||
"payload": [b"\x01", b"\x02"],
|
||||
"text": ["other", marker],
|
||||
"result": ["old", "old"],
|
||||
},
|
||||
schema=binary_schema,
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker)
|
||||
assert binary_table.search().where(predicate).to_arrow().num_rows == 2
|
||||
result = binary_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 2
|
||||
|
||||
nonfinite_table = mem_db.create_table(
|
||||
"update_expr_nonfinite",
|
||||
[{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
|
||||
)
|
||||
predicate = col("x") < float("inf")
|
||||
assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
|
||||
result = nonfinite_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 2
|
||||
|
||||
float16_table = mem_db.create_table(
|
||||
"update_expr_float16",
|
||||
[{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
|
||||
)
|
||||
predicate = col("x").cast(pa.float16()) < 2.0
|
||||
assert float16_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = float16_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
string_cast_table = mem_db.create_table(
|
||||
"update_expr_string_cast",
|
||||
[{"x": 1, "result": "old"}, {"x": 2, "result": "old"}],
|
||||
)
|
||||
predicate = col("x").cast("string") == "1"
|
||||
assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = string_cast_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
quoted_identifier_schema = pa.schema(
|
||||
[("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())]
|
||||
)
|
||||
quoted_identifier_table = mem_db.create_table(
|
||||
"update_expr_quoted_identifier",
|
||||
pa.table(
|
||||
{"payload": [b"\x01"], "odd'name": [1], "result": ["old"]},
|
||||
schema=quoted_identifier_schema,
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1)
|
||||
assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = quoted_identifier_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
decimal256_schema = pa.schema(
|
||||
[("val", pa.decimal256(40, 2)), ("result", pa.string())]
|
||||
)
|
||||
decimal256_table = mem_db.create_table(
|
||||
"update_expr_decimal256",
|
||||
pa.table(
|
||||
{
|
||||
"val": [Decimal("1.00"), Decimal("3.00")],
|
||||
"result": ["old", "old"],
|
||||
},
|
||||
schema=decimal256_schema,
|
||||
),
|
||||
)
|
||||
predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
|
||||
assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
|
||||
result = decimal256_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 1
|
||||
|
||||
binary_empty_table = mem_db.create_table(
|
||||
"update_expr_binary_empty",
|
||||
pa.table(
|
||||
{"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
|
||||
schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
|
||||
),
|
||||
)
|
||||
predicate = (col("payload") == lit(b"\x01")).isin([])
|
||||
assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
|
||||
assert predicate.to_sql() == "false"
|
||||
result = binary_empty_table.update(where=predicate, values={"result": "new"})
|
||||
assert result.rows_updated == 0
|
||||
|
||||
|
||||
def test_update_with_arrow_scalar(mem_db: DBConnection):
|
||||
schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})
|
||||
table = mem_db.create_table("my_table", schema=schema)
|
||||
@@ -3929,6 +4087,29 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path):
|
||||
table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"})
|
||||
|
||||
|
||||
def test_computed_column_blob_projection_inherits_semantics(tmp_path):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")])
|
||||
db = lancedb.connect(tmp_path)
|
||||
table = db.create_table("computed_column_blob", schema=schema)
|
||||
table.add(
|
||||
[
|
||||
{"id": 1, "image": b"hello"},
|
||||
{"id": 2, "image": b""},
|
||||
{"id": 3, "image": None},
|
||||
]
|
||||
)
|
||||
|
||||
table.add_columns(computed={"image_copy": "image", "second_copy": "image_copy"})
|
||||
assert table.refresh_column("image_copy").rows_filled == 2
|
||||
assert table.refresh_column("second_copy").rows_filled == 2
|
||||
assert table.blob_columns() == ["image", "image_copy", "second_copy"]
|
||||
|
||||
hits = table.search().with_row_id(True).limit(10).to_arrow()
|
||||
rows = sorted(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist()))
|
||||
copied = table.fetch_blobs("second_copy", [row_id for _, row_id in rows])
|
||||
assert copied.to_pylist() == [b"hello", b"", None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_computed_column_async(tmp_path):
|
||||
db = await lancedb.connect_async(tmp_path)
|
||||
|
||||
@@ -7,6 +7,7 @@ import pathlib
|
||||
from typing import Optional
|
||||
|
||||
import lance
|
||||
from lance.blob import BlobType as LanceBlobType
|
||||
from lancedb.conftest import MockTextEmbeddingFunction
|
||||
from lancedb.embeddings.base import EmbeddingFunctionConfig
|
||||
from lancedb.embeddings.registry import EmbeddingFunctionRegistry
|
||||
@@ -907,6 +908,165 @@ def test_cast_to_target_schema():
|
||||
assert output == expected
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_binary_to_blob_v2():
|
||||
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
image = output["image"].chunk(0)
|
||||
assert type(image.type) is lancedb.BlobType
|
||||
assert image.storage.to_pylist() == [
|
||||
{"data": b"hello", "uri": None, "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct():
|
||||
storage = lancedb.blob("image").type.storage_type
|
||||
target = pa.schema(
|
||||
[
|
||||
pa.field(
|
||||
"image",
|
||||
storage,
|
||||
metadata={
|
||||
b"ARROW:extension:name": b"lance.blob.v2",
|
||||
b"ARROW:extension:metadata": b"",
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())})
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
image = output["image"].chunk(0)
|
||||
assert not isinstance(image.type, pa.ExtensionType)
|
||||
assert image.to_pylist() == [
|
||||
{"data": b"hello", "uri": None, "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_nested_binary_blob():
|
||||
data = pa.table(
|
||||
{
|
||||
"info": pa.array(
|
||||
[{"blob": b"hello"}, {"blob": None}],
|
||||
type=pa.struct([pa.field("blob", pa.binary())]),
|
||||
)
|
||||
}
|
||||
)
|
||||
target = pa.schema([pa.field("info", pa.struct([lancedb.blob("blob")]))])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
blob = output["info"].chunk(0).field("blob")
|
||||
assert type(blob.type) is lancedb.BlobType
|
||||
assert blob.storage.to_pylist() == [
|
||||
{"data": b"hello", "uri": None, "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
def test_cast_to_target_schema_coerces_list_binary_blob_with_inferred_child_name():
|
||||
data = pa.table(
|
||||
{"images": pa.array([[b"a", b"b"], None], type=pa.list_(pa.binary()))}
|
||||
)
|
||||
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
images = output["images"].chunk(0)
|
||||
assert images.type.value_field.name == "image"
|
||||
assert type(images.type.value_type) is lancedb.BlobType
|
||||
assert images.to_pylist()[1] is None
|
||||
assert images.values.storage.to_pylist() == [
|
||||
{"data": b"a", "uri": None, "position": None, "size": None},
|
||||
{"data": b"b", "uri": None, "position": None, "size": None},
|
||||
]
|
||||
|
||||
|
||||
def test_list_blob_coercion_preserves_null_slots_with_nonzero_extent():
|
||||
child = pa.field("image", pa.binary())
|
||||
source = pa.ListArray.from_arrays(
|
||||
pa.array([0, 2, 4], type=pa.int32()),
|
||||
pa.array([b"a", b"b", b"dead", b"beef"], type=pa.binary()),
|
||||
mask=pa.array([False, True]),
|
||||
).cast(pa.list_(child))
|
||||
target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))])
|
||||
|
||||
output = _cast_to_target_schema(
|
||||
pa.table({"images": source}).to_reader(), target
|
||||
).read_all()
|
||||
|
||||
images = output["images"].chunk(0)
|
||||
assert images.to_pylist()[1] is None
|
||||
assert [b["data"] for b in images.to_pylist()[0]] == [b"a", b"b"]
|
||||
|
||||
|
||||
def test_fixed_size_list_blob_coercion_keeps_null_rows():
|
||||
child = pa.field("frame", pa.binary())
|
||||
source = (
|
||||
pa.FixedSizeListArray.from_arrays(
|
||||
pa.array([b"a", b"b", b"c", b"d"], type=pa.binary()), 2
|
||||
)
|
||||
.take(pa.array([0, None], type=pa.int32()))
|
||||
.cast(pa.list_(child, 2))
|
||||
)
|
||||
target = pa.schema([pa.field("frames", pa.list_(lancedb.blob("frame"), 2))])
|
||||
|
||||
output = _cast_to_target_schema(
|
||||
pa.table({"frames": source}).to_reader(), target
|
||||
).read_all()
|
||||
|
||||
frames = output["frames"].chunk(0)
|
||||
assert frames.to_pylist()[1] is None
|
||||
assert [b["data"] for b in frames.to_pylist()[0]] == [b"a", b"b"]
|
||||
|
||||
|
||||
def test_cast_to_target_schema_accepts_pylance_blob_v2():
|
||||
target_type = lancedb.BlobType()
|
||||
source = lance.blob_array([b"hello", None])
|
||||
assert type(source.type) is LanceBlobType
|
||||
assert type(source.type) is type(target_type)
|
||||
data = pa.table({"image": source})
|
||||
target = pa.schema([pa.field("image", target_type)])
|
||||
|
||||
output = _cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
image = output["image"].chunk(0)
|
||||
assert type(image.type) is LanceBlobType
|
||||
assert image.type == target_type
|
||||
assert image.storage.to_pylist() == [
|
||||
{"data": b"hello", "uri": None, "position": None, "size": None},
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
def test_cast_to_target_schema_rejects_different_blob_v2_class():
|
||||
class OtherBlobType(pa.ExtensionType):
|
||||
def __init__(self):
|
||||
super().__init__(lancedb.BlobType().storage_type, "lance.blob.v2")
|
||||
|
||||
def __arrow_ext_serialize__(self) -> bytes:
|
||||
return b""
|
||||
|
||||
@classmethod
|
||||
def __arrow_ext_deserialize__(
|
||||
cls, storage_type: pa.DataType, serialized: bytes
|
||||
) -> "OtherBlobType":
|
||||
return cls()
|
||||
|
||||
storage = lance.blob_array([b"hello"]).storage
|
||||
source = pa.ExtensionArray.from_storage(OtherBlobType(), storage)
|
||||
data = pa.table({"image": source})
|
||||
target = pa.schema([lancedb.blob("image")])
|
||||
|
||||
with pytest.raises(pa.ArrowTypeError, match="different extension type"):
|
||||
_cast_to_target_schema(data.to_reader(), target).read_all()
|
||||
|
||||
|
||||
def test_sanitize_data_stream():
|
||||
# Make sure we don't collect the whole stream when running sanitize_data
|
||||
schema = pa.schema({"a": pa.int32()})
|
||||
|
||||
@@ -130,6 +130,14 @@ impl PyExpr {
|
||||
|
||||
// ── utilities ────────────────────────────────────────────────────────────
|
||||
|
||||
/// Return the referenced column name for a bare column expression.
|
||||
fn column_name(&self) -> Option<String> {
|
||||
match &self.0 {
|
||||
DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Render the expression as a SQL string (useful for debugging).
|
||||
fn to_sql(&self) -> PyResult<String> {
|
||||
lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string()))
|
||||
|
||||
+8
-2
@@ -8,7 +8,7 @@ use lancedb::index::vector::{
|
||||
};
|
||||
use lancedb::index::{
|
||||
Index as LanceDbIndex,
|
||||
scalar::{BTreeIndexBuilder, FmIndexBuilder, FtsIndexBuilder},
|
||||
scalar::{BTreeIndexBuilder, DocumentGranularity, FmIndexBuilder, FtsIndexBuilder},
|
||||
};
|
||||
use pyo3::IntoPyObject;
|
||||
use pyo3::types::PyStringMethods;
|
||||
@@ -60,7 +60,11 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
.ngram_min_length(params.ngram_min_length)
|
||||
.ngram_max_length(params.ngram_max_length)
|
||||
.ngram_prefix_only(params.prefix_only)
|
||||
.custom_stop_words(params.custom_stop_words);
|
||||
.custom_stop_words(params.custom_stop_words)
|
||||
.document_granularity(
|
||||
DocumentGranularity::try_from(params.document_granularity.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?,
|
||||
);
|
||||
if let Some(memory_limit) = params.memory_limit {
|
||||
inner_opts = inner_opts.memory_limit_mb(memory_limit);
|
||||
}
|
||||
@@ -221,6 +225,7 @@ struct FtsParams {
|
||||
block_size: usize,
|
||||
memory_limit: Option<u64>,
|
||||
num_workers: Option<usize>,
|
||||
document_granularity: String,
|
||||
}
|
||||
|
||||
#[derive(FromPyObject)]
|
||||
@@ -481,6 +486,7 @@ mod tests {
|
||||
block_size = 128
|
||||
memory_limit = 2048
|
||||
num_workers = 7
|
||||
document_granularity = 'row'
|
||||
|
||||
config = FTS()",
|
||||
None,
|
||||
|
||||
+68
-12
@@ -1,6 +1,7 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
@@ -16,8 +17,8 @@ use arrow::pyarrow::FromPyArrow;
|
||||
use arrow::pyarrow::IntoPyArrow;
|
||||
use arrow::pyarrow::ToPyArrow;
|
||||
use lancedb::index::scalar::{
|
||||
BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur,
|
||||
Operator, PhraseQuery,
|
||||
BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery,
|
||||
MultiMatchQuery, Occur, Operator, PhraseQuery,
|
||||
};
|
||||
use lancedb::query::AnalyzePlanDistributedMetrics;
|
||||
use lancedb::query::QueryBase;
|
||||
@@ -76,8 +77,16 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
|
||||
let max_expansions = ob.getattr("max_expansions")?.extract()?;
|
||||
let operator = ob.getattr("operator")?.extract::<String>()?;
|
||||
let prefix_length = ob.getattr("prefix_length")?.extract()?;
|
||||
let document_granularity = ob
|
||||
.getattr("document_granularity")?
|
||||
.extract::<Option<String>>()?
|
||||
.map(|value| {
|
||||
DocumentGranularity::try_from(value.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(Self(
|
||||
let mut query =
|
||||
MatchQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_boost(boost)
|
||||
@@ -86,21 +95,32 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB<FtsQuery> {
|
||||
.with_operator(Operator::try_from(operator.as_str()).map_err(|e| {
|
||||
PyValueError::new_err(format!("Invalid operator: {}", e))
|
||||
})?)
|
||||
.with_prefix_length(prefix_length)
|
||||
.into(),
|
||||
))
|
||||
.with_prefix_length(prefix_length);
|
||||
if let Some(document_granularity) = document_granularity {
|
||||
query = query.with_document_granularity(document_granularity);
|
||||
}
|
||||
Ok(Self(query.into()))
|
||||
}
|
||||
"PhraseQuery" => {
|
||||
let query = ob.getattr("query")?.extract()?;
|
||||
let column = ob.getattr("column")?.extract()?;
|
||||
let slop = ob.getattr("slop")?.extract()?;
|
||||
let document_granularity = ob
|
||||
.getattr("document_granularity")?
|
||||
.extract::<Option<String>>()?
|
||||
.map(|value| {
|
||||
DocumentGranularity::try_from(value.as_str())
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(Self(
|
||||
PhraseQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_slop(slop)
|
||||
.into(),
|
||||
))
|
||||
let mut query = PhraseQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_slop(slop);
|
||||
if let Some(document_granularity) = document_granularity {
|
||||
query = query.with_document_granularity(document_granularity);
|
||||
}
|
||||
Ok(Self(query.into()))
|
||||
}
|
||||
"BoostQuery" => {
|
||||
let positive: Self = ob.getattr("positive")?.extract()?;
|
||||
@@ -167,6 +187,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
|
||||
kwargs.set_item("max_expansions", query.max_expansions)?;
|
||||
kwargs.set_item::<_, &str>("operator", query.operator.into())?;
|
||||
kwargs.set_item("prefix_length", query.prefix_length)?;
|
||||
if let Some(document_granularity) = query.document_granularity {
|
||||
let value = match document_granularity {
|
||||
DocumentGranularity::Row => "row",
|
||||
DocumentGranularity::ListElement => "list_element",
|
||||
};
|
||||
kwargs.set_item("document_granularity", value)?;
|
||||
}
|
||||
namespace
|
||||
.getattr(intern!(py, "MatchQuery"))?
|
||||
.call((query.terms, query.column.unwrap()), Some(&kwargs))
|
||||
@@ -174,6 +201,13 @@ impl<'py> IntoPyObject<'py> for PyLanceDB<FtsQuery> {
|
||||
FtsQuery::Phrase(query) => {
|
||||
let kwargs = PyDict::new(py);
|
||||
kwargs.set_item("slop", query.slop)?;
|
||||
if let Some(document_granularity) = query.document_granularity {
|
||||
let value = match document_granularity {
|
||||
DocumentGranularity::Row => "row",
|
||||
DocumentGranularity::ListElement => "list_element",
|
||||
};
|
||||
kwargs.set_item("document_granularity", value)?;
|
||||
}
|
||||
namespace
|
||||
.getattr(intern!(py, "PhraseQuery"))?
|
||||
.call((query.terms, query.column.unwrap()), Some(&kwargs))
|
||||
@@ -292,6 +326,7 @@ pub struct PyQueryRequest {
|
||||
pub filter: Option<PyQueryFilter>,
|
||||
pub full_text_search: Option<PyLanceDB<FtsQuery>>,
|
||||
pub select: PySelect,
|
||||
pub select_source_columns: Option<HashMap<String, String>>,
|
||||
pub fast_search: Option<bool>,
|
||||
pub with_row_id: Option<bool>,
|
||||
pub use_lsm: Option<bool>,
|
||||
@@ -322,6 +357,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
full_text_search: query_request
|
||||
.full_text_search
|
||||
.map(|fts| PyLanceDB(fts.query)),
|
||||
select_source_columns: PySelect::source_columns(&query_request.select),
|
||||
select: PySelect(query_request.select),
|
||||
fast_search: Some(query_request.fast_search),
|
||||
with_row_id: Some(query_request.with_row_id),
|
||||
@@ -347,6 +383,7 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
offset: vector_query.base.offset,
|
||||
filter: vector_query.base.filter.map(PyQueryFilter),
|
||||
full_text_search: None,
|
||||
select_source_columns: PySelect::source_columns(&vector_query.base.select),
|
||||
select: PySelect(vector_query.base.select),
|
||||
fast_search: Some(vector_query.base.fast_search),
|
||||
with_row_id: Some(vector_query.base.with_row_id),
|
||||
@@ -379,6 +416,25 @@ impl From<AnyQuery> for PyQueryRequest {
|
||||
#[derive(Clone)]
|
||||
pub struct PySelect(Select);
|
||||
|
||||
impl PySelect {
|
||||
fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
|
||||
match select {
|
||||
Select::Expr(pairs) => Some(
|
||||
pairs
|
||||
.iter()
|
||||
.filter_map(|(output, expr)| match expr {
|
||||
lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
|
||||
Some((output.clone(), column.name.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PySelect {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, Self::Target>;
|
||||
|
||||
+6
-2
@@ -780,15 +780,19 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None))]
|
||||
#[pyo3(signature = (data, mode, progress=None, write_parallelism=None, allow_external_blob_outside_bases=false))]
|
||||
pub fn add<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
data: PyScannable,
|
||||
mode: String,
|
||||
progress: Option<Py<PyAny>>,
|
||||
write_parallelism: Option<usize>,
|
||||
allow_external_blob_outside_bases: bool,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let mut op = self_.inner_ref()?.add(data);
|
||||
let mut op = self_
|
||||
.inner_ref()?
|
||||
.add(data)
|
||||
.allow_external_blob_outside_bases(allow_external_blob_outside_bases);
|
||||
if mode == "append" {
|
||||
op = op.mode(AddDataMode::Append);
|
||||
} else if mode == "overwrite" {
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "lancedb"
|
||||
version = "0.38.0-beta.6"
|
||||
version = "0.38.0-beta.12"
|
||||
edition.workspace = true
|
||||
description = "LanceDB: A serverless, low-latency vector database for AI applications"
|
||||
license.workspace = true
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user