Merge lancedb main into the documentation candidate

Brings the candidate up to main at bbc4aebe55, which rebuilds materialized views fully on refresh (#4373) and is version 0.40.0-beta.13, so the documentation and its examples can be checked against that SDK.
This commit is contained in:
Brendan Clement
2026-10-06 20:32:29 -07:00
82 changed files with 2761 additions and 11851 deletions
+1 -1
View File
@@ -1,5 +1,5 @@
[tool.bumpversion]
current_version = "0.40.0-beta.12"
current_version = "0.40.0-beta.13"
parse = """(?x)
(?P<major>0|[1-9]\\d*)\\.
(?P<minor>0|[1-9]\\d*)\\.
+4 -35
View File
@@ -285,7 +285,7 @@ jobs:
runs-on: ubuntu-24.04
strategy:
matrix:
msrv: ["1.91.0"] # This should match up with rust-version in Cargo.toml
msrv: ["1.94.0"] # This should match up with rust-version in Cargo.toml
env:
# Need up-to-date compilers for kernels
CC: clang-18
@@ -308,40 +308,9 @@ jobs:
# unreadable outside their own branch anyway, since GitHub scopes
# caches to the creating ref.
save-if: ${{ github.ref == 'refs/heads/main' }}
- name: Downgrade dependencies
# These packages have newer requirements for MSRV
run: |
cargo update -p aws-sdk-bedrockruntime --precise 1.77.0
cargo update -p aws-sdk-dynamodb --precise 1.68.0
cargo update -p aws-config --precise 1.6.0
cargo update -p aws-sdk-kms --precise 1.63.0
cargo update -p aws-sdk-s3 --precise 1.79.0
cargo update -p aws-sdk-sso --precise 1.62.0
cargo update -p aws-sdk-ssooidc --precise 1.63.0
cargo update -p aws-sdk-sts --precise 1.63.0
# aws-runtime/sigv4/credential-types/types and the aws-smithy-*
# crates bumped their MSRV to 1.91.1 in late 2026; pin to the last
# 1.91.0-compatible versions. The order matters — each downgrade
# only succeeds once everything that still pins it at a higher
# version has itself been downgraded.
cargo update -p aws-runtime --precise 1.5.12
cargo update -p aws-types --precise 1.3.9
cargo update -p aws-sigv4 --precise 1.3.5
cargo update -p aws-credential-types --precise 1.2.8
# aws-smithy-checksums must stay at or above 0.63.13: OpenDAL's S3
# service needs crc-fast ~1.9, and older releases pin it to ~1.3.
cargo update -p aws-smithy-checksums --precise 0.63.13
cargo update -p aws-smithy-runtime --precise 1.9.3
cargo update -p aws-smithy-http --precise 0.62.6
cargo update -p aws-smithy-eventstream --precise 0.60.14
cargo update -p aws-smithy-http-client --precise 1.1.3
cargo update -p aws-smithy-observability --precise 0.1.4
cargo update -p aws-smithy-query --precise 0.60.8
cargo update -p aws-smithy-runtime-api --precise 1.9.3
cargo update -p aws-smithy-async --precise 1.2.7
cargo update -p aws-smithy-types --precise 1.3.6
cargo update -p aws-smithy-xml --precise 0.60.11
cargo update -p home --precise 0.5.9
- name: Downgrade dependencies
# aws-smithy-types 1.6.1+ requires rustc 1.94.1
run: cargo update -p aws-smithy-types --precise 1.5.0
- name: cargo +${{ matrix.msrv }} check
env:
RUSTUP_TOOLCHAIN: ${{ matrix.msrv }}
Generated
+47 -47
View File
@@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arc-swap",
"arrow",
@@ -5217,8 +5217,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5240,7 +5240,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5254,7 +5254,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5263,8 +5263,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrayref",
"crunchy",
@@ -5274,8 +5274,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5312,8 +5312,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow",
"arrow-array",
@@ -5344,8 +5344,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow",
"arrow-array",
@@ -5362,8 +5362,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"proc-macro2",
"quote",
@@ -5372,8 +5372,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5406,8 +5406,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5439,8 +5439,8 @@ dependencies = [
[[package]]
name = "lance-geo"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"datafusion",
"geo-traits",
@@ -5454,8 +5454,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arc-swap",
"arrow",
@@ -5526,8 +5526,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5550,8 +5550,8 @@ dependencies = [
[[package]]
name = "lance-io"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow",
"arrow-array",
@@ -5593,8 +5593,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5608,8 +5608,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow",
"async-trait",
@@ -5623,8 +5623,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arc-swap",
"arrow",
@@ -5678,8 +5678,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5693,8 +5693,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow",
"arrow-array",
@@ -5734,8 +5734,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5748,8 +5748,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "13.1.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v13.1.0-beta.1#e5a3553b1699a7062ff7afe4ee268d596bbf2e0d"
version = "14.0.0-beta.1"
source = "git+https://github.com/lance-format/lance.git?tag=v14.0.0-beta.1#ae2903729f02ec8e519a4e53797a052393f25ee3"
dependencies = [
"frostem",
"icu_segmenter",
@@ -5762,7 +5762,7 @@ dependencies = [
[[package]]
name = "lancedb"
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
dependencies = [
"ahash",
"anyhow",
@@ -5859,7 +5859,7 @@ dependencies = [
[[package]]
name = "lancedb-nodejs"
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5884,7 +5884,7 @@ dependencies = [
[[package]]
name = "lancedb-python"
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
dependencies = [
"arc-swap",
"arrow",
+15 -15
View File
@@ -10,23 +10,23 @@ repository = "https://github.com/lancedb/lancedb"
description = "Serverless, low-latency vector database for AI applications"
keywords = ["lancedb", "lance", "database", "vector", "search"]
categories = ["database-implementations"]
rust-version = "1.91.0"
rust-version = "1.94.0"
[workspace.dependencies]
lance = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=13.1.0-beta.1", default-features = false, "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=13.1.0-beta.1", "tag" = "v13.1.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=14.0.0-beta.1", default-features = false, "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=14.0.0-beta.1", "tag" = "v14.0.0-beta.1", "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
+2 -4
View File
@@ -1,8 +1,6 @@
<a href="https://cloud.lancedb.com" target="_blank">
<img src="https://github.com/user-attachments/assets/92dad0a2-2a37-4ce1-b783-0d1b4f30a00c" alt="LanceDB Cloud Public Beta" width="100%" style="max-width: 100%;">
</a>
<div align="center">
[![Reverie](docs/src/assets/reverie-banner.png)](https://www.reveriesummit.com/?utm_medium=social&utm_source=github&utm_campaign=reverie-2026)
[![LanceDB](docs/src/assets/hero-header.png)](https://lancedb.com)
[![Website](https://img.shields.io/badge/-Website-100000?style=for-the-badge&labelColor=645cfb&color=645cfb)](https://lancedb.com/)
[![Blog](https://img.shields.io/badge/Blog-100000?style=for-the-badge&labelColor=645cfb&color=645cfb)](https://blog.lancedb.com/)
@@ -13,7 +11,7 @@
<img src="docs/src/assets/lancedb.png" alt="LanceDB" width="50%">
# **The Multimodal AI Lakehouse**
# **The Multimodal Lakehouse for AI**
[**How to Install** ](#how-to-install) ✦ [**Detailed Documentation**](https://docs.lancedb.com) ✦ [**Tutorials and Recipes**](https://github.com/lancedb/vectordb-recipes/tree/main) ✦ [**Contributors**](#contributors)
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.7 MiB

After

Width:  |  Height:  |  Size: 1.3 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 722 KiB

+1 -1
View File
@@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`:
<dependency>
<groupId>com.lancedb</groupId>
<artifactId>lancedb-core</artifactId>
<version>0.40.0-beta.12</version>
<version>0.40.0-beta.13</version>
</dependency>
```
+15 -4
View File
@@ -89,19 +89,30 @@ Drop an empty database. The server rejects nonempty databases.
### listDatabases()
```ts
listDatabases(options): Promise<ListDatabasesResponse>
listDatabases(options): DatabaseNames
```
List one page of databases; pass pageToken from a response for the next page.
Iterate lazily over all database names using `for await...of`.
pageToken resumes from a saved token; omitted starts at the beginning.
pageLimit limits each REST response, not the total; omitted uses the server default.
Request errors are raised during iteration and terminate the iterator.
#### Parameters
* **options** = `{}`
* **options.limit?**: `number`
* **options.pageLimit?**: `number`
* **options.pageToken?**: `string`
#### Returns
`Promise`&lt;[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)&gt;
[`DatabaseNames`](DatabaseNames.md)
#### Example
```ts
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
console.log(name);
}
```
-4
View File
@@ -183,10 +183,6 @@ Define a materialized view named `name` over the table `source`.
The view is populated before creation returns. Set `withNoData` to create
only its definition and empty backing table. The view is a normal table:
it can be queried, indexed and searched, and it appears in `tableNames`.
The source table must have stable row ids (create it with
the `newTableEnableStableRowIds` storage option); they keep the view's
provenance valid across source compactions and cannot be enabled after
a table exists.
#### Parameters
+79
View File
@@ -0,0 +1,79 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / DatabaseNames
# Class: DatabaseNames
A lazy async iterator of database names with pagination state.
Returned by [Catalog.listDatabases](Catalog.md#listdatabases).
## Implements
- `AsyncIterableIterator`&lt;`string`&gt;
## Methods
### \[asyncIterator\]()
```ts
asyncIterator: AsyncIterableIterator<string>
```
#### Returns
`AsyncIterableIterator`&lt;`string`&gt;
#### Implementation of
`AsyncIterableIterator.[asyncIterator]`
***
### next()
```ts
next(): Promise<IteratorResult<string, any>>
```
Fetch the next name, requesting another page only when needed.
#### Returns
`Promise`&lt;`IteratorResult`&lt;`string`, `any`&gt;&gt;
#### Implementation of
`AsyncIterableIterator.next`
***
### numPageResults()
```ts
numPageResults(): number
```
Number of names available without another REST request.
#### Returns
`number`
***
### pageToken()
```ts
pageToken(): undefined | string
```
Token for the next REST request. Initially this is the supplied starting token.
Returns undefined after the final page is fetched, even if names remain cached.
Drain the cache before saving a token to avoid skipping names when resuming.
A failed request terminates iteration and retains its token for a new iterator.
#### Returns
`undefined` \| `string`
+3 -10
View File
@@ -65,21 +65,14 @@ refresh(options?): Promise<RefreshMaterializedViewResult>
Recompute the view from its source.
The refresh is incremental when the source's changes can be reconciled
into the view -- rows added, changed or removed since the last one --
and otherwise rebuilds. `full` forces a rebuild; `sourceVersion`
refreshes to that source version instead of the latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser throws
rather than writing them a second time.
Every refresh rebuilds the complete view. For local views,
`sourceVersion` refreshes to that source version instead of the latest;
remote SQL refreshes do not support source-version pinning.
#### Parameters
* **options?**
* **options.full?**: `boolean`
* **options.sourceVersion?**: `number`
#### Returns
+5 -6
View File
@@ -662,10 +662,9 @@ Read the [LsmWriteSpec](../interfaces/LsmWriteSpec.md) currently installed on th
Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
spec has been set, or it was removed with [Table#unsetLsmWriteSpec](Table.md#unsetlsmwritespec)).
The returned spec mirrors what was passed to
[Table#setLsmWriteSpec](Table.md#setlsmwritespec), except that `maintainedIndexes` always
reports the concrete list resolved when the spec was set — `undefined`
never round-trips.
The spec is the one installed, including its maintained-index selection:
an absent `maintainedIndexes` for every index the table has, an empty
array for none.
#### Returns
@@ -1107,8 +1106,8 @@ All variants require the table to have an unenforced primary key
([Table#setUnenforcedPrimaryKey](Table.md#setunenforcedprimarykey)); bucket sharding additionally
requires it to be the single column being bucketed.
Omitting `maintainedIndexes` maintains every index on the table, resolved
here, failing if one cannot be maintained — name them to install anyway.
Omitting `maintainedIndexes` maintains every index the table has,
including ones created later, and skips a kind the MemWAL cannot maintain.
Naming them pins an exact set, and a still-building index is rejected
rather than quietly omitted.
+3 -1
View File
@@ -28,5 +28,7 @@ headers; opened database connections inherit authentication and client options.
```ts
const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
const db = await catalog.createDatabase("analytics", { existOk: true });
const page = await catalog.listDatabases({ limit: 20 });
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
console.log(name);
}
```
+1 -1
View File
@@ -27,6 +27,7 @@
- [Branches](classes/Branches.md)
- [Catalog](classes/Catalog.md)
- [Connection](classes/Connection.md)
- [DatabaseNames](classes/DatabaseNames.md)
- [HeaderProvider](classes/HeaderProvider.md)
- [Index](classes/Index.md)
- [Job](classes/Job.md)
@@ -108,7 +109,6 @@
- [JobEventsOptions](interfaces/JobEventsOptions.md)
- [JobFailureInfo](interfaces/JobFailureInfo.md)
- [JobInfo](interfaces/JobInfo.md)
- [ListDatabasesResponse](interfaces/ListDatabasesResponse.md)
- [ListNamespacesOptions](interfaces/ListNamespacesOptions.md)
- [ListNamespacesResponse](interfaces/ListNamespacesResponse.md)
- [ListTablesOptions](interfaces/ListTablesOptions.md)
@@ -1,23 +0,0 @@
[**@lancedb/lancedb**](../README.md) • **Docs**
***
[@lancedb/lancedb](../globals.md) / ListDatabasesResponse
# Interface: ListDatabasesResponse
## Properties
### databases
```ts
databases: string[];
```
***
### pageToken?
```ts
optional pageToken: string;
```
+2 -3
View File
@@ -34,9 +34,8 @@ Bucket and identity variants: the sharding column.
optional maintainedIndexes: string[];
```
Indexes the MemWAL keeps up to date. Omit to maintain every supported
index, resolved on install — a snapshot, so indexes created later are not
maintained. Pass `[]` for none.
Indexes the MemWAL keeps up to date. Omit for every index the table has,
including ones created later. Pass `[]` for none.
***
@@ -14,7 +14,7 @@
mode: string;
```
How the view was brought up to date: "rebuild", "incremental" or "no_op".
How the view was brought up to date: "rebuild".
***
+4 -28
View File
@@ -37,7 +37,8 @@ are ordinary connections. Dropping a database requires it to be empty.
::: lancedb.catalog.Catalog
::: lancedb.catalog.ListDatabasesResponse
::: lancedb.catalog.DatabaseNames
## Remote SQL
@@ -198,33 +199,6 @@ listing a storage directory.
::: lancedb.materialized_view.MaterializedViewDefinition
::: lancedb.vector_dedup
::: lancedb.materialized_view.VectorDedupSource
For an existing indexed source, declare and materialize a dedup result without
writing SQL. Creation captures the source version once and waits for the complete
result; remote connections use the service's MV jobs.
```python
from lancedb import vector_dedup
view = db.create_materialized_view(
"images_clean",
vector_dedup("images", column="phash", distance_threshold=4),
)
cleaned = view.table
```
To submit without waiting, call `db.create_materialized_view_async(...)` with the
same source expression and wait on the returned job. To declare only, pass
`with_no_data=True`, then use `view.refresh_async().wait()` when ready. Pass
`dataset_version=...` to `vector_dedup` for an explicit snapshot. The default
policy retains direct representatives (A-B-C without A-C keeps A and C), preserves
the source, and materializes its original columns. Only within-segment/partition
pairs are covered, using the index's distance representation.
## Views
::: lancedb.view.ViewDescription
@@ -440,6 +414,8 @@ latency on your workload before opting in.
::: lancedb.catalog.AsyncCatalog
::: lancedb.catalog.AsyncDatabaseNames
Connections represent a connection to a LanceDb database and
can be used to create, list, or open tables.
+1 -1
View File
@@ -67,7 +67,7 @@ reader usually expects:
| Value | Meaning |
| --- | --- |
| unset (null) | Maintain **every** index the MemWAL can, resolved on install |
| unset (null) | Maintain **every** index the table has, including ones created later |
| `Collections.emptyList()` | Maintain **none** |
| `Arrays.asList("id_idx")` | Maintain exactly those |
+1 -1
View File
@@ -8,7 +8,7 @@
<parent>
<groupId>com.lancedb</groupId>
<artifactId>lancedb-parent</artifactId>
<version>0.40.0-beta.12</version>
<version>0.40.0-beta.13</version>
<relativePath>../pom.xml</relativePath>
</parent>
@@ -117,9 +117,13 @@ public class LanceDbTableLsm {
/**
* Read the {@link LsmWriteSpec} currently installed on this table.
*
* <p>Empty when the LSM write path is not enabled. The returned spec mirrors what was installed,
* except that {@link LsmWriteSpec#maintainedIndexes()} always reports the concrete list resolved
* when the spec was set — a null selection never round-trips.
* <p>Empty when the LSM write path is not enabled at all — no spec has been set, or one was
* removed with {@link #unsetLsmWriteSpec()}. That is a different answer from a present spec
* whose {@link LsmWriteSpec#maintainedIndexes()} is null, which selects indexes automatically.
*
* <p>The spec read back is the one that was installed, selection included: null maintains every
* supported index the table has now or gains later, an empty list maintains none, and a
* non-empty list maintains exactly those. All three round-trip.
*/
public Optional<LsmWriteSpec> getLsmWriteSpec() {
JsonNode response = client.post(route("get_lsm_write_spec"), null);
@@ -36,6 +36,8 @@ import java.util.Map;
*/
public class LsmWriteSpec {
private static final String CONTEXT = "get_lsm_write_spec response";
/** How writes are routed to MemWAL shards. */
public enum Sharding {
/** Hash-bucket writes by a scalar column. */
@@ -124,9 +126,8 @@ public class LsmWriteSpec {
/**
* Set the indexes the MemWAL keeps up to date as rows are appended.
*
* <p>Pass {@code null} — the default for a fresh spec — to maintain every index the MemWAL can,
* resolved when the spec is installed. That is a snapshot: indexes created later are not
* maintained until the spec is unset and set again. Pass an empty list to maintain none.
* <p>Pass {@code null} — the default for a fresh spec — to maintain every index the table has,
* including ones created later. Pass an empty list to maintain none.
*
* <p>Note that {@code null} and the empty list mean opposite things here.
*/
@@ -176,8 +177,10 @@ public class LsmWriteSpec {
}
/**
* The indexes the MemWAL maintains, or null to have the server resolve every maintainable index
* on install. An empty list means none.
* Which indexes the MemWAL maintains.
*
* <p>Null selects automatically: every supported index the table has now or gains later. An empty
* list maintains none. A non-empty list maintains exactly those.
*/
public List<String> maintainedIndexes() {
return maintainedIndexes == null ? null : Collections.unmodifiableList(maintainedIndexes);
@@ -210,8 +213,8 @@ public class LsmWriteSpec {
/**
* Rebuild a spec from a {@code get_lsm_write_spec} response body.
*
* <p>The server always reports a concrete maintained-index list, so a null selection never
* round-trips.
* <p>A null {@code maintained_indexes} selects every index the table has; an empty list selects
* none.
*/
static LsmWriteSpec fromJson(JsonNode node) {
JsonNode shardingNode = node.get("sharding");
@@ -224,10 +227,19 @@ public class LsmWriteSpec {
Integer numBuckets =
shardingNode.hasNonNull("num_buckets") ? shardingNode.get("num_buckets").asInt() : null;
List<String> maintainedIndexes = new ArrayList<String>();
JsonNode indexesNode = node.get("maintained_indexes");
if (indexesNode != null && indexesNode.isArray()) {
// Absent or null is the automatic selection; anything else has to be an
// array of names. Reading a string or an object as automatic would turn a
// malformed response into "maintain every index", which is a policy, not a
// parse failure.
JsonNode indexesNode = JsonFields.optionalArray(node, "maintained_indexes", CONTEXT);
List<String> maintainedIndexes = null;
if (indexesNode != null) {
maintainedIndexes = new ArrayList<String>();
for (JsonNode index : indexesNode) {
if (!index.isTextual()) {
throw new IllegalStateException(
CONTEXT + ": maintained_indexes must hold index names, got " + index);
}
maintainedIndexes.add(index.asText());
}
}
@@ -215,6 +215,61 @@ public class LanceDbTableLsmTest {
assertEquals("true", spec.get().writerConfigDefaults().get("durable_write"));
}
/** A null selection, every index the table has, reads back as null rather than an empty list. */
@Test
public void testGetLsmWriteSpecRoundTripsEveryIndexSelection() {
enqueue(
"get_lsm_write_spec",
200,
"{\"lsm_write_spec\":{\"sharding\":{\"mode\":\"unsharded\"},"
+ "\"maintained_indexes\":null,\"writer_config_defaults\":{}}}");
Optional<LsmWriteSpec> spec = lsm.getLsmWriteSpec();
assertTrue(spec.isPresent());
assertNull(spec.get().maintainedIndexes());
}
@Test
public void testGetLsmWriteSpecRoundTripsNoIndexSelection() {
enqueue(
"get_lsm_write_spec",
200,
"{\"lsm_write_spec\":{\"sharding\":{\"mode\":\"unsharded\"},"
+ "\"maintained_indexes\":[],\"writer_config_defaults\":{}}}");
Optional<LsmWriteSpec> spec = lsm.getLsmWriteSpec();
assertTrue(spec.isPresent());
assertEquals(Collections.emptyList(), spec.get().maintainedIndexes());
}
/**
* Null is the automatic selection, so anything that is neither null nor an array has to be
* rejected: reading a string or an object as null would turn a malformed response into a policy.
*/
@Test
public void testGetLsmWriteSpecRejectsANonArrayMaintainedIndexes() {
enqueue(
"get_lsm_write_spec",
200,
"{\"lsm_write_spec\":{\"sharding\":{\"mode\":\"unsharded\"},"
+ "\"maintained_indexes\":\"id_idx\",\"writer_config_defaults\":{}}}");
assertThrows(IllegalStateException.class, () -> lsm.getLsmWriteSpec());
}
@Test
public void testGetLsmWriteSpecRejectsNonStringIndexNames() {
enqueue(
"get_lsm_write_spec",
200,
"{\"lsm_write_spec\":{\"sharding\":{\"mode\":\"unsharded\"},"
+ "\"maintained_indexes\":[7],\"writer_config_defaults\":{}}}");
assertThrows(IllegalStateException.class, () -> lsm.getLsmWriteSpec());
}
@Test
public void testGetLsmWriteSpecAbsent() {
enqueue("get_lsm_write_spec", 200, "{\"lsm_write_spec\":null}");
+2 -2
View File
@@ -6,7 +6,7 @@
<groupId>com.lancedb</groupId>
<artifactId>lancedb-parent</artifactId>
<version>0.40.0-beta.12</version>
<version>0.40.0-beta.13</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>13.1.0-beta.1</lance-core.version>
<lance-core.version>14.0.0-beta.1</lance-core.version>
<spotless.skip>false</spotless.skip>
<spotless.version>2.30.0</spotless.version>
<spotless.java.googlejavaformat.version>1.7</spotless.java.googlejavaformat.version>
+1 -1
View File
@@ -1,7 +1,7 @@
[package]
name = "lancedb-nodejs"
edition.workspace = true
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
publish = false
license.workspace = true
description.workspace = true
+136 -5
View File
@@ -80,9 +80,12 @@ describe("remote catalog", () => {
const second = await catalog.connectDatabase("other");
expect(await second.tableNames()).toEqual([]);
expect(await first.tableNames()).toEqual([]);
expect(
await catalog.listDatabases({ limit: 1, pageToken: "a/b" }),
).toEqual({ databases: ["team/search"], pageToken: "next" });
const names = catalog.listDatabases({ pageLimit: 1, pageToken: "a/b" });
expect(await names.next()).toEqual({
done: false,
value: "team/search",
});
expect(names.pageToken()).toBe("next");
await catalog.dropDatabase("team/search", { ignoreMissing: true });
expect(requests[0].url).toBe("/v1/namespace/team%2Fsearch/create");
expect(requests[0].body).toEqual({ mode: "ExistOk" });
@@ -129,11 +132,139 @@ describe("remote catalog", () => {
);
});
it("lazily traverses pages, including empty pages, and exposes cached state", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a", "b"], page_token: "empty" }],
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: [], page_token: "last" }],
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["c", "d"], page_token: "" }],
],
async (catalog, requests) => {
const names = catalog.listDatabases({ pageLimit: 2 });
expect(names[Symbol.asyncIterator]()).toBe(names);
expect(requests).toHaveLength(0);
expect(names.numPageResults()).toBe(0);
expect(names.pageToken()).toBeUndefined();
expect(await names.next()).toEqual({ done: false, value: "a" });
expect(names.numPageResults()).toBe(1);
expect(names.pageToken()).toBe("empty");
expect(await names.next()).toEqual({ done: false, value: "b" });
expect(requests).toHaveLength(1);
expect(await names.next()).toEqual({ done: false, value: "c" });
expect(requests).toHaveLength(3);
expect(names.pageToken()).toBeUndefined();
expect(names.numPageResults()).toBe(1);
const remaining = [];
for await (const name of names) remaining.push(name);
expect(remaining).toEqual(["d"]);
expect(names.numPageResults()).toBe(0);
expect(await names.next()).toEqual({ done: true, value: undefined });
expect(requests.map((request) => request.url)).toEqual([
"/v1/namespace/$/list?limit=2",
"/v1/namespace/$/list?limit=2&page_token=empty",
"/v1/namespace/$/list?limit=2&page_token=last",
]);
},
);
});
it("resumes from a saved token after draining the cache", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a", "b"], page_token: "next/token" }],
[200, { namespaces: ["c"] }],
],
async (catalog, requests) => {
const names = catalog.listDatabases({
pageToken: "start/token",
pageLimit: 2,
});
expect(names.pageToken()).toBe("start/token");
expect(await names.next()).toEqual({ done: false, value: "a" });
while (names.numPageResults() > 0) await names.next();
const resumed = catalog.listDatabases({
pageToken: names.pageToken(),
pageLimit: 2,
});
const results = [];
for await (const name of resumed) results.push(name);
expect(results).toEqual(["c"]);
expect(requests.map((request) => request.url)).toEqual([
"/v1/namespace/$/list?limit=2&page_token=start%2Ftoken",
"/v1/namespace/$/list?limit=2&page_token=next%2Ftoken",
]);
},
);
});
it("terminates on errors and retains the failed request token", async () => {
await withCatalog(
[
// biome-ignore lint/style/useNamingConvention: server wire format
[200, { namespaces: ["a"], page_token: "retry" }],
[400, {}],
[200, { namespaces: ["b"] }],
],
async (catalog, requests) => {
const names = catalog.listDatabases();
await names.next();
await expect(names.next()).rejects.toThrow();
expect(names.pageToken()).toBe("retry");
expect(names.numPageResults()).toBe(0);
expect(await names.next()).toEqual({ done: true, value: undefined });
expect(requests).toHaveLength(2);
const resumed = catalog.listDatabases({ pageToken: names.pageToken() });
expect(await resumed.next()).toEqual({ done: false, value: "b" });
expect(await resumed.next()).toEqual({ done: true, value: undefined });
},
);
});
it("handles an empty listing and concurrent advances", async () => {
await withCatalog(
[
[200, { namespaces: [] }],
[200, { namespaces: ["a", "b"] }],
],
async (catalog, requests) => {
const empty = catalog.listDatabases();
expect(await empty.next()).toEqual({ done: true, value: undefined });
expect(await empty.next()).toEqual({ done: true, value: undefined });
const names = catalog.listDatabases();
const pending = [names.next(), names.next(), names.next()];
expect(names.numPageResults()).toBe(0);
expect(names.pageToken()).toBeUndefined();
const results = await Promise.all(pending);
expect(
results
.filter((result) => !result.done)
.map((result) => result.value)
.sort(),
).toEqual(["a", "b"]);
expect(results.filter((result) => result.done)).toHaveLength(1);
expect(requests).toHaveLength(2);
},
);
});
it("validates endpoints and pagination", async () => {
await expect(connectCatalog("/tmp/catalog")).rejects.toThrow();
const catalog = await connectCatalog("http://127.0.0.1:1");
for (const limit of [0, -1, 1.5, 2147483648]) {
await expect(catalog.listDatabases({ limit })).rejects.toThrow("limit");
for (const limit of [
0,
-1,
1.5,
2147483648,
Number.NaN,
Number.POSITIVE_INFINITY,
]) {
expect(() => catalog.listDatabases({ pageLimit: limit })).toThrow(
"limit",
);
}
await expect(catalog.connectDatabase("a$b")).rejects.toThrow(
"Invalid database name",
+14 -31
View File
@@ -16,15 +16,11 @@ describe("materialized views", () => {
beforeEach(async () => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
db = await connect(tmpDir.name);
await db.createTable(
"people",
[
{ name: "ada", age: 36 },
{ name: "kid", age: 7 },
{ name: "grace", age: 85 },
],
{ storageOptions: { newTableEnableStableRowIds: "true" } },
);
await db.createTable("people", [
{ name: "ada", age: 36 },
{ name: "kid", age: 7 },
{ name: "grace", age: 85 },
]);
});
afterEach(() => tmpDir.removeCallback());
@@ -52,19 +48,9 @@ describe("materialized views", () => {
"SELECT * FROM people",
);
for (const [format, fn] of [
[3, "vector_duplicate_pairs"],
[4, "vector_dedup"],
]) {
const native = `SELECT * FROM ${fn}('images', 3, 'phash', 4)`;
expect(read(JSON.stringify({ format, query: native })).query).toBe(
native,
);
}
// A newer writer's layout is reported, never guessed at.
for (const newer of [
`{"format":5,"query":${JSON.stringify(query)}}`,
`{"format":3,"query":${JSON.stringify(query)}}`,
'{"kind":"select_v3","source_table":"people"}',
]) {
expect(() => read(newer)).toThrow(/cannot refresh/);
@@ -106,7 +92,7 @@ describe("materialized views", () => {
);
});
it("refreshes incrementally after an append", async () => {
it("fully refreshes after an append", async () => {
const view = await db.createMaterializedView("copy", "people", {
withNoData: true,
});
@@ -115,11 +101,11 @@ describe("materialized views", () => {
const people = await db.openTable("people");
await people.add([{ name: "alan", age: 41 }]);
const result = await view.refresh();
expect(result.mode).toBe("incremental");
expect(Number(result.rowsWritten)).toBe(1);
expect(result.mode).toBe("rebuild");
expect(Number(result.rowsWritten)).toBe(4);
expect(await view.table().countRows()).toBe(4);
expect((await view.refresh()).mode).toBe("no_op");
expect((await view.refresh()).mode).toBe("rebuild");
});
it("lists views and rejects non-views", async () => {
@@ -172,9 +158,7 @@ describe("materialized views", () => {
});
it("quotes bare select names", async () => {
await db.createTable("odd_names", [{ "order item": "widget" }], {
storageOptions: { newTableEnableStableRowIds: "true" },
});
await db.createTable("odd_names", [{ "order item": "widget" }]);
const view = await db.createMaterializedView("quoted", "odd_names", {
select: ["order item"],
withNoData: true,
@@ -183,10 +167,9 @@ describe("materialized views", () => {
expect(Number(result.rowsWritten)).toBe(1);
});
it("requires stable row ids on the source", async () => {
it("does not require stable row ids on the source", async () => {
await db.createTable("plain", [{ x: 1 }]);
await expect(db.createMaterializedView("v", "plain")).rejects.toThrow(
"stable row ids",
);
const view = await db.createMaterializedView("v", "plain");
expect(await view.table().countRows()).toBe(1);
});
});
+62 -12
View File
@@ -5,9 +5,9 @@ import { Connection, LocalConnection } from "./connection";
import { HeaderProvider } from "./header";
import {
JsHeaderProvider,
ListDatabasesResponse,
Catalog as NativeCatalog,
CatalogOptions as NativeCatalogOptions,
DatabaseNames as NativeDatabaseNames,
} from "./native.js";
import { OAuthConfig } from "./oauth";
@@ -21,7 +21,41 @@ export interface CatalogOptions
| (() => Record<string, string> | Promise<Record<string, string>>);
}
export type { ListDatabasesResponse } from "./native.js";
/**
* A lazy async iterator of database names with pagination state.
* Returned by {@link Catalog.listDatabases}.
*/
export class DatabaseNames implements AsyncIterableIterator<string> {
/** @hidden */
constructor(private readonly inner: NativeDatabaseNames) {}
[Symbol.asyncIterator](): AsyncIterableIterator<string> {
return this;
}
/** Fetch the next name, requesting another page only when needed. */
async next(): Promise<IteratorResult<string>> {
const value = await this.inner.next();
return value == null
? { done: true, value: undefined }
: { done: false, value };
}
/** Number of names available without another REST request. */
numPageResults(): number {
return this.inner.numPageResults();
}
/**
* Token for the next REST request. Initially this is the supplied starting token.
* Returns undefined after the final page is fetched, even if names remain cached.
* Drain the cache before saving a token to avoid skipping names when resuming.
* A failed request terminates iteration and retains its token for a new iterator.
*/
pageToken(): string | undefined {
return this.inner.pageToken() ?? undefined;
}
}
/** A remote catalog manages databases through the server's root namespace. */
export class Catalog {
@@ -56,21 +90,35 @@ export class Catalog {
await this.inner.dropDatabase(name, options.ignoreMissing);
}
/** List one page of databases; pass pageToken from a response for the next page. */
async listDatabases(
options: { limit?: number; pageToken?: string } = {},
): Promise<ListDatabasesResponse> {
/**
* Iterate lazily over all database names using `for await...of`.
* pageToken resumes from a saved token; omitted starts at the beginning.
* pageLimit limits each REST response, not the total; omitted uses the server default.
* Request errors are raised during iteration and terminate the iterator.
*
* @example
* ```ts
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
* console.log(name);
* }
* ```
*/
listDatabases(
options: { pageLimit?: number; pageToken?: string } = {},
): DatabaseNames {
if (
options.limit !== undefined &&
(!Number.isInteger(options.limit) ||
options.limit <= 0 ||
options.limit > 2147483647)
options.pageLimit !== undefined &&
(!Number.isInteger(options.pageLimit) ||
options.pageLimit <= 0 ||
options.pageLimit > 2147483647)
) {
throw new Error(
"Database list limit must be an integer between 1 and 2147483647",
);
}
return this.inner.listDatabases(options.limit, options.pageToken);
return new DatabaseNames(
this.inner.listDatabases(options.pageToken, options.pageLimit),
);
}
}
@@ -82,7 +130,9 @@ export class Catalog {
* ```ts
* const catalog = await connectCatalog("https://my-server.example", { apiKey: "secret" });
* const db = await catalog.createDatabase("analytics", { existOk: true });
* const page = await catalog.listDatabases({ limit: 20 });
* for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
* console.log(name);
* }
* ```
*/
export async function connectCatalog(
-4
View File
@@ -325,10 +325,6 @@ export abstract class Connection {
* The view is populated before creation returns. Set `withNoData` to create
* only its definition and empty backing table. The view is a normal table:
* it can be queried, indexed and searched, and it appears in `tableNames`.
* The source table must have stable row ids (create it with
* the `newTableEnableStableRowIds` storage option); they keep the view's
* provenance valid across source compactions and cannot be enabled after
* a table exists.
*/
abstract createMaterializedView(
name: string,
+1 -1
View File
@@ -646,6 +646,6 @@ export async function connectNamespace(
export {
Catalog,
CatalogOptions,
ListDatabasesResponse,
DatabaseNames,
connectCatalog,
} from "./catalog";
+6 -16
View File
@@ -9,10 +9,9 @@ export const DEFINITION_META_KEY = "mv.definition";
/**
* The newest stored layout this version reads: `{"format": N, "query": "<SQL>"}`,
* format 2 being a query with `GROUP BY`, format 3 native pairs and format 4
* native dedup results.
* format 2 being a query with `GROUP BY`.
*/
export const DEFINITION_FORMAT = 4;
export const DEFINITION_FORMAT = 2;
/**
* The query that defines a materialized view, as stored:
@@ -188,23 +187,14 @@ export class MaterializedView {
/**
* Recompute the view from its source.
*
* The refresh is incremental when the source's changes can be reconciled
* into the view -- rows added, changed or removed since the last one --
* and otherwise rebuilds. `full` forces a rebuild; `sourceVersion`
* refreshes to that source version instead of the latest.
*
* Concurrent refreshes of one view do not duplicate its rows. Two that
* plan the same source rows conflict on commit, and the loser throws
* rather than writing them a second time.
* Every refresh rebuilds the complete view. For local views,
* `sourceVersion` refreshes to that source version instead of the latest;
* remote SQL refreshes do not support source-version pinning.
*/
async refresh(options?: {
full?: boolean;
sourceVersion?: number;
}): Promise<RefreshMaterializedViewResult> {
validateNonNegativeInteger(options?.sourceVersion, "sourceVersion");
return await this.inner.refreshMaterializedView(
options?.full,
options?.sourceVersion,
);
return await this.inner.refreshMaterializedView(options?.sourceVersion);
}
}
+8 -12
View File
@@ -225,9 +225,8 @@ export interface LsmWriteSpec {
/** Bucket variant: the number of buckets, in `[1, 1024]`. */
numBuckets?: number;
/**
* Indexes the MemWAL keeps up to date. Omit to maintain every supported
* index, resolved on install — a snapshot, so indexes created later are not
* maintained. Pass `[]` for none.
* Indexes the MemWAL keeps up to date. Omit for every index the table has,
* including ones created later. Pass `[]` for none.
*/
maintainedIndexes?: string[];
/** Default `ShardWriter` configuration recorded in the MemWAL index. */
@@ -685,7 +684,6 @@ export abstract class Table {
* @ignore
*/
abstract refreshMaterializedView(
full?: boolean,
sourceVersion?: number,
): Promise<RefreshMaterializedViewResult>;
@@ -767,8 +765,8 @@ export abstract class Table {
* ({@link Table#setUnenforcedPrimaryKey}); bucket sharding additionally
* requires it to be the single column being bucketed.
*
* Omitting `maintainedIndexes` maintains every index on the table, resolved
* here, failing if one cannot be maintained — name them to install anyway.
* Omitting `maintainedIndexes` maintains every index the table has,
* including ones created later, and skips a kind the MemWAL cannot maintain.
* Naming them pins an exact set, and a still-building index is rejected
* rather than quietly omitted.
* @param {LsmWriteSpec} spec The sharding spec to install.
@@ -798,10 +796,9 @@ export abstract class Table {
*
* Resolves to `undefined` when the MemWAL LSM write path is not enabled (no
* spec has been set, or it was removed with {@link Table#unsetLsmWriteSpec}).
* The returned spec mirrors what was passed to
* {@link Table#setLsmWriteSpec}, except that `maintainedIndexes` always
* reports the concrete list resolved when the spec was set — `undefined`
* never round-trips.
* The spec is the one installed, including its maintained-index selection:
* an absent `maintainedIndexes` for every index the table has, an empty
* array for none.
* @returns {Promise<LsmWriteSpec | undefined>}
*/
abstract getLsmWriteSpec(): Promise<LsmWriteSpec | undefined>;
@@ -1420,10 +1417,9 @@ export class LocalTable extends Table {
}
async refreshMaterializedView(
full?: boolean,
sourceVersion?: number,
): Promise<RefreshMaterializedViewResult> {
return await this.inner.refreshMaterializedView(full, sourceVersion);
return await this.inner.refreshMaterializedView(sourceVersion);
}
async materializedViewDefinition(): Promise<string> {
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-darwin-arm64",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["darwin"],
"cpu": ["arm64"],
"main": "lancedb.darwin-arm64.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-gnu",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-arm64-musl",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["linux"],
"cpu": ["arm64"],
"main": "lancedb.linux-arm64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-gnu",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-gnu.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-linux-x64-musl",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["linux"],
"cpu": ["x64"],
"main": "lancedb.linux-x64-musl.node",
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-arm64-msvc",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": [
"win32"
],
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@lancedb/lancedb-win32-x64-msvc",
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"os": ["win32"],
"cpu": ["x64"],
"main": "lancedb.win32-x64-msvc.node",
+1 -1
View File
@@ -11,7 +11,7 @@
"ann"
],
"private": false,
"version": "0.40.0-beta.12",
"version": "0.40.0-beta.13",
"main": "dist/index.js",
"exports": {
".": "./dist/index.js",
+42 -19
View File
@@ -1,12 +1,12 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use lancedb::catalog::{
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
use futures::{StreamExt, future::poll_fn};
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
use napi::bindgen_prelude::*;
use napi_derive::napi;
@@ -25,10 +25,37 @@ pub struct CatalogOptions {
pub oauth_config: Option<OAuthConfig>,
}
#[napi(object)]
pub struct ListDatabasesResponse {
pub databases: Vec<String>,
pub page_token: Option<String>,
/// A lazy iterator over database names with inspectable pagination state.
#[napi]
pub struct DatabaseNames {
inner: Mutex<lancedb::catalog::DatabaseNames>,
next_lock: futures::lock::Mutex<()>,
}
#[napi]
impl DatabaseNames {
/// Number of names cached without another REST request.
#[napi]
pub fn num_page_results(&self) -> u32 {
self.inner.lock().unwrap().num_page_results() as u32
}
/// Continuation token for the next REST request.
#[napi]
pub fn page_token(&self) -> Option<String> {
self.inner.lock().unwrap().page_token().map(str::to_owned)
}
/// Fetch the next name, returning None when exhausted.
#[napi]
pub async fn next(&self) -> Result<Option<String>> {
// Serialize advances but allow synchronous inspection between polls.
let _guard = self.next_lock.lock().await;
poll_fn(|cx| self.inner.lock().unwrap().poll_next_unpin(cx))
.await
.transpose()
.default_error()
}
}
#[napi]
@@ -106,18 +133,14 @@ impl Catalog {
.default_error()
}
#[napi]
pub async fn list_databases(
pub fn list_databases(
&self,
limit: Option<u32>,
page_token: Option<String>,
) -> Result<ListDatabasesResponse> {
let mut request = ListDatabasesRequest::default();
request.limit = limit;
request.page_token = page_token;
let response = self.inner.list_databases(request).await.default_error()?;
Ok(ListDatabasesResponse {
databases: response.databases,
page_token: response.page_token,
})
page_limit: Option<u32>,
) -> DatabaseNames {
DatabaseNames {
inner: Mutex::new(self.inner.list_databases(page_token, page_limit)),
next_lock: futures::lock::Mutex::new(()),
}
}
}
+5 -10
View File
@@ -469,13 +469,12 @@ impl Table {
#[napi(catch_unwind)]
pub async fn refresh_materialized_view(
&self,
full: Option<bool>,
source_version: Option<i64>,
) -> napi::Result<RefreshMaterializedViewResult> {
let view = lancedb::MaterializedView::from_table(self.inner_ref()?.clone())
.await
.default_error()?;
let mut builder = view.refresh().full(full.unwrap_or(false));
let mut builder = view.refresh();
if let Some(version) = source_version {
let version = u64::try_from(version).map_err(|_| {
napi::Error::from_reason("sourceVersion must be a non-negative integer")
@@ -492,11 +491,9 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.default_error()?;
view.definition().to_json().map_err(|err| {
napi::Error::from_reason(format!(
"failed to serialize materialized-view definition: {err}"
))
})
Ok(lancedb::materialized_view::definition_metadata_from_sql(
view.definition_sql(),
))
}
#[napi(catch_unwind)]
@@ -1608,7 +1605,7 @@ impl From<lancedb::function::FunctionErrors> for FunctionErrors {
#[napi(object)]
pub struct RefreshMaterializedViewResult {
/// How the view was brought up to date: "rebuild", "incremental" or "no_op".
/// How the view was brought up to date: "rebuild".
pub mode: String,
pub rows_written: i64,
pub source_version: i64,
@@ -1619,8 +1616,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
fn from(value: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match value.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
+2 -2
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb-python"
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
publish = false
edition.workspace = true
description = "Python bindings for LanceDB"
@@ -8,7 +8,7 @@ license.workspace = true
repository.workspace = true
keywords.workspace = true
categories.workspace = true
rust-version = "1.91.0"
rust-version = "1.94.0"
[lib]
name = "_lancedb"
-3
View File
@@ -46,7 +46,6 @@ from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewDefinition,
vector_dedup as vector_dedup,
)
from .view import ViewDescription as ViewDescription
from .table import AsyncTable, Table
@@ -62,7 +61,6 @@ from .namespace import (
from .catalog import (
AsyncCatalog,
Catalog,
ListDatabasesResponse,
connect_catalog,
connect_catalog_async,
)
@@ -578,7 +576,6 @@ async def connect_async(
__all__ = [
"Catalog",
"AsyncCatalog",
"ListDatabasesResponse",
"connect_catalog",
"connect_catalog_async",
"AsyncMaterializedView",
+11 -9
View File
@@ -255,8 +255,6 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
*,
vector_source_json: Optional[str] = None,
) -> Table: ...
async def create_materialized_view_async(
self,
@@ -266,8 +264,6 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
*,
vector_source_json: Optional[str] = None,
) -> Job: ...
async def list_materialized_views(self) -> List[str]: ...
async def drop_materialized_view(
@@ -500,10 +496,10 @@ class Table:
limit: Optional[int] = None,
) -> FunctionErrors: ...
async def refresh_materialized_view(
self, full: bool = False, source_version: Optional[int] = None
self, source_version: Optional[int] = None
) -> RefreshMaterializedViewResult: ...
async def refresh_materialized_view_async(
self, full: bool = False, source_version: Optional[int] = None
self, source_version: Optional[int] = None
) -> Job: ...
async def materialized_view_definition(self) -> str: ...
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
@@ -946,9 +942,15 @@ class Catalog:
async def drop_database(
self, name: str, *, ignore_missing: bool = False
) -> None: ...
async def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> tuple[list[str], Optional[str]]: ...
def list_databases(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> DatabaseNames: ...
class DatabaseNames:
def num_page_results(self) -> int: ...
def page_token(self) -> Optional[str]: ...
def __aiter__(self) -> DatabaseNames: ...
async def __anext__(self) -> str: ...
async def connect_catalog(
endpoint: str,
+91 -19
View File
@@ -3,9 +3,8 @@
"""Remote catalogs manage databases through a server's root namespace."""
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Optional, Union
from typing import Any, AsyncIterator, Iterator, Optional, Union
from . import _lancedb
from .background_loop import LOOP
@@ -14,12 +13,66 @@ from .remote import ClientConfig, OAuthConfig
from .remote.db import RemoteDBConnection
@dataclass
class ListDatabasesResponse:
"""A page of database names and an optional continuation token."""
class AsyncDatabaseNames(AsyncIterator[str]):
"""An async iterator of database names with pagination state.
databases: list[str]
page_token: Optional[str] = None
Returned by [list_databases][lancedb.catalog.AsyncCatalog.list_databases].
"""
def __init__(self, inner: _lancedb.DatabaseNames):
self._inner = inner
def __aiter__(self) -> "AsyncDatabaseNames":
return self
async def __anext__(self) -> str:
return await self._inner.__anext__()
def num_page_results(self) -> int:
"""Return the number of names available without another REST request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the token for the next REST request.
Before the first request, this is the supplied starting token (``None``
starts at the beginning). After the final page is fetched, it is ``None``,
even if names remain cached. Drain the cache before saving a token to
avoid skipping those names when resuming. A failed request terminates
iteration and leaves its token available for resuming a new iterator.
"""
return self._inner.page_token()
class DatabaseNames(Iterator[str]):
"""A synchronous iterator of database names with pagination state.
Returned by [Catalog.list_databases][lancedb.catalog.Catalog.list_databases].
"""
def __init__(self, inner: AsyncDatabaseNames):
self._inner = inner
def __iter__(self) -> "DatabaseNames":
return self
def __next__(self) -> str:
try:
return LOOP.run(self._inner.__anext__())
except StopAsyncIteration:
raise StopIteration from None
def num_page_results(self) -> int:
"""Return the number of names available without another REST request."""
return self._inner.num_page_results()
def page_token(self) -> Optional[str]:
"""Return the token for the next REST request.
See [page_token][lancedb.catalog.AsyncDatabaseNames.page_token]
for initial, cached, final-page, and error behavior.
"""
return self._inner.page_token()
class AsyncCatalog:
@@ -50,14 +103,25 @@ class AsyncCatalog:
"""Connect to an existing database by its logical name."""
return AsyncConnection(await self._inner.connect_database(name))
async def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> ListDatabasesResponse:
"""List a page of databases. Pass the returned token for the next page."""
names, token = await self._inner.list_databases(
limit=limit, page_token=page_token
def list_databases(
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> AsyncDatabaseNames:
"""Iterate lazily over all database names using ``async for``.
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
``page_limit`` sets the maximum names per REST response (not the total);
``None`` uses the server default. Request errors are raised during iteration.
Examples
--------
```python
async for name in catalog.list_databases():
print(name)
```
"""
return AsyncDatabaseNames(
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
)
return ListDatabasesResponse(names, token)
async def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
"""Drop an empty database. A nonempty database is an error."""
@@ -73,7 +137,8 @@ class Catalog:
```python
catalog = lancedb.connect_catalog("https://my-server.example", api_key="secret")
db = catalog.create_database("analytics", exist_ok=True)
page = catalog.list_databases(limit=20)
for name in catalog.list_databases():
print(name)
```
"""
@@ -118,10 +183,17 @@ class Catalog:
)
def list_databases(
self, *, limit: Optional[int] = None, page_token: Optional[str] = None
) -> ListDatabasesResponse:
"""List a page of databases. Pass the returned token for the next page."""
return LOOP.run(self._inner.list_databases(limit=limit, page_token=page_token))
self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None
) -> DatabaseNames:
"""Iterate lazily over database names, fetching more results as needed.
``page_token`` resumes from a saved token; ``None`` starts at the beginning.
``page_limit`` sets the maximum names per REST response (not the total);
``None`` uses the server default. Request errors are raised during iteration.
"""
return DatabaseNames(
self._inner.list_databases(page_token=page_token, page_limit=page_limit)
)
def drop_database(self, name: str, *, ignore_missing: bool = False) -> None:
"""Drop an empty database. A nonempty database is an error."""
+4 -29
View File
@@ -57,7 +57,6 @@ from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewSource,
VectorDedupSource,
SelectArg,
normalize_select,
)
@@ -613,21 +612,12 @@ class DBConnection(EnforceOverrides):
table. The view is a normal table: it can be queried, indexed and
searched, and it appears in ``table_names``.
Ordinary source tables must have stable row ids (create them with the
``new_table_enable_stable_row_ids`` storage option): they keep the
view's provenance valid across source compactions, and cannot be
enabled after a table exists. Native dedup sources also support physical
row IDs interpreted against their pinned snapshot.
Parameters
----------
name: str
The name of the view.
source: str or VectorDedupSource
The name of the source table, or an indexed dedup source declared
with [vector_dedup][lancedb.vector_dedup]. Dedup sources capture a
fixed snapshot and preserve the source table. They accept no
additional select, where or limit options.
source: str
The name of the source table.
select: list or dict, optional
The view's columns: column names, ``(alias, SQL expression)``
pairs, or a dict of the same. Omitting it selects every source
@@ -1654,10 +1644,7 @@ class LanceDBConnection(DBConnection):
Examples
--------
>>> import lancedb
>>> db = lancedb.connect(
... "./.lancedb",
... storage_options={"new_table_enable_stable_row_ids": "true"},
... )
>>> db = lancedb.connect("./.lancedb")
>>> data = [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}]
>>> table = db.create_table("people", data)
>>> view = db.create_materialized_view(
@@ -2625,19 +2612,7 @@ class AsyncConnection(object):
where: Optional[str],
limit: Optional[int],
) -> Tuple[str, Dict[str, str]]:
if not isinstance(source, VectorDedupSource):
return source, {}
if select is not None or where is not None or limit is not None:
raise ValueError(
"vector_dedup cannot be combined with select, where or limit"
)
version = source.dataset_version
if version is None:
table = await self.open_table(source.source)
version = await table.version()
return source.source, {
"vector_source_json": source._native_source_json(version)
}
return source, {}
async def create_materialized_view(
self,
+18 -122
View File
@@ -6,11 +6,11 @@ maintained by refresh. See ``DBConnection.create_materialized_view``."""
from __future__ import annotations
import json
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
import json
from dataclasses import dataclass
from .background_loop import LOOP
from .job import AsyncJob, Job, _typed_job
@@ -21,7 +21,6 @@ if TYPE_CHECKING:
from .table import AsyncTable, LanceTable
DEFINITION_META_KEY = b"mv.definition"
_MAX_FLOAT32 = float.fromhex("0x1.fffffep+127")
SelectArg = Union[
str,
@@ -31,104 +30,12 @@ SelectArg = Union[
]
@dataclass(frozen=True)
class VectorDedupSource:
"""A declarative source of retained original rows from an indexed snapshot.
Construct with [vector_dedup][lancedb.vector_dedup] and pass it as the
``source`` of ``create_materialized_view`` or
``create_materialized_view_async``. An omitted ``dataset_version`` is
captured once when creation is submitted, not when this object is made.
Refreshes keep that captured version. The source table is never modified.
"""
source: str
column: str
distance_threshold: float
dataset_version: Optional[int] = None
def __post_init__(self):
for name in ("source", "column"):
value = getattr(self, name)
if not isinstance(value, str) or not value:
raise ValueError(f"{name} must be a non-empty string")
if self.dataset_version is not None and (
type(self.dataset_version) is not int
or not 0 < self.dataset_version < 2**64
):
raise ValueError("dataset_version must be a positive uint64 integer")
value = self.distance_threshold
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError("distance_threshold must be a number")
try:
value = float(value)
except OverflowError as exc:
raise ValueError("distance_threshold must fit a finite float32") from exc
if not math.isfinite(value) or abs(value) > _MAX_FLOAT32:
raise ValueError("distance_threshold must fit a finite float32")
object.__setattr__(self, "distance_threshold", value)
def _native_source_json(self, version: int) -> str:
return json.dumps(
{
"kind": "Dedup",
"dataset_version": version,
"column": self.column,
"distance_threshold": repr(self.distance_threshold),
},
allow_nan=False,
)
MaterializedViewSource = str
def vector_dedup(
source: str,
*,
column: str,
distance_threshold: float,
dataset_version: Optional[int] = None,
) -> VectorDedupSource:
"""Declare an indexed dedup source for a materialized view.
The default policy keeps the smallest unassigned snapshot row ID and removes
only its direct qualifying neighbors: an A-B-C chain without an A-C edge
retains A and C. Rows with no pair, including null vectors, survive.
Only same-segment/same-IVF-partition pairs are considered; quantized distances
use reconstructed index vectors. The index must cover the selected snapshot.
``source`` is an exact table name in the connection's root namespace, just
as for ``create_materialized_view``. ``dataset_version=None`` captures the
source's version when creation is submitted. Pass an explicit positive
version to reproduce a previous snapshot. The result contains original
source columns and rows. Do not combine this source with ``select``,
``where`` or ``limit``; query the resulting table for further transforms.
Creation follows the usual MV contract: by default it waits for the result.
Use ``with_no_data=True`` to declare only, or
``create_materialized_view_async`` to get a job handle. On Cloud/Enterprise,
execution uses server-side MV registry jobs; local connections run in process.
Examples
--------
Given a connected ``db`` and an indexed ``images.phash`` column::
from lancedb import vector_dedup
view = db.create_materialized_view(
"images_clean",
vector_dedup("images", column="phash", distance_threshold=4),
)
cleaned = view.table
"""
return VectorDedupSource(source, column, distance_threshold, dataset_version)
MaterializedViewSource = Union[str, VectorDedupSource]
DEFINITION_FORMAT = 4
DEFINITION_FORMAT = 2
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
format 2 being a query with ``GROUP BY``, format 3 native pairs and format 4
native dedup results. A ``kind`` key beside it is for
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
readers older than the format number."""
@@ -272,38 +179,31 @@ class AsyncMaterializedView:
return _definition_from_json(raw)
async def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source.
The refresh is incremental when the source's changes can be
reconciled into the view -- rows added, changed or removed since the
last one -- and otherwise rebuilds. ``full=True`` forces a rebuild;
Every refresh rebuilds the complete view. For local views,
``source_version`` refreshes to that source version instead of the
latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser raises
rather than writing them a second time.
latest; remote SQL refreshes do not support source-version pinning.
"""
return await self._table._inner.refresh_materialized_view(
full=full, source_version=source_version
source_version=source_version
)
async def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "AsyncJob[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
The job may already be complete for a local view. On LanceDB Cloud
and Enterprise, its ``id`` is the server job identifier returned by
the refresh endpoint.
The job may already be complete for a local view. Remote refreshes are
submitted through the SQL service.
"""
from ._lancedb import RefreshMaterializedViewResult
return _typed_job(
await self._table._inner.refresh_materialized_view_async(
full=full, source_version=source_version
source_version=source_version
),
RefreshMaterializedViewResult.from_json,
)
@@ -335,22 +235,18 @@ class MaterializedView:
return LOOP.run(self._async.definition())
def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source. See
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
return LOOP.run(self._async.refresh(source_version=source_version))
def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "Job[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
See
[AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async].
"""
return Job(
LOOP.run(
self._async.refresh_async(full=full, source_version=source_version)
)
)
return Job(LOOP.run(self._async.refresh_async(source_version=source_version)))
+4 -6
View File
@@ -608,6 +608,9 @@ class RemoteDBConnection(DBConnection):
One of "error", "drop", "fill".
fill_value: float
The value to use when filling vectors. Only used if on_bad_vectors="fill".
embedding_functions: list of EmbeddingFunctionConfig, optional
The embedding functions to store in the table schema and use to
generate vectors when creating the table or adding data.
Returns
-------
@@ -682,12 +685,6 @@ class RemoteDBConnection(DBConnection):
if namespace_path is None:
namespace_path = []
validate_table_name(name)
if embedding_functions is not None:
logging.warning(
"embedding_functions is not yet supported on LanceDB Cloud."
"Please vote https://github.com/lancedb/lancedb/issues/626 "
"for this feature."
)
from .table import RemoteTable
@@ -700,6 +697,7 @@ class RemoteDBConnection(DBConnection):
schema=schema,
on_bad_vectors=on_bad_vectors,
fill_value=fill_value,
embedding_functions=embedding_functions,
)
)
return RemoteTable(
+23
View File
@@ -63,6 +63,7 @@ from lancedb.table import _normalize_progress
from ..query import (
AnalyzePlanDistributedMetrics,
DocumentGranularity,
FullTextQuery,
LanceQueryBuilder,
LanceTakeQueryBuilder,
LanceVectorQueryBuilder,
@@ -77,6 +78,7 @@ from ..table import (
Tags,
)
from ..types import BaseTokenizerType
from ..util import infer_vector_column_name
class RemoteTable(Table):
@@ -734,6 +736,12 @@ class RemoteTable(Table):
- If the table has multiple vector columns then the *vector_column_name*
needs to be specified. Otherwise, an error is raised.
query_type: str, default "auto"
The type of search: "auto", "vector", "fts", or "hybrid".
With "auto", a string query uses the embedding function for the
inferred or specified vector column if available, otherwise it uses
full-text search.
fast_search: bool, optional
Skip a flat search of unindexed data. This may improve
search performance but search results will not include unindexed data.
@@ -753,6 +761,21 @@ class RemoteTable(Table):
- and also the "_distance" column which is the distance between the query
vector and the returned vector.
"""
if isinstance(query, FullTextQuery):
query_type = "fts"
elif (
query_type == "auto"
and isinstance(query, str)
and not self.embedding_functions
):
query_type = "fts"
vector_column_name = infer_vector_column_name(
schema=self.schema,
query_type=query_type,
query=query,
vector_column_name=vector_column_name,
)
return LanceQueryBuilder.create(
self,
query,
+10 -5
View File
@@ -5445,11 +5445,16 @@ class AsyncTable:
async def get_lsm_write_spec(self) -> Optional["LsmWriteSpec"]:
"""Read the LsmWriteSpec currently installed on this table.
Returns ``None`` when the MemWAL LSM write path is not enabled (no
spec has been set, or it was removed with `unset_lsm_write_spec`).
The returned spec mirrors what was passed to `set_lsm_write_spec`,
except that ``maintained_indexes`` always reports the concrete list
resolved when the spec was set — ``None`` never round-trips.
Returns ``None`` when the LSM write path is not enabled at all — no
spec has been set, or one was removed with `unset_lsm_write_spec`.
That is a different answer from a spec whose ``maintained_indexes``
is ``None``, which is an installed spec selecting indexes
automatically.
The spec read back is the one that was installed, selection
included: ``None`` maintains every supported index the table has now
or gains later, ``[]`` maintains none, and a non-empty list maintains
exactly those. All three round-trip.
"""
return await self._inner.get_lsm_write_spec()
+129 -4
View File
@@ -55,7 +55,7 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
(200, {}),
(200, {"tables": []}),
(200, {"tables": []}),
(200, {"namespaces": ["team/search"], "page_token": "next"}),
(200, {"namespaces": ["team/search"]}),
(204, None),
]
)
@@ -85,13 +85,12 @@ def test_catalog_sync_scope_and_serialization(catalog_server):
with pytest.raises(ValueError, match="sql_host_override must use"):
connection.execute_query_async("SELECT 1")
assert restored.table_names() == []
page = catalog.list_databases(limit=1, page_token="a/b")
assert page == lancedb.ListDatabasesResponse(["team/search"], "next")
assert list(catalog.list_databases()) == ["team/search"]
catalog.drop_database("team/search", ignore_missing=True)
assert requests[0][0] == "/v1/namespace/team%2Fsearch/create"
assert requests[0][2] == {"mode": "ExistOk"}
assert requests[1][0] == "/v1/namespace/team%2Fsearch/describe"
assert requests[4][0] == "/v1/namespace/$/list?limit=1&page_token=a%2Fb"
assert requests[4][0] == "/v1/namespace/$/list"
assert requests[5][2] == {"mode": "Skip", "behavior": "Restrict"}
for i, (_, headers, _) in enumerate(requests):
headers = {key.lower(): value for key, value in headers.items()}
@@ -138,3 +137,129 @@ async def test_catalog_async_and_errors(catalog_server):
def test_catalog_requires_remote_endpoint(endpoint):
with pytest.raises(ValueError, match="endpoint"):
lancedb.connect_catalog(endpoint)
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
@pytest.mark.parametrize("page_token,page_limit", [(None, None), ("start/token", 2)])
async def test_database_names_iteration(
catalog_server, asynchronous, page_token, page_limit
):
endpoint, requests, responses = catalog_server
responses.extend(
[
(200, {"namespaces": ["first", "second"], "page_token": "a/b"}),
(200, {"namespaces": [], "page_token": "next"}),
(200, {"namespaces": ["last", "final"], "page_token": ""}),
]
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_token=page_token, page_limit=page_limit)
assert requests == []
assert names.num_page_results() == 0
assert names.page_token() == page_token
for cached, expected in [(1, "first"), (0, "second")]:
assert (await names.__anext__() if asynchronous else next(names)) == expected
assert names.num_page_results() == cached
assert names.page_token() == "a/b"
assert len(requests) == 1
assert (await names.__anext__() if asynchronous else next(names)) == "last"
assert names.num_page_results() == 1
assert names.page_token() is None
assert (await names.__anext__() if asynchronous else next(names)) == "final"
assert names.num_page_results() == 0
assert names.page_token() is None
for _ in range(2):
with pytest.raises(StopAsyncIteration if asynchronous else StopIteration):
if asynchronous:
await names.__anext__()
else:
next(names)
from urllib.parse import parse_qs, urlsplit
for request, token in zip(requests, [page_token, "a/b", "next"]):
url = urlsplit(request[0])
assert url.path == "/v1/namespace/$/list"
expected = {}
if token is not None:
expected["page_token"] = [token]
if page_limit is not None:
expected["limit"] = [str(page_limit)]
assert parse_qs(url.query) == expected
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("fail", [False, True])
@pytest.mark.asyncio
async def test_database_names_empty_or_error(catalog_server, asynchronous, fail):
endpoint, requests, responses = catalog_server
responses.append(
(401, {"error": "unauthorized"}) if fail else (200, {"namespaces": []})
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases()
if fail:
with pytest.raises(HttpError):
if asynchronous:
await names.__anext__()
else:
next(names)
assert ([name async for name in names] if asynchronous else list(names)) == []
assert len(requests) == 1
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.asyncio
async def test_database_names_resume(catalog_server, asynchronous):
endpoint, requests, responses = catalog_server
responses.extend(
[
(200, {"namespaces": ["first"], "page_token": "resume/token"}),
(200, {"namespaces": ["second"]}),
]
)
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_limit=1)
assert (await names.__anext__() if asynchronous else next(names)) == "first"
assert names.num_page_results() == 0
resumed = catalog.list_databases(page_token=names.page_token(), page_limit=1)
assert resumed.page_token() == "resume/token"
assert len(requests) == 1
assert ([name async for name in resumed] if asynchronous else list(resumed)) == [
"second"
]
assert resumed.page_token() is None
assert resumed.num_page_results() == 0
assert requests[1][0] == "/v1/namespace/$/list?limit=1&page_token=resume%2Ftoken"
@pytest.mark.parametrize("asynchronous", [False, True])
@pytest.mark.parametrize("page_limit", [0, 2**32 - 1])
@pytest.mark.asyncio
async def test_database_names_invalid_limit(catalog_server, asynchronous, page_limit):
endpoint, requests, _ = catalog_server
catalog = (
await lancedb.connect_catalog_async(endpoint)
if asynchronous
else lancedb.connect_catalog(endpoint)
)
names = catalog.list_databases(page_limit=page_limit)
with pytest.raises(ValueError, match="limit"):
if asynchronous:
await names.__anext__()
else:
next(names)
assert names.num_page_results() == 0
assert requests == []
+17 -191
View File
@@ -97,24 +97,6 @@ def test_remote_list_uses_namespace_route():
assert requests == ["/v1/namespace/$/materialized_view/list"]
def test_remote_create_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
job = db.create_materialized_view_async("adults", "people", where="age >= 18")
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/adults/create",
{"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False},
)
]
def test_remote_drop_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
@@ -221,29 +203,29 @@ def test_definition_round_trips(tmp_path):
)
def test_incremental_refresh_after_append(tmp_path):
def test_full_refresh_after_append(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").add([{"name": "alan", "age": 41}])
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 4
assert view.table.count_rows() == 4
assert view.refresh().mode == "no_op"
assert view.refresh().mode == "rebuild"
def test_incremental_refresh_after_update(tmp_path):
def test_full_refresh_after_update(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 3
rows = view.table.search().to_list()
assert sorted(row["age"] for row in rows) == [8, 36, 85]
@@ -331,8 +313,8 @@ async def test_async_create_and_refresh_jobs(tmp_path):
@pytest.mark.asyncio
async def test_async_incremental(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
async def test_async_refresh_rebuilds(tmp_path):
db = await lancedb.connect_async(tmp_path)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view("copy", "people", with_no_data=True)
await view.refresh()
@@ -340,15 +322,15 @@ async def test_async_incremental(tmp_path):
table = await db.open_table("people")
await table.add([{"name": "alan", "age": 41}])
result = await view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 2
def test_source_requires_stable_row_ids(tmp_path):
def test_source_does_not_require_stable_row_ids(tmp_path):
db = lancedb.connect(tmp_path)
db.create_table("plain", [{"x": 1}])
with pytest.raises(Exception, match="stable row ids"):
db.create_materialized_view("v", "plain")
view = db.create_materialized_view("v", "plain")
assert view.table.count_rows() == 1
def test_bare_select_names_are_quoted(tmp_path):
@@ -373,7 +355,7 @@ def test_scalar_select_is_one_column(tmp_path):
view = db.create_materialized_view("just_name", "people", select="name")
view.refresh()
rows = view.table.search().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"]
@@ -384,7 +366,7 @@ async def test_async_scalar_select_is_one_column(tmp_path):
view = await db.create_materialized_view("just_name", "people", select="name")
await view.refresh()
rows = await view.table.query().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
def test_limit_above_i64_max_is_refused(tmp_path):
@@ -520,166 +502,10 @@ def test_stored_queries_and_legacy_layouts_are_read():
"SELECT * FROM people"
)
for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]:
native = f"SELECT * FROM {function}('images', 3, 'phash', 4)"
assert read({"format": fmt, "query": native}).query == native
# A newer writer's layout is reported, never guessed at.
for newer in (
{"format": 5, "query": query},
{"format": 3, "query": query},
{"kind": "select_v3", "source_table": "people"},
):
with pytest.raises(NotImplementedError, match="cannot refresh"):
read(newer)
def indexed_dedup_source(db, name="images"):
import pyarrow as pa
from lancedb.index import IvfFlat
vectors = pa.FixedSizeListArray.from_arrays(
pa.array(
[v for i in range(256) for v in (float(i if i < 3 else i * 10), 0.0)],
type=pa.float32(),
),
2,
)
table = db.create_table(name, pa.table({"id": range(256), "vector": vectors}))
table.create_index("vector", config=IvfFlat(num_partitions=1))
return table
@pytest.mark.parametrize("stable", [False, True])
def test_vector_dedup_source_captures_snapshot_and_preserves_source(tmp_path, stable):
db = lancedb.connect(
tmp_path,
storage_options={"new_table_enable_stable_row_ids": str(stable).lower()},
)
source = indexed_dedup_source(db, "images.with.dot")
version = source.version
expression = lancedb.vector_dedup(
source.name, column="vector", distance_threshold=1
)
view = db.create_materialized_view("clean", expression)
assert "vector_dedup(" in view.definition.query
assert f", {version}, 'vector', 1.0)" in view.definition.query
assert view.table.count_rows() == 255
assert [r["id"] for r in view.table.search().where("id < 3").to_list()] == [0, 2]
assert source.version == version
assert source.count_rows() == 256
source.delete("id = 0")
view.refresh_async().wait()
assert view.table.search().where("id = 0").to_list() == [
{"id": 0, "vector": [0.0, 0.0]}
]
assert expression.dataset_version is None # reusable immutable declaration
def test_vector_dedup_create_job_and_explicit_snapshot(tmp_path):
db = lancedb.connect(tmp_path)
source = indexed_dedup_source(db)
version = source.version
source.delete("id = 0")
job = db.create_materialized_view_async(
"clean",
lancedb.vector_dedup(
"images", column="vector", distance_threshold=1, dataset_version=version
),
)
assert job.wait() is None
view = db.open_materialized_view("clean")
assert view.table.count_rows() == 255
assert len(view.table.search().where("id = 0").to_list()) == 1
@pytest.mark.asyncio
async def test_async_vector_dedup_declaration_and_creation(tmp_path):
sync_db = lancedb.connect(tmp_path)
source = indexed_dedup_source(sync_db)
db = await lancedb.connect_async(tmp_path)
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
view = await db.create_materialized_view("clean", expression, with_no_data=True)
assert await view.table.count_rows() == 0
# The snapshot is captured at declaration even if refresh comes later.
source.delete("id = 0")
job = await view.refresh_async()
await job.wait()
assert await view.table.count_rows() == 255
another = await db.create_materialized_view_async(
"later", expression, with_no_data=True
)
await another.wait()
later = await db.open_materialized_view("later")
assert (await later.definition()).query != (await view.definition()).query
assert await later.table.count_rows() == 0
@pytest.mark.parametrize(
"options", [{"select": ["id"]}, {"select": []}, {"where": "id > 0"}, {"limit": 1}]
)
def test_vector_dedup_rejects_extra_clauses_before_opening_source(tmp_path, options):
db = lancedb.connect(tmp_path)
expression = lancedb.vector_dedup("missing", column="vector", distance_threshold=1)
with pytest.raises(ValueError, match="cannot be combined"):
db.create_materialized_view("bad", expression, **options)
assert db.list_tables().tables == []
@pytest.mark.parametrize("version", [0, -1, True, 1.5, "2", 2**64])
def test_vector_dedup_rejects_invalid_snapshot(version):
with pytest.raises(ValueError, match="dataset_version"):
lancedb.vector_dedup(
"images", column="vector", distance_threshold=1, dataset_version=version
)
@pytest.mark.parametrize(
"threshold", [float("nan"), float("inf"), -float("inf"), 1e100, 10**1000, True, "4"]
)
def test_vector_dedup_rejects_invalid_threshold(threshold):
with pytest.raises((ValueError, TypeError), match="distance_threshold"):
lancedb.vector_dedup("images", column="vector", distance_threshold=threshold)
def test_vector_dedup_descriptor_is_frozen():
from dataclasses import FrozenInstanceError
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
with pytest.raises(FrozenInstanceError):
expression.column = "other"
@pytest.mark.parametrize("with_no_data", [False, True])
def test_remote_vector_dedup_uses_mv_endpoint_without_sql_client(with_no_data):
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
sql_host_override="grpc://127.0.0.1:1",
client_config={"retry_config": {"retries": 0}},
)
job = db.create_materialized_view_async(
"clean",
lancedb.vector_dedup(
"images.with'dot",
column="vector'field",
distance_threshold=4,
dataset_version=7,
),
with_no_data=with_no_data,
)
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/clean/create",
{
"query": (
"SELECT * FROM vector_dedup('`images.with''dot`', "
"7, 'vector''field', 4.0)"
),
"with_no_data": with_no_data,
},
)
]
+243 -21
View File
@@ -17,8 +17,9 @@ from packaging.version import Version
import lancedb
import numpy as np
from lancedb.conftest import MockTextEmbeddingFunction
from lancedb.query import AsyncQuery, ColumnOrdering
from lancedb.conftest import MockNonNormTextEmbeddingFunction, MockTextEmbeddingFunction
from lancedb.embeddings import EmbeddingFunctionConfig, EmbeddingFunctionRegistry
from lancedb.query import AsyncQuery, ColumnOrdering, LanceVectorQueryBuilder
from lancedb.remote import ClientConfig
from lancedb.remote.errors import HttpError, RetryError
import pytest
@@ -596,6 +597,106 @@ def test_remote_permutation_is_picklable():
]
@pytest.mark.parametrize("with_schema", [True, False])
def test_create_table_embedding_functions(with_schema):
func = MockNonNormTextEmbeddingFunction.create()
config = EmbeddingFunctionConfig(
source_column="text", vector_column="vector", function=func
)
schema = pa.schema(
[pa.field("text", pa.string()), pa.field("vector", pa.list_(pa.float32(), 10))]
)
received = {}
def handler(request):
if request.path == "/v1/table/test/describe/":
# Echo the metadata actually sent by create_table, as the server does.
metadata = received["create"].schema.metadata or {}
send_json(
request,
{
"version": 1,
"schema": {
"fields": [
{
"name": "text",
"type": {"type": "string"},
"nullable": True,
},
{
"name": "vector",
"type": {
"type": "fixed_size_list",
"fields": [
{
"name": "item",
"type": {"type": "float"},
"nullable": True,
}
],
"length": 10,
},
"nullable": True,
},
],
"metadata": {
k.decode(): v.decode() for k, v in metadata.items()
},
},
},
)
elif request.path in (
"/v1/table/test/create/?mode=create",
"/v1/table/test/insert/",
):
if request.headers.get("Transfer-Encoding") == "chunked":
body = bytearray()
while True:
size = int(request.rfile.readline(), 16)
if size == 0:
request.rfile.readline()
break
body.extend(request.rfile.read(size))
request.rfile.read(2)
else:
body = request.rfile.read(int(request.headers["Content-Length"]))
operation = "create" if "/create/" in request.path else "insert"
received[operation] = pa.ipc.open_stream(body).read_all()
send_json(request, {})
else:
request.send_response(404)
request.end_headers()
with mock_lancedb_connection(handler) as db:
table = db.create_table(
"test",
schema=schema if with_schema else None,
data=None if with_schema else [{"text": "hello world"}],
embedding_functions=[config],
)
metadata = pa.schema(
[],
metadata=EmbeddingFunctionRegistry.get_instance().get_table_metadata(
[config]
),
).metadata
assert received["create"].schema.metadata == metadata
assert table.schema.metadata == metadata
assert table.schema.field("vector").type == pa.list_(pa.float32(), 10)
assert table.embedding_functions["vector"].source_column == "text"
if with_schema:
assert received["create"].num_rows == 0
else:
assert received["create"]["vector"].to_pylist() == [
func.compute_source_embeddings(["hello world"])[0].tolist()
]
table.add([{"text": "goodbye world"}])
assert received["insert"]["vector"].to_pylist() == [
func.compute_source_embeddings(["goodbye world"])[0].tolist()
]
def test_create_table_exist_ok():
def handler(request):
if request.path == "/v1/table/test/create/?mode=exist_ok":
@@ -1502,14 +1603,51 @@ def test_checkpoint_lsm_sync():
@contextlib.contextmanager
def query_test_table(query_handler, *, server_version=Version("0.1.0")):
def query_test_table(
query_handler,
*,
server_version=Version("0.1.0"),
vector_columns=None,
embedding_functions=None,
):
if vector_columns is None:
vector_columns = {"vector": 3}
schema = {
"fields": [
{"name": "text", "type": {"type": "string"}, "nullable": False},
*[
{
"name": name,
"type": {
"type": "fixed_size_list",
"fields": [
{
"name": "item",
"type": {"type": "float"},
"nullable": True,
}
],
"length": dim,
},
"nullable": False,
}
for name, dim in vector_columns.items()
],
]
}
if embedding_functions:
metadata = EmbeddingFunctionRegistry.get_instance().get_table_metadata(
embedding_functions
)
schema["metadata"] = {key: value.decode() for key, value in metadata.items()}
def handler(request):
if request.path == "/v1/table/test/describe/":
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.send_header("phalanx-version", str(server_version))
request.end_headers()
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
request.wfile.write(json.dumps({"version": 1, "schema": schema}).encode())
elif request.path == "/v1/table/test/query/":
content_len = int(request.headers.get("Content-Length"))
body = request.rfile.read(content_len)
@@ -1560,6 +1698,7 @@ def test_query_sync_minimal():
"upper_bound": None,
"ef": None,
"vector": [1.0, 2.0, 3.0],
"vector_column": "vector",
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
@@ -1574,7 +1713,84 @@ def test_query_sync_minimal():
assert data == expected
def test_query_sync_empty_query():
@pytest.mark.parametrize("query_type", ["auto", "vector", "hybrid"])
@pytest.mark.parametrize("vector_column_name", [None, "embedding"])
def test_query_sync_text_embedding(query_type, vector_column_name):
embedding_func = MockTextEmbeddingFunction.create()
config = EmbeddingFunctionConfig(
source_column="text", vector_column="embedding", function=embedding_func
)
query = "quick brown fox"
expected_vector = embedding_func.compute_query_embeddings(query)[0]
queries = []
def handler(body):
queries.append(body)
if "full_text_query" in body:
assert query_type == "hybrid"
assert body["full_text_query"]["query"] == query
return pa.table({"id": [1], "_rowid": [1], "_score": [1.0]})
assert body["vector_column"] == "embedding"
assert body["vector"] == pytest.approx(expected_vector)
return pa.table({"id": [1], "_rowid": [1], "_distance": [0.0]})
with query_test_table(
handler,
vector_columns={
"embedding": embedding_func.ndims(),
**({"other": embedding_func.ndims()} if vector_column_name else {}),
},
embedding_functions=[config],
) as table:
assert list(table.embedding_functions) == ["embedding"]
builder = table.search(
query, query_type=query_type, vector_column_name=vector_column_name
)
if query_type != "hybrid":
assert isinstance(builder, LanceVectorQueryBuilder)
results = builder.to_list()
assert results[0]["id"] == 1
if query_type != "hybrid":
assert results[0]["_distance"] == 0.0
assert len(queries) == (2 if query_type == "hybrid" else 1)
@pytest.mark.parametrize("query_type", ["auto", "vector", "hybrid"])
def test_query_sync_ambiguous_vector_columns(query_type):
config = EmbeddingFunctionConfig(
source_column="text",
vector_column="v1",
function=MockTextEmbeddingFunction.create(),
)
queries = []
def handler(body):
queries.append(body)
return pa.table({"id": [1]})
with query_test_table(
handler,
vector_columns={"v1": 10, "v2": 10},
embedding_functions=[config],
) as table:
with pytest.raises(ValueError, match=r"Candidates: \['v1', 'v2'\]"):
table.search("puppy", query_type=query_type).to_list()
assert queries == []
@pytest.mark.parametrize("vector_column, dim", [("v1", 3), ("v2", 4)])
def test_query_sync_infers_vector_column_by_dimension(vector_column, dim):
def handler(body):
assert body["vector_column"] == vector_column
assert body["vector"] == [1.0] * dim
return pa.table({"id": [1]})
with query_test_table(handler, vector_columns={"v1": 3, "v2": 4}) as table:
assert table.search([1.0] * dim).to_list() == [{"id": 1}]
@pytest.mark.parametrize("vector_columns", [{}, {"v1": 3, "v2": 3}])
def test_query_sync_empty_query(vector_columns):
def handler(body):
assert body == {
"k": 10,
@@ -1587,7 +1803,7 @@ def test_query_sync_empty_query():
return pa.table({"id": [1, 2, 3]})
with query_test_table(handler) as table:
with query_test_table(handler, vector_columns=vector_columns) as table:
data = table.search(None).where("true").select(["id"]).limit(10).to_list()
expected = [{"id": 1}, {"id": 2}, {"id": 3}]
assert data == expected
@@ -1837,7 +2053,9 @@ def test_query_sync_batch_queries(server_version):
assert results == [{"id": 1, "query_index": 0}, {"id": 1, "query_index": 1}]
def test_query_sync_fts():
@pytest.mark.parametrize("vector_columns", [{}, {"vector": 3}, {"v1": 3, "v2": 3}])
@pytest.mark.parametrize("query_type", ["auto", "fts"])
def test_query_sync_fts(vector_columns, query_type):
def handler(body):
assert body == {
"full_text_query": {
@@ -1852,8 +2070,8 @@ def test_query_sync_fts():
return pa.table({"id": [1, 2, 3]})
with query_test_table(handler) as table:
(table.search("puppy", query_type="fts").to_list())
with query_test_table(handler, vector_columns=vector_columns) as table:
(table.search("puppy", query_type=query_type).to_list())
def handler(body):
assert body == {
@@ -1880,9 +2098,11 @@ def test_query_sync_fts():
return pa.table({"id": [1, 2, 3]})
with query_test_table(handler) as table:
with query_test_table(handler, vector_columns=vector_columns) as table:
(
table.search("puppy", query_type="fts", fts_columns=["name", "description"])
table.search(
"puppy", query_type=query_type, fts_columns=["name", "description"]
)
.with_row_id(True)
.limit(42)
.to_list()
@@ -1920,7 +2140,9 @@ def test_query_sync_fts_document_granularity():
}
)
with query_test_table(handler, server_version=Version("0.6.0")) as table:
with query_test_table(
handler, server_version=Version("0.6.0"), vector_columns={}
) as table:
result = table.search(
MatchQuery(
"alpha",
@@ -1955,6 +2177,7 @@ def test_query_sync_hybrid():
"prefilter": True,
"refine_factor": None,
"vector": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
"vector_column": "vector",
"nprobes": 20,
"minimum_nprobes": 20,
"maximum_nprobes": 20,
@@ -1966,15 +2189,14 @@ def test_query_sync_hybrid():
}
return pa.table({"_rowid": [1, 2, 3], "_distance": [0.1, 0.2, 0.3]})
with query_test_table(handler) as table:
embedding_func = MockTextEmbeddingFunction()
embedding_config = MagicMock()
embedding_config.function = embedding_func
embedding_funcs = MagicMock()
embedding_funcs.get = MagicMock(return_value=embedding_config)
table.embedding_functions = embedding_funcs
config = EmbeddingFunctionConfig(
source_column="text",
vector_column="vector",
function=MockTextEmbeddingFunction.create(),
)
with query_test_table(
handler, vector_columns={"vector": 10}, embedding_functions=[config]
) as table:
(table.search("puppy", query_type="hybrid").limit(42).to_list())
+45 -15
View File
@@ -1,12 +1,13 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::sync::{Arc, Mutex};
use std::time::Duration;
use lancedb::catalog::{
CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest, ListDatabasesRequest,
};
use pyo3::exceptions::PyValueError;
use futures::{StreamExt, future::poll_fn};
use lancedb::catalog::{CatalogConnection, CreateDatabaseRequest, DropDatabaseRequest};
use pyo3::exceptions::{PyStopAsyncIteration, PyValueError};
use pyo3::{Bound, PyAny, PyRef, PyResult, Python, pyclass, pyfunction, pymethods};
use crate::connection::{Connection, PyClientConfig};
@@ -67,19 +68,48 @@ impl Catalog {
})
}
#[pyo3(signature = (*, limit=None, page_token=None))]
fn list_databases<'py>(
self_: PyRef<'py, Self>,
limit: Option<u32>,
page_token: Option<String>,
) -> PyResult<Bound<'py, PyAny>> {
#[pyo3(signature = (*, page_token=None, page_limit=None))]
fn list_databases(&self, page_token: Option<String>, page_limit: Option<u32>) -> DatabaseNames {
DatabaseNames {
inner: Arc::new(Mutex::new(
self.inner.list_databases(page_token, page_limit),
)),
next_lock: Arc::new(tokio::sync::Mutex::new(())),
}
}
}
#[pyclass]
pub struct DatabaseNames {
inner: Arc<Mutex<lancedb::catalog::DatabaseNames>>,
next_lock: Arc<tokio::sync::Mutex<()>>,
}
#[pymethods]
impl DatabaseNames {
fn num_page_results(&self, py: Python<'_>) -> usize {
py.detach(|| self.inner.lock().unwrap().num_page_results())
}
fn page_token(&self, py: Python<'_>) -> Option<String> {
py.detach(|| self.inner.lock().unwrap().page_token().map(str::to_owned))
}
fn __aiter__(self_: PyRef<'_, Self>) -> PyRef<'_, Self> {
self_
}
fn __anext__(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner.clone();
let next_lock = self_.next_lock.clone();
future_into_py(self_.py(), async move {
let mut request = ListDatabasesRequest::default();
request.limit = limit;
request.page_token = page_token;
let response = inner.list_databases(request).await.infer_error()?;
Ok((response.databases, response.page_token))
// Serialize advances while allowing synchronous state inspection
// between polls, including while a REST request is pending.
let _guard = next_lock.lock().await;
poll_fn(|cx| inner.lock().unwrap().poll_next_unpin(cx))
.await
.ok_or_else(|| PyStopAsyncIteration::new_err(""))?
.infer_error()
})
}
}
+2 -18
View File
@@ -407,8 +407,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -417,17 +416,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -443,8 +435,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
@@ -453,17 +444,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
+1
View File
@@ -64,6 +64,7 @@ pub fn _lancedb(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
env_logger::init_from_env(env);
m.add_class::<Connection>()?;
m.add_class::<catalog::Catalog>()?;
m.add_class::<catalog::DatabaseNames>()?;
m.add_function(wrap_pyfunction!(catalog::connect_catalog, m)?)?;
m.add_class::<Session>()?;
m.add_class::<Table>()?;
+11 -17
View File
@@ -267,7 +267,7 @@ fn fmt_maintained(maintained: &Option<Vec<String>>) -> String {
/// Constructed via the `bucket(...)`, `identity(...)`, or `unsharded()`
/// classmethods, then optionally chain `with_maintained_indexes(...)` and
/// `with_writer_config_defaults(...)`. A fresh spec maintains every index the
/// MemWAL supports, resolved on install.
/// table has, including ones created later.
#[pyclass(module = "lancedb._lancedb", from_py_object)]
#[derive(Clone, Debug)]
pub struct LsmWriteSpec {
@@ -307,9 +307,9 @@ impl LsmWriteSpec {
}
}
/// Set which indexes the MemWAL maintains. `None` (the default)
/// resolves every supported index on install; a list is verbatim,
/// and an empty list maintains nothing.
/// Set which indexes the MemWAL maintains. `None` (the default) is
/// every index the table has, including ones created later; a list is
/// verbatim, and an empty list maintains nothing.
#[pyo3(signature = (indexes))]
pub fn with_maintained_indexes(&self, indexes: Option<Vec<String>>) -> Self {
Self {
@@ -592,8 +592,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
fn from(result: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match result.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
@@ -1817,10 +1815,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1828,7 +1825,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1837,10 +1834,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view_async(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1848,7 +1844,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1863,11 +1859,9 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
view.definition().to_json().map_err(|err| {
PyRuntimeError::new_err(format!(
"failed to serialize materialized-view definition: {err}"
))
})
Ok(lancedb::materialized_view::definition_metadata_from_sql(
view.definition_sql(),
))
})
}
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "lancedb"
version = "0.40.0-beta.12"
version = "0.40.0-beta.13"
edition.workspace = true
description = "LanceDB: A serverless, low-latency vector database for AI applications"
license.workspace = true
+100 -5
View File
@@ -15,7 +15,12 @@
//! ```
use std::fmt;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use futures::future::BoxFuture;
use futures::{FutureExt, Stream};
use crate::Result;
use crate::connection::Connection;
@@ -118,6 +123,73 @@ pub struct ListDatabasesResponse {
pub page_token: Option<String>,
}
/// A lazy stream of database names with inspectable pagination state.
///
/// The continuation token advances when a page is fetched, so resuming from it
/// skips any names still cached in this stream. Drain the cache before saving a
/// token if those names must be included. Errors terminate the stream; the token
/// for the failed request remains available for resuming with a new stream.
pub struct DatabaseNames {
catalog: Arc<dyn Catalog>,
names: std::vec::IntoIter<String>,
page_token: Option<String>,
page_limit: Option<u32>,
done: bool,
pending: Option<BoxFuture<'static, Result<ListDatabasesResponse>>>,
}
impl DatabaseNames {
/// Number of names available without another REST request.
pub fn num_page_results(&self) -> usize {
self.names.len()
}
/// Token for the next REST request, or `None` after the last page is fetched.
///
/// Before the first request, returns the supplied starting token (`None`
/// means start from the beginning). A final page can still have cached names.
pub fn page_token(&self) -> Option<&str> {
self.page_token.as_deref()
}
}
impl Stream for DatabaseNames {
type Item = Result<String>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
if let Some(name) = this.names.next() {
return Poll::Ready(Some(Ok(name)));
}
if this.done {
return Poll::Ready(None);
}
let pending = this.pending.get_or_insert_with(|| {
let catalog = this.catalog.clone();
let request = ListDatabasesRequest {
page_token: this.page_token.clone(),
limit: this.page_limit,
};
async move { catalog.list_databases(request).await }.boxed()
});
let response = futures::ready!(pending.as_mut().poll(cx));
this.pending = None;
match response {
Ok(response) => {
this.page_token = response.page_token.filter(|token| !token.is_empty());
this.done = this.page_token.is_none();
this.names = response.databases.into_iter();
}
Err(error) => {
this.done = true;
return Poll::Ready(Some(Err(error)));
}
}
}
}
}
/// A backend that manages databases. Implementations own database lifecycle semantics.
#[async_trait::async_trait]
pub trait Catalog: Send + Sync + std::fmt::Debug + 'static {
@@ -188,12 +260,35 @@ impl CatalogConnection {
self.catalog.drop_database(request.into()).await
}
/// List a page of databases. Pass [`ListDatabasesRequest::default`] for the first page.
pub async fn list_databases(
/// Iterate lazily over all database names, fetching more results as needed.
///
/// `page_token` starts at a saved continuation token; `None` starts at the
/// beginning. `page_limit` limits each REST response (not the whole stream);
/// `None` uses the server default. Errors terminate the listing.
///
/// ```
/// # async fn example(catalog: &lancedb::catalog::CatalogConnection) -> lancedb::Result<()> {
/// use futures::TryStreamExt;
/// let mut names = catalog.list_databases(None, None);
/// while let Some(name) = names.try_next().await? {
/// println!("{name}");
/// }
/// # Ok(())
/// # }
/// ```
pub fn list_databases(
&self,
request: ListDatabasesRequest,
) -> Result<ListDatabasesResponse> {
self.catalog.list_databases(request).await
page_token: Option<String>,
page_limit: Option<u32>,
) -> DatabaseNames {
DatabaseNames {
catalog: self.catalog.clone(),
names: Vec::new().into_iter(),
page_token,
page_limit,
done: false,
pending: None,
}
}
/// Connect to a database without creating it if it is missing.
+50 -1
View File
@@ -3,6 +3,8 @@
//! Handles to operations a server may run asynchronously.
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use arrow_array::RecordBatch;
@@ -250,6 +252,18 @@ where
pub(crate) fn spawned(task: JoinHandle<Result<T>>) -> Self {
Self::new_typed(Box::new(SpawnedJob::new(task)))
}
pub(crate) fn spawned_with_cancellation<F, Fut>(task: JoinHandle<Result<T>>, cancel: F) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let cancellation: Cancellation = Arc::new(move || Box::pin(cancel()));
Self::new_typed(Box::new(SpawnedJob::new_with_cancellation(
task,
cancellation,
)))
}
}
impl<T> Job<T>
@@ -501,10 +515,28 @@ impl Outcome {
struct SpawnedJob {
outcome: watch::Receiver<Option<Outcome>>,
abort: AbortHandle,
cancellation: Option<Cancellation>,
}
type CancellationFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
type Cancellation = Arc<dyn Fn() -> CancellationFuture + Send + Sync>;
impl SpawnedJob {
fn new<T>(task: JoinHandle<Result<T>>) -> Self
where
T: Serialize + Send + 'static,
{
Self::new_inner(task, None)
}
fn new_with_cancellation<T>(task: JoinHandle<Result<T>>, cancellation: Cancellation) -> Self
where
T: Serialize + Send + 'static,
{
Self::new_inner(task, Some(cancellation))
}
fn new_inner<T>(task: JoinHandle<Result<T>>, cancellation: Option<Cancellation>) -> Self
where
T: Serialize + Send + 'static,
{
@@ -526,7 +558,11 @@ impl SpawnedJob {
};
let _ = tx.send(Some(outcome));
});
Self { outcome, abort }
Self {
outcome,
abort,
cancellation,
}
}
}
@@ -556,6 +592,19 @@ impl JobHandle for SpawnedJob {
}
async fn cancel(&self) -> Result<()> {
if let Some(cancellation) = &self.cancellation {
let mut outcome = self.outcome.clone();
tokio::select! {
biased;
settled = outcome.wait_for(Option::is_some) => {
settled.map_err(|_| Error::Runtime {
message: "job outcome was dropped before cancellation completed".to_string(),
})?;
return Ok(());
}
result = cancellation() => result?,
}
}
self.abort.abort();
Ok(())
}
File diff suppressed because it is too large Load Diff
@@ -1,730 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Differential refresh testing.
//!
//! The refresh contract is a property: after any sequence of source
//! mutations, a view maintained by default (incremental-where-possible)
//! refreshes equals the definition evaluated against the source directly,
//! and so does a forced rebuild. The oracle is an independent read of the
//! source -- plain column scan, filter applied in Rust -- so it shares
//! nothing with the refresh path it checks.
//!
//! The oracle runs after every step, not just at the end: a later mutation
//! that forces a rebuild would silently heal an incremental error, and those
//! transient errors are exactly the bugs this exists to catch.
use arrow_array::{Float32Array, Int32Array, RecordBatch};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use futures::{StreamExt, TryStreamExt};
use lance::dataset::NewColumnTransform;
use std::sync::Arc;
use super::MaterializedView;
use super::refresh::RefreshMode;
use crate::connect;
use crate::connection::Connection;
use crate::query::{ExecutableQuery, QueryBase, Select};
use crate::table::{CompactionOptions, OptimizeAction, Table};
/// One source mutation, one per correctness-relevant class.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SrcOp {
/// Fresh non-colliding ids of both parities, so every other op has
/// view-resident rows to act on: the only op that should refresh
/// incrementally.
AppendNew,
/// Deletion in surviving fragments must break the pure-append check.
DeleteEven,
/// An in-place update; on the filtered shape it crosses the predicate,
/// so rows must leave the view.
UpdateOddScore,
/// Fragment rewrite/renumber must break the pure-append check.
Compact,
/// A column the view does not read must NOT force a rebuild.
AddColumn,
/// merge_insert commits an Update whose by-source arm deletes rows, so a
/// classifier that reads Update as "changed only" loses those deletions.
MergeDropLargest,
/// merge_insert that both changes existing rows and inserts new ones in
/// one transaction.
MergeUpsert,
}
const ALL_OPS: [SrcOp; 7] = [
SrcOp::AppendNew,
SrcOp::DeleteEven,
SrcOp::UpdateOddScore,
SrcOp::Compact,
SrcOp::AddColumn,
SrcOp::MergeDropLargest,
SrcOp::MergeUpsert,
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Shape {
/// SELECT id, score.
Identity,
/// SELECT id, score WHERE score > 50: additionally sensitive to rows
/// crossing the predicate.
Filtered,
/// SELECT id, score LIMIT 4. Which rows are held depends on the order
/// they were first materialized, so the oracle checks containment and
/// the cap rather than equality.
Limited,
}
impl Shape {
fn filter(&self) -> Option<&'static str> {
match self {
Self::Identity | Self::Limited => None,
Self::Filtered => Some("score > 50"),
}
}
fn matches(&self, score: f32) -> bool {
match self {
Self::Identity | Self::Limited => true,
Self::Filtered => score > 50.0,
}
}
fn limit(&self) -> Option<usize> {
match self {
Self::Limited => Some(4),
_ => None,
}
}
}
struct Case {
conn: Connection,
source: Table,
view: MaterializedView,
shape: Shape,
next_id: i32,
added_columns: u32,
}
fn rows_batch(ids: &[i32]) -> RecordBatch {
let scores: Vec<f32> = ids.iter().map(|id| (*id * 10) as f32).collect();
RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, true),
ArrowField::new("score", DataType::Float32, true),
])),
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(Float32Array::from(scores)),
],
)
.unwrap()
}
fn merge_batch(ids: &[i32]) -> RecordBatch {
let scores: Vec<f32> = ids.iter().map(|id| (*id * 10 + 5) as f32).collect();
RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, true),
ArrowField::new("score", DataType::Float32, true),
])),
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(Float32Array::from(scores)),
],
)
.unwrap()
}
impl Case {
async fn new(shape: Shape) -> Self {
let conn = connect("memory://").execute().await.unwrap();
let source = conn
.create_table("src", rows_batch(&[1, 2, 3, 4]))
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
let mut builder = conn
.create_materialized_view("view", "src")
.select([("id", "id"), ("score", "score")]);
if let Some(filter) = shape.filter() {
builder = builder.only_if(filter);
}
if let Some(limit) = shape.limit() {
builder = builder.limit(limit as u64);
}
let view = builder.execute().await.unwrap();
Self {
conn,
source,
view,
shape,
next_id: 100,
added_columns: 0,
}
}
async fn apply(&mut self, op: SrcOp) {
match op {
SrcOp::AppendNew => {
// Mixed parity: the middle id is odd, so UpdateOddScore always
// has a filter-matching appended row to evict.
let ids = vec![self.next_id, self.next_id + 101, self.next_id + 202];
self.next_id += 303;
self.source.add(rows_batch(&ids)).execute().await.unwrap();
}
SrcOp::DeleteEven => {
self.source.delete("id % 2 = 0").await.unwrap();
}
SrcOp::UpdateOddScore => {
self.source
.update()
.column("score", "-1.0")
.only_if("id % 2 = 1")
.execute()
.await
.unwrap();
}
SrcOp::Compact => {
self.source
.optimize(OptimizeAction::Compact {
options: CompactionOptions::default(),
remap_options: None,
})
.await
.unwrap();
}
SrcOp::MergeDropLargest => {
let mut ids = self.source_ids().await;
ids.sort_unstable();
ids.pop();
if ids.is_empty() {
return;
}
let batch = rows_batch(&ids);
let reader =
arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
let mut merge = self.source.merge_insert(&["id"]);
merge.when_not_matched_by_source_delete(None);
merge.execute(Box::new(reader)).await.unwrap();
}
SrcOp::MergeUpsert => {
let mut ids = self.source_ids().await;
ids.sort_unstable();
// One row that exists (updated in place) and one that does not.
let existing = ids.first().copied().unwrap_or(self.next_id);
let fresh = self.next_id;
self.next_id += 1;
let batch = merge_batch(&[existing, fresh]);
let reader =
arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
let mut merge = self.source.merge_insert(&["id"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
merge.execute(Box::new(reader)).await.unwrap();
}
SrcOp::AddColumn => {
self.added_columns += 1;
let field = ArrowField::new(
format!("extra_{}", self.added_columns),
DataType::Int32,
true,
);
self.source
.add_columns()
.transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(
vec![field],
))))
.execute()
.await
.unwrap();
}
}
}
async fn source_ids(&self) -> Vec<i32> {
read_rows(
self.source
.query()
.select(Select::columns(&["id", "score"])),
)
.await
.into_iter()
.map(|(id, _)| id)
.collect()
}
/// The definition's result, read independently of the refresh path:
/// plain column scan, filter applied here, sorted.
async fn oracle(&self) -> Vec<(i32, i32)> {
let mut rows = read_rows(
self.source
.query()
.select(Select::columns(&["id", "score"])),
)
.await
.into_iter()
.filter(|(_, score)| self.shape.matches(*score as f32))
.collect::<Vec<_>>();
rows.sort_unstable();
rows
}
async fn view_rows(&self) -> Vec<(i32, i32)> {
let mut rows = read_rows(
self.view
.table()
.query()
.select(Select::columns(&["id", "score"])),
)
.await;
rows.sort_unstable();
rows
}
async fn check(&self, label: &str) -> Result<(), String> {
let expected = self.oracle().await;
let actual = self.view_rows().await;
let Some(cap) = self.shape.limit() else {
if expected != actual {
return Err(format!(
"{label}: view diverged from oracle\n expected: {expected:?}\n actual: {actual:?}"
));
}
return Ok(());
};
// A capped view holds some subset of the definition's result, never
// more than the cap, and never the same row twice.
if actual.len() > cap {
return Err(format!(
"{label}: view holds {} rows, over its cap of {cap}: {actual:?}",
actual.len()
));
}
let mut unique = actual.clone();
unique.dedup();
if unique.len() != actual.len() {
return Err(format!("{label}: view holds a row twice: {actual:?}"));
}
if let Some(stray) = actual.iter().find(|row| !expected.contains(row)) {
return Err(format!(
"{label}: view holds {stray:?}, which the definition does not select: {expected:?}"
));
}
// Below the cap the view must be complete, or a row was lost.
if actual.len() < cap.min(expected.len()) {
return Err(format!(
"{label}: view holds {} of {} selectable rows under a cap of {cap}: {actual:?}",
actual.len(),
expected.len()
));
}
Ok(())
}
}
async fn read_rows(query: impl ExecutableQuery) -> Vec<(i32, i32)> {
let batches = query
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
batches
.iter()
.flat_map(|batch| {
let ids = batch["id"].as_any().downcast_ref::<Int32Array>().unwrap();
let scores = batch["score"]
.as_any()
.downcast_ref::<Float32Array>()
.unwrap();
// Scores are integer-valued by construction; compare exactly.
(0..batch.num_rows())
.map(|i| (ids.value(i), scores.value(i) as i32))
.collect::<Vec<_>>()
})
.collect()
}
/// Drive one mutation sequence: refresh + oracle-check after every step,
/// then a forced rebuild checked against the same oracle.
async fn run_sequence(ops: &[SrcOp], shape: Shape) -> Result<(), String> {
let label = format!("{shape:?} {ops:?}");
let mut case = Case::new(shape).await;
case.view
.refresh()
.execute()
.await
.map_err(|e| format!("{label}: initial refresh failed: {e}"))?;
case.check(&format!("{label} (initial)")).await?;
for (step, op) in ops.iter().enumerate() {
case.apply(*op).await;
case.view
.refresh()
.execute()
.await
.map_err(|e| format!("{label}: refresh at step {step} failed: {e}"))?;
case.check(&format!("{label} (step {step}, {op:?})"))
.await?;
}
case.view
.refresh()
.full(true)
.execute()
.await
.map_err(|e| format!("{label}: final full refresh failed: {e}"))?;
case.check(&format!("{label} (final rebuild)")).await?;
// Silence the unused-connection lint without dropping it mid-case.
let _ = &case.conn;
Ok(())
}
/// Every op sequence up to `max_len`.
fn all_sequences(max_len: u32) -> Vec<Vec<SrcOp>> {
let mut sequences = Vec::new();
for len in 1..=max_len {
for mut index in 0..ALL_OPS.len().pow(len) {
let mut ops = Vec::with_capacity(len as usize);
for _ in 0..len {
ops.push(ALL_OPS[index % ALL_OPS.len()]);
index /= ALL_OPS.len();
}
sequences.push(ops);
}
}
sequences
}
async fn run_exhaustive(max_len: u32) {
let mut cases = Vec::new();
for shape in [Shape::Identity, Shape::Filtered, Shape::Limited] {
for ops in all_sequences(max_len) {
cases.push((ops, shape));
}
}
let failures: Vec<String> = futures::stream::iter(cases)
.map(|(ops, shape)| async move { run_sequence(&ops, shape).await.err() })
.buffer_unordered(8)
.filter_map(|failure| async move { failure })
.collect()
.await;
assert!(
failures.is_empty(),
"{} sequences diverged; first: {}",
failures.len(),
failures[0]
);
}
#[tokio::test(flavor = "multi_thread")]
async fn differential_exhaustive() {
run_exhaustive(3).await;
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "longer sweep; run manually"]
async fn differential_exhaustive_deep() {
run_exhaustive(4).await;
}
/// Named interleavings that double as repro handles. The mode assertions pin
/// the classifier, which value comparison alone cannot: a wrongly rebuilt
/// view still matches the oracle.
#[tokio::test(flavor = "multi_thread")]
async fn differential_named_regressions() {
// An append is the one op that must stay incremental.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AppendNew).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
case.check("append stays incremental").await.unwrap();
// A column the view does not read must not force a rebuild.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AddColumn).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 0);
// Compaction rearranges rows without changing them: the watermark
// advances and nothing rebuilds.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AppendNew).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::Compact).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 0);
case.check("compaction alone").await.unwrap();
// Fragment bookkeeping stays coherent across the compaction: the next
// append is separable and computed alone.
case.apply(SrcOp::AppendNew).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 3);
case.check("compact then append").await.unwrap();
// A row updated to no longer match the filter must leave the view --
// and the fixture must prove the eviction happened, not merely that the
// end state matches: an update that never touched a view-resident row
// would also "match".
let mut case = Case::new(Shape::Filtered).await;
case.apply(SrcOp::AppendNew).await;
case.view.refresh().execute().await.unwrap();
let before = case.view_rows().await.len();
case.apply(SrcOp::UpdateOddScore).await;
case.view.refresh().execute().await.unwrap();
let after = case.view_rows().await.len();
assert!(
after < before,
"no view-resident row was evicted ({before} -> {after}); the fixture \
no longer exercises the filtered-update transition"
);
case.check("update crosses the filter").await.unwrap();
}
// ---------------------------------------------------------------------------
// Concurrency
// ---------------------------------------------------------------------------
//
// The sequential cases above cannot observe a cross-process race: the
// per-view refresh lock is process-local, so a second refresh in this
// process queues behind the first. What is missing is not more op
// sequences but a second process. These cases add one, and assert the same
// property the harness always asserts -- the view holds each row once.
/// Rows the definition selects from the source: every id but the first,
/// read straight from the source, sharing nothing with the refresh path.
async fn concurrency_oracle(conn: &Connection) -> Vec<i32> {
let batches: Vec<RecordBatch> = conn
.open_table("src")
.execute()
.await
.unwrap()
.query()
.select(Select::columns(&["id"]))
.execute()
.await
.unwrap()
.try_collect()
.await
.unwrap();
let mut ids = Vec::new();
for batch in &batches {
let column = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..batch.num_rows() {
if column.value(i) > 1 {
ids.push(column.value(i));
}
}
}
ids.sort_unstable();
ids
}
/// The view's ids, sorted.
async fn concurrency_view_ids(conn: &Connection) -> Vec<i32> {
let batches: Vec<RecordBatch> = conn
.open_table("mv")
.execute()
.await
.unwrap()
.query()
.select(Select::columns(&["id"]))
.execute()
.await
.unwrap()
.try_collect()
.await
.unwrap();
let mut ids = Vec::new();
for batch in &batches {
let column = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..batch.num_rows() {
ids.push(column.value(i));
}
}
ids.sort_unstable();
ids
}
/// One refresh of the view at `MV_RACE_DIR`, in its own process.
///
/// Setup happens before the start barrier so warm-up does not stagger the
/// two processes. What makes the race certain rather than likely is the
/// second barrier inside `refresh()` itself, which holds every participant
/// between staging and commit.
#[tokio::test]
#[ignore = "spawned as a child process by the concurrency cases"]
async fn cross_process_refresh_child() {
let Ok(dir) = std::env::var("MV_RACE_DIR") else {
return;
};
let dir = std::path::PathBuf::from(dir);
let tag = std::env::var("MV_RACE_TAG").unwrap();
let conn = connect(dir.to_str().unwrap()).execute().await.unwrap();
let table = conn.open_table("mv").execute().await.unwrap();
let _ = table.schema().await.unwrap();
let _ = table.count_rows(None).await.unwrap();
let source = conn.open_table("src").execute().await.unwrap();
let _ = source.count_rows(None).await.unwrap();
let view = MaterializedView::from_table(table).await.unwrap();
std::fs::write(dir.join(format!("ready-{tag}")), b"1").unwrap();
while !dir.join("START").exists() {
std::thread::sleep(std::time::Duration::from_millis(2));
}
let outcome = match view.refresh().execute().await {
Ok(result) => format!("committed rows={}", result.rows_written),
Err(err) if is_commit_conflict(&err) => "conflicted".to_string(),
Err(err) => format!("failed {err}"),
};
std::fs::write(dir.join(format!("outcome-{tag}")), outcome).unwrap();
}
/// Whether a refresh lost its commit to a concurrent one, as opposed to
/// failing for any other reason.
fn is_commit_conflict(err: &crate::Error) -> bool {
let text = err.to_string();
text.contains("Retryable commit conflict") || text.contains("preempted by concurrent")
}
/// Two processes refreshing one view concurrently must leave the view
/// equal to the oracle: each selected row present exactly once.
///
/// Both plan the same incremental delta from one watermark. A refresh is
/// meant to land on the generation it planned or leave nothing behind, so
/// at most one of them may write.
#[tokio::test]
async fn concurrent_refreshes_hold_each_row_once() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().to_str().unwrap().to_string();
let conn = connect(&path).execute().await.unwrap();
conn.create_table("src", rows_batch(&[1, 2, 3, 4]))
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
let view = conn
.create_materialized_view("mv", "src")
.select([("id", "id"), ("score", "score")])
.only_if("id > 1")
.execute()
.await
.unwrap();
// Seed the watermark so the racing refreshes are both incremental.
view.refresh().execute().await.unwrap();
// Large enough that a refresh is real work rather than a formality.
let ids: Vec<i32> = (100..200_100).collect();
conn.open_table("src")
.execute()
.await
.unwrap()
.add(rows_batch(&ids))
.execute()
.await
.unwrap();
let tags = ["a", "b"];
let exe = std::env::current_exe().unwrap();
let children: Vec<std::process::Child> = tags
.iter()
.map(|tag| {
std::process::Command::new(&exe)
.args([
"--exact",
"materialized_view::differential::cross_process_refresh_child",
"--ignored",
"--nocapture",
])
.env("MV_RACE_DIR", dir.path())
.env("MV_RACE_SYNC", dir.path())
.env("MV_RACE_PEERS", "2")
.env("MV_RACE_TAG", tag)
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.unwrap()
})
.collect();
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(180);
while tags
.iter()
.any(|tag| !dir.path().join(format!("ready-{tag}")).exists())
{
assert!(
std::time::Instant::now() < deadline,
"children never became ready"
);
std::thread::sleep(std::time::Duration::from_millis(10));
}
std::fs::write(dir.path().join("START"), b"1").unwrap();
for (tag, mut child) in tags.iter().zip(children) {
let status = loop {
match child.try_wait().unwrap() {
Some(status) => break status,
None if std::time::Instant::now() >= deadline => {
child.kill().unwrap();
panic!("child {tag} never finished");
}
None => std::thread::sleep(std::time::Duration::from_millis(10)),
}
};
assert!(status.success(), "child {tag} exited {status}");
}
// Both refreshes reached the commit boundary before either committed --
// the in-refresh barrier guarantees it -- so exactly one may win.
let outcomes: Vec<String> = tags
.iter()
.map(|tag| {
std::fs::read_to_string(dir.path().join(format!("outcome-{tag}")))
.unwrap_or_else(|_| panic!("child {tag} recorded no outcome"))
})
.collect();
for tag in tags {
assert!(
dir.path().join(format!("planned-{tag}")).exists(),
"child {tag} never reached the commit boundary, so nothing was synchronized"
);
}
let committed = outcomes.iter().filter(|o| o.contains("committed")).count();
let conflicted = outcomes.iter().filter(|o| o.contains("conflicted")).count();
assert_eq!(
(committed, conflicted),
(1, 1),
"exactly one refresh may win the generation both planned: {outcomes:?}"
);
let expected = concurrency_oracle(&conn).await;
let actual = concurrency_view_ids(&conn).await;
assert_eq!(
actual.len(),
expected.len(),
"the view holds {} rows, the oracle {}: a losing refresh left rows behind",
actual.len(),
expected.len()
);
assert_eq!(actual, expected, "the view does not match the oracle");
}
@@ -1,478 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Native pair generation as partitioned materialized-view refresh work.
use std::sync::Arc;
use arrow_schema::{Schema, SchemaRef};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_sql::sqlparser::ast::{
Expr, FunctionArg, FunctionArgExpr, TableFunctionArgs, Value,
};
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::transaction::Operation;
use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams};
use lance_table::format::Fragment;
use serde::{Deserialize, Serialize};
use super::grouped_units::GroupedRefreshPlan;
use super::refresh::{
RefreshMaterializedViewResult, ensure_incarnation, ensure_no_mem_wal, open_source,
};
use super::{MaterializedViewDefinition, Planned, ViewProjection};
use crate::table::Table;
use crate::table::datafusion::udtf::duplicate_pairs::{
DUPLICATE_PAIRS_OPERATOR_VERSION, DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec,
duplicate_pair_schema, plan_duplicate_pairs,
};
use crate::{Error, Result};
pub(super) const FUNCTION_NAME: &str = "vector_duplicate_pairs";
pub(super) const DEDUP_FUNCTION_NAME: &str = "vector_dedup";
/// The result of a pinned native vector source.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum VectorSourceKind {
/// The qualifying pair relation.
Pairs,
/// Original source rows retained by greedy direct representatives.
Dedup,
}
/// A native table-function source, pinned for the lifetime of the view.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct VectorSource {
pub kind: VectorSourceKind,
pub dataset_version: u64,
pub column: String,
/// Validated numeric SQL spelling; execution and task identity use f32 bits.
pub distance_threshold: String,
}
impl VectorSource {
pub(super) fn function_name(&self) -> &'static str {
match self.kind {
VectorSourceKind::Pairs => FUNCTION_NAME,
VectorSourceKind::Dedup => DEDUP_FUNCTION_NAME,
}
}
pub(crate) fn config(&self) -> Result<DuplicatePairsConfig> {
let threshold: f32 = self
.distance_threshold
.parse()
.map_err(|_| invalid("invalid pair threshold"))?;
if self.dataset_version == 0 || self.column.is_empty() || !threshold.is_finite() {
return Err(invalid(
"duplicate pairs require a positive snapshot, a column and a finite threshold",
));
}
Ok(DuplicatePairsConfig {
dataset_version: self.dataset_version,
column: self.column.clone(),
distance_threshold: threshold,
})
}
}
fn invalid(message: impl Into<String>) -> Error {
Error::InvalidInput {
message: message.into(),
}
}
pub(super) fn source(
args: &TableFunctionArgs,
kind: VectorSourceKind,
) -> Result<(Vec<String>, String, VectorSource)> {
if args.settings.is_some() {
return Err(invalid("duplicate pairs do not accept function settings"));
}
let exprs = args
.args
.iter()
.map(|arg| match arg {
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Ok(expr),
_ => Err(invalid(
"duplicate pairs require four positional literal arguments",
)),
})
.collect::<Result<Vec<_>>>()?;
let [table, version, column, threshold] = exprs.as_slice() else {
return Err(invalid(
"native vector sources require (table, dataset_version, column, distance_threshold)",
));
};
let string = |expr: &Expr| match expr {
Expr::Value(value) => match &value.value {
Value::SingleQuotedString(value) => Ok(value.clone()),
_ => Err(invalid("table and column must be string literals")),
},
_ => Err(invalid("table and column must be string literals")),
};
let table = string(table)?;
let plain = MaterializedViewDefinition::from_sql(&format!("SELECT * FROM {table}"))?;
if plain.vector_source.is_some()
|| plain.lateral.is_some()
|| plain.filter.is_some()
|| plain.is_grouped()
|| plain.limit.is_some()
{
return Err(invalid("duplicate pairs must name one source table"));
}
let dataset_version = match version {
Expr::Value(value) => match &value.value {
Value::Number(number, _) => number
.parse()
.map_err(|_| invalid("invalid source version"))?,
_ => return Err(invalid("source version must be an integer literal")),
},
_ => return Err(invalid("source version must be an integer literal")),
};
let config = VectorSource {
kind,
dataset_version,
column: string(column)?,
distance_threshold: threshold.to_string(),
};
config.config()?;
Ok((plain.source_namespace, plain.source_table, config))
}
pub(super) fn check_shape(definition: &MaterializedViewDefinition) -> Result<()> {
if definition
.vector_source
.as_ref()
.is_some_and(|s| s.kind == VectorSourceKind::Dedup)
{
if !definition.selects_star()
|| definition.filter.is_some()
|| definition.limit.is_some()
|| definition.lateral.is_some()
|| definition.is_grouped()
{
return Err(invalid(
"a vector_dedup view must select * without other clauses",
));
}
return Ok(());
}
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect::<Vec<_>>();
if (!definition.selects_star() && definition.projections != projections)
|| definition.filter.is_some()
|| definition.limit.is_some()
|| definition.lateral.is_some()
|| definition.is_grouped()
{
return Err(invalid(
"a native pair view selects *, or row_id_a, row_id_b, distance, without other clauses",
));
}
Ok(())
}
pub(super) fn plan(
source_schema: SchemaRef,
definition: &MaterializedViewDefinition,
) -> Result<Planned> {
check_shape(definition)?;
let config = definition.vector_source.as_ref().expect("native source");
config.config()?;
source_schema.field_with_name(&config.column)?;
if config.kind == VectorSourceKind::Dedup {
return Ok(Planned {
definition: definition.clone(),
fields: source_schema
.fields()
.iter()
.map(|f| super::without_declarations(f))
.collect(),
lineage: source_schema
.fields()
.iter()
.map(|f| (f.name().clone(), vec![f.name().clone()]))
.collect(),
inputs: source_schema
.fields()
.iter()
.map(|f| f.name().clone())
.collect(),
});
}
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect();
let mut definition = definition.clone();
let inputs = vec![config.column.clone()];
definition.projections = projections;
Ok(Planned {
definition,
fields: schema.fields().iter().map(|f| f.as_ref().clone()).collect(),
lineage: Default::default(),
inputs,
})
}
/// A saved native refresh plan, including every physical partition identity.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct DuplicatePairsRefreshPlan {
pub snapshot: GroupedRefreshPlan,
pub source_uri: String,
pub definition_sql: String,
pub operator_version: u32,
pub threshold_bits: u32,
pub tasks: Vec<DuplicatePairTask>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenPairsPartition {
pub(super) unit: u32,
pub(super) plan_id: String,
pub(super) task: DuplicatePairTask,
pub(super) fragments: Vec<Fragment>,
pub(super) rows: u64,
}
impl DuplicatePairsRefreshPlan {
pub(super) fn identity(&self) -> Result<String> {
use sha2::{Digest, Sha256};
let encoded = serde_json::to_vec(&(
&self.snapshot,
&self.source_uri,
&self.definition_sql,
self.operator_version,
self.threshold_bits,
))
.map_err(|e| invalid(e.to_string()))?;
Ok(format!("{:x}", Sha256::digest(encoded)))
}
}
async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let native = view
.as_native()
.ok_or_else(|| invalid("native pair refresh requires a local view"))?;
native.dataset.ensure_mutable()?;
native.dataset.reload().await?;
let view_ds = native.dataset.get().await?.as_ref().clone();
let Some(super::StoredDefinition::Query(definition)) =
super::read_definition(&view_ds.schema().metadata)?
else {
return Err(invalid("view does not carry a supported definition"));
};
let config = definition
.vector_source
.as_ref()
.ok_or_else(|| invalid("not a native pair view"))?
.config()?;
let source = Arc::new(
open_source(view, &definition, None)
.await?
.checkout_version(config.dataset_version)
.await?,
);
ensure_no_mem_wal(&source, "source table", &definition.source_table).await?;
let planned = plan(Arc::new(Schema::from(source.schema())), &definition)?;
let expected = Schema::new(planned.fields);
let physical = Schema::from(view_ds.schema());
if physical.fields().len() != expected.fields().len()
|| physical
.fields()
.iter()
.zip(expected.fields())
.any(|(a, b)| {
a.name() != b.name()
|| a.data_type() != b.data_type()
|| a.is_nullable() != b.is_nullable()
})
{
return Err(invalid("native pair view schema changed since creation"));
}
Ok((view_ds, planned.definition, source))
}
pub(super) async fn plan_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<DuplicatePairsRefreshPlan> {
let (view_ds, definition, source) = open(view).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if pinned.is_some_and(|v| v != config.dataset_version) {
return Err(invalid(
"refresh source version differs from the native pair view's fixed snapshot",
));
}
let tasks = plan_duplicate_pairs(source.clone(), &config).await?;
let units = u32::try_from(tasks.len()).map_err(|_| invalid("too many index partitions"))?;
let incarnation = view_ds
.schema()
.metadata
.get(super::INCARNATION_META_KEY)
.cloned()
.ok_or_else(|| invalid("native pair view has no incarnation"))?;
Ok(DuplicatePairsRefreshPlan {
snapshot: GroupedRefreshPlan {
units,
source_version: config.dataset_version,
view_version: view_ds.version().version,
incarnation,
},
source_uri: source.uri().to_string(),
definition_sql: definition.to_sql(),
operator_version: DUPLICATE_PAIRS_OPERATOR_VERSION,
threshold_bits: config.distance_threshold.to_bits(),
tasks,
})
}
pub(super) async fn open_planned(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let (view_ds, definition, source) = open(view).await?;
ensure_incarnation(&view_ds, Some(&plan.snapshot.incarnation), view.name()).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if source.uri() != plan.source_uri
|| source.version().version != plan.snapshot.source_version
|| definition.to_sql() != plan.definition_sql
|| config.distance_threshold.to_bits() != plan.threshold_bits
|| plan.operator_version != DUPLICATE_PAIRS_OPERATOR_VERSION
|| plan.tasks.len() != plan.snapshot.units as usize
{
return Err(invalid(
"native pair task belongs to a different snapshot, configuration or operator",
));
}
Ok((view_ds, definition, source))
}
pub(super) async fn write_unit(
view: &Table,
unit: u32,
plan: &DuplicatePairsRefreshPlan,
) -> Result<WrittenPairsPartition> {
let (view_ds, definition, source) = open_planned(view, plan).await?;
let task = plan
.tasks
.get(unit as usize)
.ok_or_else(|| invalid("pair partition is outside the plan"))?
.clone();
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
let exec = DuplicatePairsExec::try_new(source, config, vec![task.clone()])?;
let stream = exec.execute(0, Arc::new(TaskContext::default()))?;
let rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let count = rows.clone();
let stream = stream.inspect_ok(move |batch| {
count.fetch_add(
batch.num_rows() as u64,
std::sync::atomic::Ordering::Relaxed,
);
});
let stream: SendableRecordBatchStream = Box::pin(
datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
stream,
),
);
let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds)))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append { fragments } = txn.operation else {
return Err(invalid(
"native pair partition did not stage append fragments",
));
};
Ok(WrittenPairsPartition {
unit,
plan_id: plan.identity()?,
task,
fragments,
rows: rows.load(std::sync::atomic::Ordering::Relaxed),
})
}
pub(super) async fn commit(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
units: Vec<WrittenPairsPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let (_, definition, source) = open_planned(view, plan).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if plan_duplicate_pairs(source.clone(), &config).await? != plan.tasks {
return Err(invalid(
"pair plan does not cover the pinned index partitions exactly",
));
}
let identity = plan.identity()?;
let mut ordered = vec![None; plan.snapshot.units as usize];
for written in units {
if written.plan_id != identity
|| plan.tasks.get(written.unit as usize) != Some(&written.task)
{
return Err(invalid("pair partition receipt belongs to another plan"));
}
let slot = ordered
.get_mut(written.unit as usize)
.ok_or_else(|| invalid("unknown pair partition receipt"))?;
if slot.is_some() {
return Err(invalid("duplicate pair partition receipt"));
}
*slot = Some(written);
}
let units = ordered
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid("missing pair partition receipt"))?;
let rows = units.iter().map(|u| u.rows).sum();
let fragments = units.into_iter().flat_map(|u| u.fragments).collect();
super::grouped_units::commit_fragments(
view,
&plan.snapshot,
&definition,
source.manifest.timestamp_nanos,
fragments,
rows,
expected,
)
.await
}
+13 -122
View File
@@ -12,7 +12,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use arrow_array::cast::AsArray;
use arrow_array::{Array, FixedSizeListArray, UInt32Array};
use arrow_buffer::NullBuffer;
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef};
use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef};
use datafusion::catalog::default_table_source::provider_as_source;
use datafusion::datasource::empty::EmptyTable;
use datafusion::execution::FunctionRegistry;
@@ -36,20 +36,16 @@ use datafusion_sql::sqlparser::parser::Parser;
use futures::StreamExt;
use lance::Dataset;
use lance::index::{DatasetIndexExt, DatasetIndexInternalExt};
use lance_core::ROW_ID;
use lance_datafusion::exec::SessionContextExt;
use lance_index::metrics::NoOpMetricsCollector;
use lance_index::vector::VectorIndex;
use lance_index::vector::ivf::{IvfTransformer, new_ivf_transformer};
use lance_linalg::distance::DistanceType;
use lance_linalg::kernels::normalize_fsl;
use roaring::RoaringBitmap;
use uuid::Uuid;
use super::refresh::to_view_batch;
use super::{
MaterializedViewDefinition, Planned, SOURCE_ROW_ID_COLUMN, ViewProjection, ensure_immutable,
plan_filter, query, without_declarations,
MaterializedViewDefinition, Planned, ViewProjection, plan_filter, query, without_declarations,
};
use crate::{Error, Result};
@@ -68,50 +64,6 @@ pub(super) struct IvfBinding {
pub(super) transformer: Arc<IvfTransformer>,
/// A cosine index assigns L2 over unit vectors, as lance's index path does.
pub(super) normalize: bool,
/// Partitions of the model; the ids `ivf_partition` returns are below it.
pub(super) partitions: u32,
}
/// The segments of the one IVF index on a column, opened, and the fragments
/// they cover between them; every other fragment's rows are unindexed.
/// `covered` is `None` when a segment does not say which fragments it holds
/// (`fragment_bitmap: None`, which lance defines as unknown rather than
/// empty): its rows would be read once by the partition reader and again as
/// unindexed, so no caller may split this index.
pub(super) struct IvfSegments {
pub(super) segments: Vec<IvfSegment>,
pub(super) covered: Option<RoaringBitmap>,
}
/// One opened segment and the fragments it holds rows for. Its postings
/// still name rows of fragments it has since lost -- a column rewrite
/// attaches a new file and takes the fragment out of this bitmap -- so a
/// reader of its partitions keeps only the rows it still owns.
pub(super) struct IvfSegment {
pub(super) index: Arc<dyn VectorIndex>,
pub(super) fragments: Option<RoaringBitmap>,
}
impl IvfSegments {
pub(super) fn is_unindexed(&self, fragment_id: u64) -> Option<bool> {
Some(!self.covered.as_ref()?.contains(fragment_id as u32))
}
}
/// Coverage after folding in one more segment. Lance defines a segment
/// without a fragment bitmap as unknown coverage, not empty, and one such
/// segment is enough to lose the coverage of the whole index.
fn coverage(
covered: Option<RoaringBitmap>,
bitmap: Option<&RoaringBitmap>,
) -> Option<RoaringBitmap> {
match (covered, bitmap) {
(Some(mut covered), Some(bitmap)) => {
covered |= bitmap;
Some(covered)
}
_ => None,
}
}
/// `bindings` maps a source column to its index; planning runs unbound.
@@ -278,16 +230,6 @@ pub(super) async fn bind(
}
async fn bind_column(source: &Dataset, column: &str) -> Result<IvfBinding> {
bind_column_with_segments(source, column)
.await
.map(|(binding, _)| binding)
}
/// The binding plus the opened segments it was read from.
pub(super) async fn bind_column_with_segments(
source: &Dataset,
column: &str,
) -> Result<(IvfBinding, IvfSegments)> {
let field = source
.schema()
.field(column)
@@ -298,8 +240,6 @@ pub(super) async fn bind_column_with_segments(
// on its own fragments carries its own centroids, and grouping every row
// by one segment's model would bucket the other segments' rows wrongly.
let mut found: Option<(String, Uuid, DistanceType, FixedSizeListArray)> = None;
let mut segments = Vec::new();
let mut covered = Some(RoaringBitmap::new());
for index in source.load_indices().await?.iter() {
if index.fields != [field.id] {
continue;
@@ -314,12 +254,6 @@ pub(super) async fn bind_column_with_segments(
continue;
};
let metric = vector.metric_type();
covered = coverage(covered, index.fragment_bitmap.as_ref());
let fragments = index.fragment_bitmap.clone();
segments.push(IvfSegment {
index: vector.clone(),
fragments,
});
match &found {
None => found = Some((index.name.clone(), index.uuid, metric, centroids)),
Some((name, _, seen_metric, seen)) if *name == index.name => {
@@ -369,32 +303,20 @@ pub(super) async fn bind_column_with_segments(
// Lance's index path assigns cosine by L2 over unit vectors.
let normalize = metric == DistanceType::Cosine;
let distance = if normalize { DistanceType::L2 } else { metric };
let partitions = u32::try_from(centroids.len()).map_err(|_| Error::InvalidInput {
message: format!("{IVF_PARTITION}({column}): too many partitions"),
})?;
Ok((
IvfBinding {
index: uuid,
transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])),
normalize,
partitions,
},
IvfSegments { segments, covered },
))
Ok(IvfBinding {
index: uuid,
transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])),
normalize,
})
}
pub(super) fn empty_source(source_schema: &SchemaRef) -> Arc<dyn TableSource> {
let mut fields = source_schema.fields().to_vec();
fields.push(Arc::new(ArrowField::new(ROW_ID, DataType::UInt64, false)));
provider_as_source(Arc::new(EmptyTable::new(Arc::new(ArrowSchema::new(
fields,
)))))
provider_as_source(Arc::new(EmptyTable::new(source_schema.clone())))
}
/// The query DataFusion runs: the view's projections plus the group's
/// smallest source row id, which stands as the row's provenance. `filtered`
/// puts the view's predicate in the query, for a caller whose rows did not
/// come from a lance scan that already applied it.
/// The query DataFusion runs. `filtered` puts the view's predicate in the
/// query for a caller whose rows did not come from a Lance scan that already
/// applied it.
fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String {
let items: Vec<String> = definition
.projections
@@ -406,7 +328,7 @@ fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String {
_ => String::new(),
};
format!(
"SELECT {}, min({ROW_ID}) AS {ROW_ID} FROM {SOURCE}{predicate} GROUP BY {}",
"SELECT {} FROM {SOURCE}{predicate} GROUP BY {}",
items.join(", "),
definition.group_by.join(", ")
)
@@ -499,11 +421,6 @@ pub(super) fn plan(
};
let mut projections: Vec<ViewProjection> = Vec::with_capacity(definition.projections.len());
for p in &definition.projections {
if p.output == SOURCE_ROW_ID_COLUMN || p.output == ROW_ID {
return Err(Error::InvalidInput {
message: format!("view column name '{}' is reserved", p.output),
});
}
if projections.iter().any(|seen| seen.output == p.output) {
return Err(Error::ColumnAlreadyExists {
name: p.output.clone(),
@@ -549,14 +466,11 @@ pub(super) fn plan(
message: format!("invalid grouped view: {e}"),
})?;
for expr in &exprs {
ensure_immutable(expr, |message| Error::InvalidInput {
message: format!("invalid grouped view: {message}"),
})?;
inputs.extend(
expr.column_refs()
.into_iter()
.map(|c| c.name.clone())
.filter(|name| name != ROW_ID && source_schema.field_with_name(name).is_ok()),
.filter(|name| source_schema.field_with_name(name).is_ok()),
);
}
inputs.sort();
@@ -566,7 +480,6 @@ pub(super) fn plan(
.schema()
.fields()
.iter()
.filter(|f| f.name() != ROW_ID)
.map(|f| without_declarations(f))
.collect();
Ok(Planned {
@@ -586,7 +499,6 @@ pub(super) async fn stream(
rows_written: Arc<AtomicU64>,
) -> Result<SendableRecordBatchStream> {
let mut scanner = source.scan();
scanner.with_row_id();
if let Some(filter) = &definition.filter {
scanner.filter(filter)?;
}
@@ -611,24 +523,3 @@ pub(super) async fn stream(
});
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped)))
}
#[cfg(test)]
mod tests {
use super::*;
/// A segment that does not name the fragments it holds leaves the whole
/// index's coverage unknown: reading "everything it does not cover" would
/// read its rows a second time.
#[test]
fn one_segment_without_a_bitmap_loses_the_coverage() {
let first = RoaringBitmap::from_iter([0u32, 1]);
let second = RoaringBitmap::from_iter([2u32]);
let both = coverage(
coverage(Some(RoaringBitmap::new()), Some(&first)),
Some(&second),
);
assert_eq!(both, Some(RoaringBitmap::from_iter([0u32, 1, 2])));
assert_eq!(coverage(both, None), None);
assert_eq!(coverage(None, Some(&first)), None);
}
}
File diff suppressed because it is too large Load Diff
@@ -1,192 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Shared plan/write/commit surface for materialized views with independent units.
use serde::{Deserialize, Serialize};
use super::duplicate_pairs::{
self, DuplicatePairsRefreshPlan, VectorSourceKind, WrittenPairsPartition,
};
use super::vector_dedup::{self, VectorDedupPlan, WrittenDedupUnit};
use super::{GroupedRefreshPlan, RefreshMaterializedViewResult, StoredDefinition, WrittenUnit};
use crate::table::Table;
use crate::{Error, Result};
/// Serializable work description consumed by a host's existing job workers.
/// The grouped variant preserves the previous saved-plan representation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PartitionedRefreshPlan {
VectorDedup(VectorDedupPlan),
DuplicatePairs(DuplicatePairsRefreshPlan),
Grouped(GroupedRefreshPlan),
}
impl PartitionedRefreshPlan {
/// The shared source snapshot and destination generation fence.
pub fn snapshot(&self) -> &GroupedRefreshPlan {
match self {
Self::VectorDedup(plan) => &plan.snapshot,
Self::DuplicatePairs(plan) => &plan.snapshot,
Self::Grouped(plan) => plan,
}
}
/// Selection tasks that must complete before remaining tasks can start.
pub fn dependency_units(&self) -> u32 {
match self {
Self::VectorDedup(plan) => plan.dependencies(),
_ => 0,
}
}
/// Number of independent units in this immutable plan.
pub fn units(&self) -> u32 {
self.snapshot().units
}
}
/// One completed unit's uncommitted fragments. Failed attempts produce no receipt.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WrittenPartition {
VectorDedup(WrittenDedupUnit),
DuplicatePairs(WrittenPairsPartition),
Grouped(WrittenUnit),
}
impl WrittenPartition {
/// Number of pair or grouped rows staged by this unit.
pub fn rows(&self) -> u64 {
match self {
Self::VectorDedup(unit) => unit.rows,
Self::DuplicatePairs(unit) => unit.rows,
Self::Grouped(unit) => unit.rows(),
}
}
/// Uncommitted output files, retained only after task completion.
pub fn fragments(&self) -> &[lance_table::format::Fragment] {
match self {
Self::VectorDedup(unit) => &unit.fragments,
Self::DuplicatePairs(unit) => &unit.fragments,
Self::Grouped(unit) => unit.fragments(),
}
}
}
/// Plan native pair or indexed grouped work; return `None` for ordinary views.
/// `pinned`, when supplied, must match a native view's declared source version.
pub async fn plan_partitioned_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<Option<PartitionedRefreshPlan>> {
if let Some(StoredDefinition::Query(definition)) =
super::read_definition(view.schema().await?.metadata())?
&& definition.vector_source.is_some()
{
if definition
.vector_source
.as_ref()
.is_some_and(|s| s.kind == VectorSourceKind::Dedup)
{
return vector_dedup::plan_refresh(view, pinned)
.await
.map(|p| Some(PartitionedRefreshPlan::VectorDedup(p)));
}
return duplicate_pairs::plan_refresh(view, pinned)
.await
.map(|p| Some(PartitionedRefreshPlan::DuplicatePairs(p)));
}
super::plan_grouped_refresh(view, pinned)
.await
.map(|plan| plan.map(PartitionedRefreshPlan::Grouped))
}
/// Stream one unit into uncommitted fragments. Persist the returned receipt only
/// on success. A failed or interrupted call must be retried from the unit's start.
pub async fn write_refresh_partition(
view: &Table,
unit: u32,
plan: &PartitionedRefreshPlan,
) -> Result<WrittenPartition> {
write_refresh_partition_with_inputs(view, unit, plan, &[]).await
}
/// Write a unit after its durable selection dependencies have completed.
/// Supply exactly the first `plan.dependency_units()` receipts for a dependent
/// unit. A missing, duplicate or foreign receipt fails before output is staged.
pub async fn write_refresh_partition_with_inputs(
view: &Table,
unit: u32,
plan: &PartitionedRefreshPlan,
dependencies: &[WrittenPartition],
) -> Result<WrittenPartition> {
match plan {
PartitionedRefreshPlan::VectorDedup(plan) => {
let dependencies = dependencies
.iter()
.map(|unit| match unit {
WrittenPartition::VectorDedup(unit) => Ok(unit.clone()),
_ => Err(Error::InvalidInput {
message: "foreign dedup dependency".into(),
}),
})
.collect::<Result<Vec<_>>>()?;
vector_dedup::write_unit(view, unit, plan, &dependencies)
.await
.map(WrittenPartition::VectorDedup)
}
PartitionedRefreshPlan::DuplicatePairs(plan) => {
duplicate_pairs::write_unit(view, unit, plan)
.await
.map(WrittenPartition::DuplicatePairs)
}
PartitionedRefreshPlan::Grouped(plan) => super::write_grouped_unit(view, unit, plan)
.await
.map(WrittenPartition::Grouped),
}
}
/// Publish a complete set of unit receipts through the common MV commit sink.
/// `expected` fences the destination incarnation. No partial set is publishable.
pub async fn commit_partitioned_refresh(
view: &Table,
plan: &PartitionedRefreshPlan,
units: Vec<WrittenPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let foreign = || Error::InvalidInput {
message: "refresh receipt does not match the plan kind".into(),
};
match plan {
PartitionedRefreshPlan::VectorDedup(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::VectorDedup(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
vector_dedup::commit(view, plan, units, expected).await
}
PartitionedRefreshPlan::DuplicatePairs(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::DuplicatePairs(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
duplicate_pairs::commit(view, plan, units, expected).await
}
PartitionedRefreshPlan::Grouped(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::Grouped(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
super::commit_grouped_refresh(view, plan, units, expected).await
}
}
}
+1 -84
View File
@@ -264,36 +264,7 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let mut from = select.from.iter();
let mut vector_source = None;
let (source_namespace, source_table) = match from.next().map(|f| &f.relation) {
Some(TableFactor::Table {
name,
args: Some(args),
alias: None,
with_hints,
version: None,
with_ordinality: false,
partitions,
json_path: None,
sample: None,
index_hints,
..
}) if matches!(
name.to_string().as_str(),
super::duplicate_pairs::FUNCTION_NAME | super::duplicate_pairs::DEDUP_FUNCTION_NAME
) && with_hints.is_empty()
&& partitions.is_empty()
&& index_hints.is_empty() =>
{
let kind = if name.to_string() == super::duplicate_pairs::FUNCTION_NAME {
super::VectorSourceKind::Pairs
} else {
super::VectorSourceKind::Dedup
};
let (namespace, table, config) = super::duplicate_pairs::source(args, kind)?;
vector_source = Some(Box::new(config));
(namespace, table)
}
Some(TableFactor::Table { args: Some(_), .. }) => {
return Err(invalid(
"a view reads a table; a Function in FROM position follows it: \
@@ -445,7 +416,6 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let definition = MaterializedViewDefinition {
vector_source,
source_table,
source_namespace,
lateral,
@@ -454,9 +424,6 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
group_by,
limit,
};
if definition.vector_source.is_some() {
super::duplicate_pairs::check_shape(&definition)?;
}
check_grouping(&definition)?;
Ok(definition)
}
@@ -516,18 +483,7 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
.chain(std::iter::once(&definition.source_table))
.map(|part| ident_sql(part))
.collect();
if let Some(native) = &definition.vector_source {
sql.push_str(&format!(
"{}('{}', {}, '{}', {})",
native.function_name(),
table.join(".").replace('\'', "''"),
native.dataset_version,
native.column.replace('\'', "''"),
native.distance_threshold
));
} else {
sql.push_str(&table.join("."));
}
sql.push_str(&table.join("."));
if let Some(lateral) = &definition.lateral {
match &lateral.source {
LateralSource::Unnest { column } => {
@@ -555,45 +511,6 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
mod tests {
use super::*;
#[test]
fn native_pairs_are_a_table_source_with_literal_snapshot_arguments() {
for threshold in ["4", "4.0", "-1", "0.05"] {
let sql =
format!("SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', {threshold})");
let definition = parse(&sql).unwrap();
assert_eq!(definition.source_table, "images");
assert!(definition.lateral.is_none());
assert_eq!(
definition.vector_source.as_ref().unwrap().dataset_version,
3
);
assert_eq!(parse(&render(&definition)).unwrap(), definition);
let metadata = std::collections::HashMap::from([(
super::super::DEFINITION_META_KEY.to_string(),
definition.to_json().unwrap(),
)]);
assert_eq!(
super::super::read_definition(&metadata).unwrap(),
Some(super::super::StoredDefinition::Query(definition))
);
}
for sql in [
"SELECT row_id_a FROM vector_duplicate_pairs('images', 3, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 0, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 1e100)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', cutoff)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) LIMIT 1",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) WHERE distance < 2",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) AS p",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4), f(phash) AS p",
] {
assert!(
parse(sql).is_err(),
"accepted unsupported native view: {sql}"
);
}
}
#[test]
fn a_query_round_trips_through_its_canonical_form() {
for (sql, canonical) in [
File diff suppressed because it is too large Load Diff
@@ -1,636 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Two dependent waves in the ordinary refresh fleet: index partition selection,
//! then source fragment materialization. Only the latter publishes data files.
use std::collections::BTreeMap;
use std::io::Cursor;
use std::sync::Arc;
use arrow_array::{Array, RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion::prelude::{SessionConfig, SessionContext, col};
use datafusion_common::DataFusionError;
use datafusion_execution::{
SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder,
};
use datafusion_physical_plan::{ExecutionPlan, stream::RecordBatchStreamAdapter};
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::rowids::{get_row_id_index, load_row_id_sequence};
use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap};
use lance::dataset::transaction::Operation;
use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams};
use lance_core::datatypes::BlobHandling;
use lance_core::utils::address::RowAddress;
use lance_datafusion::exec::{OneShotExec, SessionContextExt};
use lance_table::format::Fragment;
use roaring::{RoaringBitmap, RoaringTreemap};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;
use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, VectorSourceKind};
use super::{GroupedRefreshPlan, RefreshMaterializedViewResult};
use crate::table::Table;
use crate::table::datafusion::udtf::duplicate_pairs::{DuplicatePairsExec, plan_duplicate_pairs};
use crate::{Error, Result};
const POLICY_VERSION: u32 = 1;
const SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MATERIALIZE_BATCH_ROWS: usize = 1024;
const MASK_FILE: &str = "dedup.bitmap";
const MASK_MAGIC: &[u8] = b"LANCEDB_DEDUP_1";
fn invalid(message: impl Into<String>) -> Error {
Error::InvalidInput {
message: message.into(),
}
}
fn io_error(error: std::io::Error) -> Error {
invalid(format!("invalid dedup bitmap: {error}"))
}
/// Immutable snapshot, partition tasks and source fragments for both waves.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct VectorDedupPlan {
pub snapshot: GroupedRefreshPlan,
native: DuplicatePairsRefreshPlan,
source_fragments: Vec<u32>,
policy_version: u32,
}
impl VectorDedupPlan {
pub(super) fn dependencies(&self) -> u32 {
self.native.snapshot.units
}
fn identity(&self) -> Result<String> {
Ok(format!(
"{:x}",
Sha256::digest(serde_json::to_vec(self).map_err(|e| invalid(e.to_string()))?)
))
}
}
/// Small receipt for a durable selection mask. Row IDs and pairs are not stored
/// in the host's task registry. UUID directories use ordinary index-file GC.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct MaskFile {
source_fragment: u32,
object_id: Uuid,
bytes: usize,
digest: String,
}
/// A selection task produces masks; a materialization task produces data files.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenDedupUnit {
unit: u32,
plan_id: String,
masks: Vec<MaskFile>,
pub(super) fragments: Vec<Fragment>,
pub(super) rows: u64,
}
pub(super) async fn plan_refresh(view: &Table, pinned: Option<u64>) -> Result<VectorDedupPlan> {
let native = duplicate_pairs::plan_refresh(view, pinned).await?;
let (_, _, source) = duplicate_pairs::open_planned(view, &native).await?;
let source_fragments = source
.get_fragments()
.iter()
.map(|f| f.id() as u32)
.collect::<Vec<_>>();
let mut snapshot = native.snapshot.clone();
snapshot.units = snapshot
.units
.checked_add(
u32::try_from(source_fragments.len())
.map_err(|_| invalid("too many source fragments"))?,
)
.ok_or_else(|| invalid("too many dedup refresh units"))?;
Ok(VectorDedupPlan {
snapshot,
native,
source_fragments,
policy_version: POLICY_VERSION,
})
}
async fn open(
view: &Table,
plan: &VectorDedupPlan,
) -> Result<(Dataset, super::MaterializedViewDefinition, Arc<Dataset>)> {
let (view_ds, definition, source) = duplicate_pairs::open_planned(view, &plan.native).await?;
let mut snapshot = plan.native.snapshot.clone();
snapshot.units = snapshot
.units
.checked_add(
u32::try_from(plan.source_fragments.len())
.map_err(|_| invalid("too many dedup fragments"))?,
)
.ok_or_else(|| invalid("too many dedup units"))?;
if plan.policy_version != POLICY_VERSION
|| snapshot != plan.snapshot
|| definition.vector_source.as_ref().map(|s| s.kind) != Some(VectorSourceKind::Dedup)
|| source
.get_fragments()
.iter()
.map(|f| f.id() as u32)
.collect::<Vec<_>>()
!= plan.source_fragments
{
return Err(invalid("dedup task belongs to another snapshot or policy"));
}
Ok((view_ds, definition, source))
}
fn edge_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt64, false),
Field::new("b", DataType::UInt64, false),
]))
}
/// Sorting canonical edges makes the first unremoved endpoint the smallest
/// unassigned representative. The sort spills; only vertex masks stay in RAM.
pub async fn select_representatives(
stream: SendableRecordBatchStream,
memory: usize,
) -> Result<(RoaringTreemap, RoaringTreemap)> {
let schema = edge_schema();
let batch_schema = schema.clone();
let canonical = stream
.map_ok(move |batch| {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
RecordBatch::try_new(
batch_schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)),
)),
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)),
)),
],
)
.map_err(DataFusionError::from)
})
.and_then(futures::future::ready);
let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical));
let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS);
// Leave enough room to merge spilled runs while the final in-memory run
// is still resident. Small batches also bound each merge input buffer.
config.options_mut().execution.sort_spill_reservation_bytes = memory / 2;
let ctx = SessionContext::new_with_config_rt(
config,
RuntimeEnvBuilder::new()
.with_memory_limit(memory, 1.0)
.build_arc()?,
);
let mut sorted = ctx
.read_one_shot(canonical)?
.sort_by(vec![col("a"), col("b")])?
.execute_stream()
.await?;
let mut removed = RoaringTreemap::new();
let mut representatives = RoaringTreemap::new();
while let Some(batch) = sorted.try_next().await? {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
for (&a, &b) in a.values().iter().zip(b.values()) {
if a != b && !removed.contains(a) && !removed.contains(b) {
representatives.insert(a);
removed.insert(b);
}
}
}
Ok((removed, representatives))
}
#[derive(Default)]
struct Masks {
removed: RoaringBitmap,
representatives: RoaringBitmap,
}
async fn persist_masks(
view: &Dataset,
source: &Dataset,
removed: RoaringTreemap,
representatives: RoaringTreemap,
) -> Result<Vec<MaskFile>> {
let row_index = get_row_id_index(source).await?;
let mut masks: BTreeMap<u32, Masks> = BTreeMap::new();
for (ids, is_removed) in [(removed, true), (representatives, false)] {
for id in ids {
let address = match &row_index {
Some(index) => index
.get(id)?
.ok_or_else(|| invalid("pair row id missing from pinned source"))?,
None => RowAddress::new_from_u64(id),
};
let mask = masks.entry(address.fragment_id()).or_default();
if is_removed {
mask.removed.insert(address.row_offset());
} else {
mask.representatives.insert(address.row_offset());
}
}
}
let store = view.object_store(None).await?;
let mut files = Vec::new();
for (source_fragment, mask) in masks {
if source.get_fragment(source_fragment as usize).is_none() {
return Err(invalid("pair row address is outside the pinned snapshot"));
}
let mut bytes = MASK_MAGIC.to_vec();
mask.removed.serialize_into(&mut bytes).map_err(io_error)?;
mask.representatives
.serialize_into(&mut bytes)
.map_err(io_error)?;
let object_id = Uuid::new_v4();
store
.put(
&view
.indices_dir()
.join(object_id.to_string())
.join(MASK_FILE),
&bytes,
)
.await?;
files.push(MaskFile {
source_fragment,
object_id,
bytes: bytes.len(),
digest: format!("{:x}", Sha256::digest(&bytes)),
});
}
Ok(files)
}
fn validate_receipts<'a>(
plan: &VectorDedupPlan,
receipts: &'a [WrittenDedupUnit],
count: u32,
) -> Result<Vec<&'a WrittenDedupUnit>> {
let identity = plan.identity()?;
let mut ordered = vec![None; count as usize];
for receipt in receipts {
if receipt.plan_id != identity {
return Err(invalid("dedup receipt belongs to another plan"));
}
let slot = ordered
.get_mut(receipt.unit as usize)
.ok_or_else(|| invalid("unexpected dedup receipt"))?;
if slot.is_some() {
return Err(invalid("duplicate dedup receipt"));
}
if receipt.unit < plan.dependencies() {
if !receipt.fragments.is_empty()
|| receipt.rows != 0
|| receipt
.masks
.iter()
.any(|m| !plan.source_fragments.contains(&m.source_fragment))
{
return Err(invalid("invalid selection receipt"));
}
} else if !receipt.masks.is_empty() {
return Err(invalid("invalid materialization receipt"));
}
*slot = Some(receipt);
}
ordered
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid("missing dedup dependency receipt"))
}
async fn load_masks(
view: &Dataset,
source_fragment: u32,
dependencies: &[&WrittenDedupUnit],
) -> Result<RoaringBitmap> {
let store = view.object_store(None).await?;
let mut masks = Masks::default();
for file in dependencies
.iter()
.flat_map(|u| &u.masks)
.filter(|m| m.source_fragment == source_fragment)
{
let reader = store
.open(
&view
.indices_dir()
.join(file.object_id.to_string())
.join(MASK_FILE),
)
.await?;
if reader.size().await? != file.bytes {
return Err(invalid("dedup mask length mismatch"));
}
let bytes = reader.get_range(0..file.bytes).await?;
if !bytes.starts_with(MASK_MAGIC) || format!("{:x}", Sha256::digest(&bytes)) != file.digest
{
return Err(invalid("dedup mask checksum mismatch"));
}
let mut cursor = Cursor::new(&bytes[MASK_MAGIC.len()..]);
masks.removed |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?;
masks.representatives |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?;
if cursor.position() as usize != bytes.len() - MASK_MAGIC.len() {
return Err(invalid("dedup mask has trailing data"));
}
}
// Native index scopes normally have disjoint rows. Fail closed if an
// overlapping scope would delete a representative used by another scope.
if !masks.removed.is_disjoint(&masks.representatives) {
return Err(invalid(
"overlapping index scopes disagree on retained representatives",
));
}
Ok(masks.removed)
}
pub(super) async fn write_unit(
view: &Table,
unit: u32,
plan: &VectorDedupPlan,
dependencies: &[WrittenDedupUnit],
) -> Result<WrittenDedupUnit> {
let (view_ds, definition, source) = open(view, plan).await?;
let mut written = WrittenDedupUnit {
unit,
plan_id: plan.identity()?,
masks: Vec::new(),
fragments: Vec::new(),
rows: 0,
};
if unit < plan.dependencies() {
let config = definition
.vector_source
.as_ref()
.expect("dedup source")
.config()?;
let exec = DuplicatePairsExec::try_new(
source.clone(),
config,
vec![plan.native.tasks[unit as usize].clone()],
)?;
let stream = exec.execute(0, Arc::new(TaskContext::default()))?;
let (removed, representatives) = select_representatives(stream, SORT_MEMORY_BYTES).await?;
written.masks = persist_masks(&view_ds, &source, removed, representatives).await?;
return Ok(written);
}
let source_fragment = *plan
.source_fragments
.get((unit - plan.dependencies()) as usize)
.ok_or_else(|| invalid("dedup unit is outside the plan"))?;
let dependencies = validate_receipts(plan, dependencies, plan.dependencies())?;
let removed = load_masks(&view_ds, source_fragment, &dependencies).await?;
let fragment = source
.get_fragment(source_fragment as usize)
.ok_or_else(|| invalid("source fragment missing"))?;
// The scanner prefilter uses snapshot row IDs. Our persisted masks use
// physical offsets so each writer only loads its own fragment's artifacts.
let blocked = if source.manifest.uses_stable_row_ids() {
let sequence = load_row_id_sequence(&source, fragment.metadata()).await?;
if removed
.max()
.is_some_and(|offset| offset as u64 >= sequence.len())
{
return Err(invalid("dedup mask offset is outside the source fragment"));
}
RowAddrTreeMap::from_iter(sequence.select(removed.iter().map(|offset| offset as usize)))
} else {
RowAddrTreeMap::from_iter(
removed
.iter()
.map(|offset| u64::from(RowAddress::new_from_parts(source_fragment, offset))),
)
};
let mut scanner = source.scan();
scanner
.with_fragments(vec![fragment.metadata().clone()])
.with_row_addr_prefilter(RowAddrMask::from_block(blocked))
.blob_handling(BlobHandling::AllBinary)
.batch_size(MATERIALIZE_BATCH_ROWS)
.strict_batch_size(true);
let stream = scanner.try_into_stream().await?;
let schema = lance_io::stream::RecordBatchStream::schema(&stream);
let rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let count = rows.clone();
let stream = stream
.map_err(|e| DataFusionError::External(Box::new(e)))
.inspect_ok(move |batch| {
count.fetch_add(
batch.num_rows() as u64,
std::sync::atomic::Ordering::Relaxed,
);
});
let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(schema, stream));
// Materialized binary payloads must become the destination's logical blob
// structs before the Lance writer encodes them. Reuse the normal add path's
// coercions, including nulls and nested fields.
let coerced = crate::table::datafusion::cast::cast_to_table_schema(
Arc::new(OneShotExec::new(stream)),
&Schema::from(view_ds.schema()),
)?;
let stream = coerced.execute(0, Arc::new(TaskContext::default()))?;
let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds)))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append { fragments } = txn.operation else {
return Err(invalid(
"dedup materialization did not stage append fragments",
));
};
written.fragments = fragments;
written.rows = rows.load(std::sync::atomic::Ordering::Relaxed);
Ok(written)
}
pub(super) async fn commit(
view: &Table,
plan: &VectorDedupPlan,
units: Vec<WrittenDedupUnit>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let (_, definition, source) = open(view, plan).await?;
let config = definition
.vector_source
.as_ref()
.expect("dedup source")
.config()?;
if plan_duplicate_pairs(source.clone(), &config).await? != plan.native.tasks {
return Err(invalid("dedup plan does not cover the pinned index"));
}
validate_receipts(plan, &units, plan.snapshot.units)?;
let rows = units.iter().map(|u| u.rows).sum();
let fragments = units.into_iter().flat_map(|u| u.fragments).collect();
super::grouped_units::commit_fragments(
view,
&plan.snapshot,
&definition,
source.manifest.timestamp_nanos,
fragments,
rows,
expected,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
fn edges(values: &[(u64, u64)], batch_size: usize) -> SendableRecordBatchStream {
let schema = edge_schema();
let batches = values
.chunks(batch_size)
.map(|edges| {
Ok(RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.0))),
Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.1))),
],
)
.unwrap())
})
.collect::<Vec<_>>();
Box::pin(RecordBatchStreamAdapter::new(
schema,
futures::stream::iter(batches),
))
}
#[tokio::test]
async fn direct_representatives_handle_chains_orientation_and_uint64() {
let high = u64::MAX - 2;
let pairs = [
(2, 1),
(1, 0),
(high + 2, high + 1),
(high, high + 1),
(1, 0),
];
let (removed, representatives) =
select_representatives(edges(&pairs, 1), SORT_MEMORY_BYTES)
.await
.unwrap();
assert_eq!(removed.iter().collect::<Vec<_>>(), [1, high + 1]);
assert_eq!(representatives.iter().collect::<Vec<_>>(), [0, high]);
}
#[tokio::test]
async fn masks_reject_overlapping_representatives_and_corrupt_artifacts() {
let dir = tempfile::tempdir().unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt64, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))],
)
.unwrap();
let ds = Dataset::write(
arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema),
dir.path().join("source.lance").to_str().unwrap(),
None,
)
.await
.unwrap();
let first = persist_masks(
&ds,
&ds,
[1].into_iter().collect(),
[0].into_iter().collect(),
)
.await
.unwrap();
let second = persist_masks(
&ds,
&ds,
[2].into_iter().collect(),
[1].into_iter().collect(),
)
.await
.unwrap();
let receipt = |masks| WrittenDedupUnit {
unit: 0,
plan_id: "test".into(),
masks,
fragments: vec![],
rows: 0,
};
let first = receipt(first);
let second = receipt(second);
assert_eq!(
load_masks(&ds, 0, &[&first]).await.unwrap(),
[1].into_iter().collect()
);
assert!(
load_masks(&ds, 0, &[&first, &second])
.await
.unwrap_err()
.to_string()
.contains("overlapping")
);
let file = &first.masks[0];
ds.object_store(None)
.await
.unwrap()
.put(
&ds.indices_dir()
.join(file.object_id.to_string())
.join(MASK_FILE),
&vec![0; file.bytes],
)
.await
.unwrap();
assert!(
load_masks(&ds, 0, &[&first])
.await
.unwrap_err()
.to_string()
.contains("checksum")
);
}
#[tokio::test]
async fn external_sort_handles_more_edges_than_its_budget() {
let pairs = (0..262144)
.rev()
.map(|i| ((i + 1) as u64, i as u64))
.collect::<Vec<_>>();
let (removed, reps) = select_representatives(edges(&pairs, 4096), 2 * 1024 * 1024)
.await
.unwrap();
assert_eq!(removed.len(), 131072);
assert_eq!(reps.len(), 131072);
assert!(removed.iter().all(|id| id % 2 == 1));
assert!(reps.iter().all(|id| id % 2 == 0));
}
}
+88 -5
View File
@@ -11,7 +11,7 @@ use arrow_array::{
Array, Float16Array, Float32Array, Float64Array, RecordBatch, UInt64Array,
cast::AsArray,
make_array,
types::{Int64Type, UInt64Type},
types::{Float32Type, Int64Type, UInt64Type},
};
use arrow_schema::{DataType, SchemaRef};
use datafusion_common::{DataFusionError, Result as DataFusionResult};
@@ -1024,6 +1024,8 @@ impl Query {
///
/// By default, there is no embedding model, and the input should be
/// vector/slice of floats.
/// NaN and infinite components are rejected with [`Error::InvalidInput`]
/// after conversion to Float32, for both local and remote tables.
///
/// If there is only one vector column (a column whose data type is a
/// fixed size list of floats) then the column does not need to be specified.
@@ -1050,8 +1052,7 @@ impl Query {
/// * `vector` - The vector that will be used for search.
pub fn nearest_to(self, vector: impl IntoQueryVector) -> Result<VectorQuery> {
let mut vector_query = self.into_vector();
let query_vector = vector.to_query_vector(&DataType::Float32, "default")?;
vector_query.request.query_vector.push(query_vector);
vector_query.request.add_query_vector(vector)?;
if vector_query.request.base.limit.is_none() {
vector_query.request.base.limit = Some(DEFAULT_TOP_K);
@@ -1154,6 +1155,23 @@ impl Default for VectorQueryRequest {
}
impl VectorQueryRequest {
fn add_query_vector(&mut self, vector: impl IntoQueryVector) -> Result<()> {
let vector = vector.to_query_vector(&DataType::Float32, "default")?;
let array =
vector
.as_primitive_opt::<Float32Type>()
.ok_or_else(|| Error::InvalidInput {
message: "VectorQuery vector must be of type Float32".into(),
})?;
if array.values().iter().any(|value| !value.is_finite()) {
return Err(Error::InvalidInput {
message: "query vector must contain only finite values".into(),
});
}
self.query_vector.push(vector);
Ok(())
}
pub fn from_plain_query(query: QueryRequest) -> Self {
Self {
base: query,
@@ -1222,9 +1240,11 @@ impl VectorQuery {
/// The output data will contain an additional column `query_index` which
/// will contain the index of the query vector that was used to generate the
/// result.
///
/// NaN and infinite components are rejected with [`Error::InvalidInput`]
/// after conversion to Float32, for both local and remote tables.
pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result<Self> {
let query_vector = vector.to_query_vector(&DataType::Float32, "default")?;
self.request.query_vector.push(query_vector);
self.request.add_query_vector(vector)?;
Ok(self)
}
@@ -2336,6 +2356,69 @@ mod tests {
use crate::{Table, connect, database::CreateTableMode, index::Index};
#[tokio::test]
async fn test_query_vector_non_finite() {
let conn = connect("memory://non_finite_vectors")
.execute()
.await
.unwrap();
let table = conn
.create_table("my_table", make_test_batches())
.execute()
.await
.unwrap();
let assert_invalid = |result: Result<VectorQuery>| {
let err = result.unwrap_err();
assert!(matches!(err, Error::InvalidInput { .. }));
assert!(err.to_string().contains("only finite values"));
};
// Float64 values can be finite before conversion but overflow Float32.
for value in [
f64::NAN,
f64::INFINITY,
f64::NEG_INFINITY,
f64::MAX,
f64::MIN,
] {
assert_invalid(table.query().nearest_to(&[0.1, value as f32]));
assert_invalid(table.query().nearest_to(vec![0.1, value]));
assert_invalid(
table
.query()
.nearest_to(&[0.1, 0.2])
.unwrap()
.add_query_vector(vec![0.1, value]),
);
let arrays: [Arc<dyn Array>; 3] = [
Arc::new(Float16Array::from(vec![
f16::from_f32(0.1),
f16::from_f64(value),
])),
Arc::new(Float32Array::from(vec![0.1, value as f32])),
Arc::new(Float64Array::from(vec![0.1, value])),
];
for array in arrays {
assert_invalid(table.query().nearest_to(array.clone()));
assert_invalid(
table
.query()
.nearest_to(&[0.1, 0.2])
.unwrap()
.add_query_vector(array),
);
}
}
table
.query()
.nearest_to(&[f32::MIN, f32::MAX])
.unwrap()
.add_query_vector(&[0.0, -0.0])
.unwrap();
}
#[tokio::test]
async fn test_setters_getters() {
let batches = make_test_batches();
+124
View File
@@ -416,6 +416,7 @@ mod tests {
RemoteCatalog::try_new(&endpoint, options).unwrap(),
));
let page = catalog
.catalog()
.list_databases(ListDatabasesRequest::default().limit(1).page_token("a/b"))
.await
.unwrap();
@@ -443,6 +444,7 @@ mod tests {
.unwrap();
assert!(
catalog
.catalog()
.list_databases(ListDatabasesRequest::default())
.await
.unwrap()
@@ -538,6 +540,128 @@ mod tests {
);
}
#[tokio::test]
async fn database_names_follow_pages_lazily() {
use futures::TryStreamExt;
let (endpoint, task) = server(vec![
(
200,
json!({"namespaces": ["first", "second"], "page_token": "a/b"}),
),
(200, json!({"namespaces": [], "page_token": "next"})),
(200, json!({"namespaces": ["last"], "page_token": ""})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(None, None);
drop(catalog);
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), None);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
assert_eq!(names.num_page_results(), 1);
assert_eq!(names.page_token(), Some("a/b"));
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), Some("a/b"));
assert!(!task.is_finished());
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("last"));
assert_eq!(names.try_next().await.unwrap(), None);
assert_eq!(names.try_next().await.unwrap(), None);
let requests = task.await.unwrap();
assert_eq!(requests[0].line, "GET /v1/namespace/$/list HTTP/1.1");
assert_eq!(
requests[1].line,
"GET /v1/namespace/$/list?page_token=a%2Fb HTTP/1.1"
);
assert_eq!(
requests[2].line,
"GET /v1/namespace/$/list?page_token=next HTTP/1.1"
);
}
#[tokio::test]
async fn database_names_stop_after_error() {
use futures::StreamExt;
let (endpoint, task) = server(vec![
(200, json!({"namespaces": ["first"], "page_token": "next"})),
(401, json!({"error": "unauthorized"})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(None, None);
assert_eq!(names.next().await.unwrap().unwrap(), "first");
assert!(names.next().await.unwrap().is_err());
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.page_token(), Some("next"));
assert!(names.next().await.is_none());
task.await.unwrap();
}
#[tokio::test]
async fn database_names_can_resume_with_page_limit() {
use futures::TryStreamExt;
let (endpoint, task) = server(vec![
(
200,
json!({"namespaces": ["first", "second"], "page_token": "next/token"}),
),
(200, json!({"namespaces": ["third", "fourth"]})),
])
.await;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new(endpoint, RemoteCatalogOptions::default()).unwrap(),
));
let mut names = catalog.list_databases(Some("start/token".into()), Some(2));
assert_eq!(names.page_token(), Some("start/token"));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("first"));
assert_eq!(names.page_token(), Some("next/token"));
assert_eq!(names.num_page_results(), 1);
assert_eq!(names.try_next().await.unwrap().as_deref(), Some("second"));
assert_eq!(names.num_page_results(), 0);
let token = names.page_token().map(str::to_owned);
drop(names);
let mut resumed = catalog.list_databases(token, Some(2));
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("third"));
assert_eq!(resumed.page_token(), None);
assert_eq!(resumed.num_page_results(), 1);
assert_eq!(resumed.try_next().await.unwrap().as_deref(), Some("fourth"));
assert_eq!(resumed.num_page_results(), 0);
assert_eq!(resumed.try_next().await.unwrap(), None);
let requests = task.await.unwrap();
assert_eq!(
requests[0].line,
"GET /v1/namespace/$/list?limit=2&page_token=start%2Ftoken HTTP/1.1"
);
assert_eq!(
requests[1].line,
"GET /v1/namespace/$/list?limit=2&page_token=next%2Ftoken HTTP/1.1"
);
}
#[tokio::test]
async fn database_names_reject_invalid_page_limit() {
use futures::TryStreamExt;
let catalog = CatalogConnection::new(Arc::new(
RemoteCatalog::try_new("http://127.0.0.1:1", RemoteCatalogOptions::default()).unwrap(),
));
for limit in [0, u32::MAX] {
let mut names = catalog.list_databases(None, Some(limit));
assert!(matches!(
names.try_next().await,
Err(Error::InvalidInput { .. })
));
assert_eq!(names.num_page_results(), 0);
assert_eq!(names.try_next().await.unwrap(), None);
}
}
#[tokio::test]
async fn catalog_validates_before_sending_requests() {
for endpoint in [
+120 -62
View File
@@ -47,6 +47,10 @@ use super::table::RemoteTable;
use super::util::parse_server_version;
use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id};
fn quote_sql_identifier(identifier: &str) -> String {
format!("\"{}\"", identifier.replace('"', "\"\""))
}
// Request structure for the remote clone table API
#[derive(serde::Serialize)]
struct RemoteCloneTableRequest {
@@ -860,42 +864,22 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
&self,
request: CreateMaterializedViewRequest,
) -> Result<Job> {
let identifier = build_table_identifier(&request.name, &request.namespace_path)?;
let req = self
.client
.post(&format!("/v1/materialized_view/{identifier}/create"))
.json(&serde_json::json!({
"query": request.query,
"with_no_data": request.with_no_data,
}));
let (request_id, response) = self.client.send(req).await?;
let response = self.client.check_response(&request_id, response).await?;
let status = response.status();
let body = response.text().await.err_to_http(request_id.clone())?;
let job_id = extract_job_id(&body);
if request.with_no_data {
return Ok(match job_id {
Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))),
None => Job::new_done(),
return Err(Error::NotSupported {
message: "remote materialized views are always populated by their SQL job"
.to_string(),
});
}
if status != StatusCode::ACCEPTED {
return Err(Error::Http {
source: "materialized-view creation with data must return 202 Accepted".into(),
request_id,
status_code: Some(status),
});
}
let job_id = job_id.ok_or_else(|| Error::Http {
source: "materialized-view creation response did not contain a valid job_id".into(),
request_id,
status_code: Some(status),
let client = self.sql_client.clone().ok_or_else(|| Error::NotSupported {
message: "SQL is unavailable for this remote database client".to_string(),
})?;
Ok(Job::new(Box::new(RemoteJob::new(
self.client.clone(),
job_id,
))))
let namespace = request.namespace_path;
let statement = format!(
"CREATE MATERIALIZED VIEW {} AS {}",
quote_sql_identifier(&request.name),
request.query
);
Ok(client.submit_as_job(statement, namespace, || async { Ok(()) }))
}
async fn drop_materialized_view_async(
@@ -1491,6 +1475,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
}
async fn create_table(&self, mut request: CreateTableRequest) -> Result<Arc<dyn BaseTable>> {
let data_schema = request.data.schema();
let body = stream_as_body(request.data.scan_as_stream())?;
let identifier = build_table_identifier(&request.name, &request.namespace_path)?;
@@ -1521,7 +1506,17 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
managed_versioning: None,
};
let req = (callback)(req);
self.open_table(req).await
let table = self.open_table(req).await?;
let table_schema = table.schema().await?;
if table_schema.as_ref() != data_schema.as_ref() {
return Err(Error::Schema {
message: "Provided schema does not match existing table schema"
.to_string(),
});
}
Ok(table)
}
// This should not happen, as we explicitly set the mode to overwrite and the server
@@ -1547,12 +1542,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let version = parse_server_version(&request_id, &rsp)?;
let table_identifier = build_table_identifier(&request.name, &request.namespace_path)?;
let cache_key = build_cache_key(&request.name, &request.namespace_path);
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name.clone(),
request.namespace_path.clone(),
table_identifier,
version.clone(),
self.sql_client.clone(),
));
self.table_cache.insert(cache_key, version).await;
@@ -1589,12 +1585,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let version = parse_server_version(&request_id, &rsp)?;
let cache_key = build_cache_key(&request.target_table_name, &request.target_namespace_path);
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.target_table_name.clone(),
request.target_namespace_path.clone(),
table_identifier,
version.clone(),
self.sql_client.clone(),
));
self.table_cache.insert(cache_key, version).await;
@@ -1607,12 +1604,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
// Every open gets its own checkout, schema cache, and freshness state.
if let Some(version) = self.table_cache.get(&cache_key).await {
Ok(Arc::new(RemoteTable::new(
Ok(Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name,
request.namespace_path,
identifier,
version,
self.sql_client.clone(),
)))
} else {
// Describe the table to confirm it exists before moving on.
@@ -1625,12 +1623,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let rsp = self.client.check_response(&request_id, rsp).await?;
let version = parse_server_version(&request_id, &rsp)?;
let describe_body = rsp.text().await.ok();
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name.clone(),
request.namespace_path.clone(),
identifier,
version.clone(),
self.sql_client.clone(),
));
// This describe already carries the schema, so hand it to the table
// instead of making the first schema read fetch it again. A version or
@@ -1913,6 +1912,7 @@ mod tests {
use arrow_array::{Int32Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use lance_namespace_impls::{DynamicContextProvider, OperationInfo};
use rstest::rstest;
use crate::connection::ConnectBuilder;
use crate::database::Database;
@@ -1957,29 +1957,11 @@ mod tests {
}
#[tokio::test]
async fn test_create_materialized_view_uses_item_route_and_job() {
let db = super::RemoteDatabase::new_mock(|request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/materialized_view/analytics$adults/create"
);
let body = request
.body()
.and_then(reqwest::Body::as_bytes)
.and_then(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).ok())
.unwrap();
assert_eq!(
body["query"],
"SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10"
);
assert_eq!(body["with_no_data"], false);
http::Response::builder()
.status(202)
.body(serde_json::json!({"job_id": "j1-mv-create"}).to_string())
.unwrap()
async fn test_create_materialized_view_requires_sql_client() {
let db = super::RemoteDatabase::new_mock(|request| -> http::Response<String> {
panic!("unexpected REST request: {}", request.url().path())
});
let job = db
let error = db
.create_materialized_view_async(CreateMaterializedViewRequest {
name: "adults".into(),
namespace_path: vec!["analytics".into()],
@@ -1988,8 +1970,8 @@ mod tests {
with_no_data: false,
})
.await
.unwrap();
assert_eq!(job.id(), Some("j1-mv-create"));
.unwrap_err();
assert!(error.to_string().contains("SQL is unavailable"));
}
#[tokio::test]
@@ -2538,6 +2520,75 @@ mod tests {
);
}
#[rstest]
#[case::matching(Field::new("a", DataType::Int32, false), true)]
#[case::different_name(Field::new("x", DataType::Int32, false), false)]
#[case::different_type(Field::new("a", DataType::Int64, false), false)]
#[case::different_nullability(Field::new("a", DataType::Int32, true), false)]
#[tokio::test]
async fn test_create_table_exist_ok_validates_schema(
#[case] existing_field: Field,
#[case] matches: bool,
#[values(false, true)] empty: bool,
#[values(false, true)] cached: bool,
) {
let existing_schema = Schema::new(vec![existing_field]);
let description = serde_json::json!({
"version": 1,
"schema": lance::arrow::json::JsonSchema::try_from(&existing_schema).unwrap(),
})
.to_string();
let mut db = super::RemoteDatabase::new_mock(move |request| {
assert_eq!(request.method(), &reqwest::Method::POST);
match request.url().path() {
"/v1/table/table1/create/" => {
assert_eq!(request.url().query(), Some("mode=exist_ok"));
http::Response::builder()
.status(400)
.body("Table table1 already exists".to_string())
.unwrap()
}
"/v1/table/table1/describe/" => http::Response::builder()
.status(200)
.body(description.clone())
.unwrap(),
path => panic!("unexpected path: {path}"),
}
});
db.table_cache = moka::future::Cache::new(10);
let conn = Connection::new(
Arc::new(db),
Arc::new(crate::embeddings::MemoryRegistry::new()),
);
if cached {
conn.open_table("table1").execute().await.unwrap();
}
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
let builder = if empty {
conn.create_empty_table("table1", schema.clone())
} else {
let data = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
conn.create_table("table1", data)
};
let result = builder
.mode(CreateTableMode::exist_ok(|b| b))
.execute()
.await;
if matches {
let table = result.unwrap();
assert_eq!(table.name(), "table1");
assert_eq!(table.schema().await.unwrap(), schema);
} else {
assert!(matches!(result, Err(Error::Schema { message })
if message == "Provided schema does not match existing table schema"));
}
}
#[tokio::test]
async fn test_create_table_modes() {
let test_cases = [
@@ -2577,7 +2628,14 @@ mod tests {
.status(400)
.body("Table table1 already exists")
.unwrap(),
"/v1/table/table1/describe/" => http::Response::builder().status(200).body("").unwrap(),
"/v1/table/table1/describe/" => http::Response::builder()
.status(200)
.body(
r#"{"version": 1, "schema": {"fields": [
{"name": "a", "type": {"type": "int32"}, "nullable": false}
]}}"#,
)
.unwrap(),
_ => {
panic!("unexpected path: {:?}", request.url().path());
}
+40 -1
View File
@@ -3,6 +3,7 @@
use std::collections::HashMap;
use std::fs;
use std::future::{Future, pending};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::{Duration, Instant};
@@ -20,12 +21,14 @@ use arrow_schema::{Schema, SchemaRef};
use futures::TryStreamExt;
use http::header::{HeaderMap, HeaderName, HeaderValue};
use prost::Message;
use tokio::sync::{Mutex, Notify, OnceCell, mpsc};
use serde::{Serialize, de::DeserializeOwned};
use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch};
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use uuid::Uuid;
use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
use crate::error::{Error, Result};
use crate::job::Job;
use crate::remote::client::{ClientConfig, TlsConfig};
use crate::remote::retry::ResolvedRetryConfig;
use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus};
@@ -209,6 +212,42 @@ impl SqlClient {
.await
}
pub(super) fn submit_as_job<T, F, Fut>(
&self,
statement: String,
namespace: Vec<String>,
finish: F,
) -> Job<T>
where
T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
{
let client = self.clone();
let (query_tx, query_rx) = watch::channel(None::<Arc<Query>>);
let task = tokio::spawn(async move {
let query = Arc::new(client.submit(&statement, &namespace).await?);
query_tx.send_replace(Some(query.clone()));
let mut reader = query.reader().await?;
while reader.try_next().await?.is_some() {}
finish().await
});
Job::spawned_with_cancellation(task, move || {
let mut query_rx = query_rx.clone();
async move {
loop {
let query = query_rx.borrow().clone();
if let Some(query) = query {
return query.cancel().await;
}
if query_rx.changed().await.is_err() {
pending::<()>().await;
}
}
}
})
}
pub(super) async fn describe(&self, query_id: Uuid) -> Result<QueryDescription> {
let query = self
.queries
+151
View File
@@ -21,7 +21,10 @@ use super::*;
use crate::database::Database;
use crate::remote::RemoteCatalogOptions;
use crate::remote::client::HeaderProvider;
use crate::remote::client::test_utils::client_with_handler;
use crate::remote::db::RemoteDatabase;
use crate::remote::table::RemoteTable;
use crate::table::BaseTable;
#[derive(Debug, Default)]
struct DelayedHeaderProvider {
@@ -73,6 +76,7 @@ struct TestSqlService {
first_continuation_count: Arc<AtomicUsize>,
transient_poll_failures: Arc<AtomicUsize>,
headers: Arc<std::sync::Mutex<Vec<CapturedHeaders>>>,
gate_refresh_release: Arc<Notify>,
result: RecordBatch,
large_result: RecordBatch,
dictionary_result: RecordBatch,
@@ -121,6 +125,7 @@ impl Default for TestSqlService {
first_continuation_count: Arc::new(AtomicUsize::new(0)),
transient_poll_failures: Arc::new(AtomicUsize::new(0)),
headers: Arc::new(std::sync::Mutex::new(Vec::new())),
gate_refresh_release: Arc::new(Notify::new()),
result,
large_result,
dictionary_result,
@@ -279,6 +284,7 @@ impl FlightService for TestSqlService {
) -> std::result::Result<Response<<Self as FlightService>::DoGetStream>, Status> {
self.do_get_count.fetch_add(1, Ordering::SeqCst);
let ticket = request.get_ref().ticket.as_ref();
let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"";
let empty = ticket == b"SELECT empty";
let slow = ticket == b"SELECT slow get";
let large = ticket == b"SELECT large message";
@@ -290,7 +296,11 @@ impl FlightService for TestSqlService {
self.result.clone()
};
let schema = result.schema();
let gate_refresh_release = self.gate_refresh_release.clone();
let input = futures::stream::once(async move {
if gate_refresh {
gate_refresh_release.notified().await;
}
if slow {
tokio::time::sleep(Duration::from_millis(250)).await;
}
@@ -446,6 +456,147 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() {
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn remote_refresh_twice_keeps_latest_selector() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let incoming = futures::stream::try_unfold(listener, |listener| async {
let (socket, _) = listener.accept().await?;
Ok::<_, std::io::Error>(Some((socket, listener)))
});
let service = TestSqlService::default();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(FlightServiceServer::new(service))
.serve_with_incoming_shutdown(incoming, async {
let _ = shutdown_rx.await;
}),
);
let sql_client = SqlClient::new(
"analytics".into(),
None,
"test-key".into(),
None,
Some(format!("grpc://{address}")),
ClientConfig::default(),
);
let selectors = Arc::new(std::sync::Mutex::new(Vec::new()));
let captured = selectors.clone();
let client = client_with_handler(move |request| {
let body = request
.body()
.and_then(reqwest::Body::as_bytes)
.and_then(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).ok())
.unwrap();
captured.lock().unwrap().push(body["version"].clone());
let text = match request.url().path() {
"/v1/table/ns$view/describe/" => r#"{"version": 42, "schema": {"fields": []}}"#,
"/v1/table/ns$view/count_rows/" => "2",
other => panic!("unexpected REST request: {other}"),
};
http::Response::builder()
.status(200)
.body(text.to_string())
.unwrap()
});
let table = RemoteTable::new_with_sql_client(
client,
"view".into(),
vec!["ns".into()],
"ns$view".into(),
Default::default(),
Some(sql_client),
);
let first = table
.refresh_materialized_view_async(None)
.await
.unwrap()
.wait()
.await
.unwrap();
let second = table
.refresh_materialized_view_async(None)
.await
.unwrap()
.wait()
.await
.unwrap();
assert_eq!(first.version, 42);
assert_eq!(second.version, 42);
assert!(
selectors
.lock()
.unwrap()
.iter()
.all(serde_json::Value::is_null)
);
shutdown_tx.send(()).unwrap();
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn cancelling_remote_refresh_cancels_sql_query() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let incoming = futures::stream::try_unfold(listener, |listener| async {
let (socket, _) = listener.accept().await?;
Ok::<_, std::io::Error>(Some((socket, listener)))
});
let service = TestSqlService::default();
let cancel_count = service.cancel_count.clone();
let do_get_count = service.do_get_count.clone();
let gate_refresh_release = service.gate_refresh_release.clone();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(FlightServiceServer::new(service))
.serve_with_incoming_shutdown(incoming, async {
let _ = shutdown_rx.await;
}),
);
let sql_client = SqlClient::new(
"analytics".into(),
None,
"test-key".into(),
None,
Some(format!("grpc://{address}")),
ClientConfig::default(),
);
let client = client_with_handler(|request| -> http::Response<String> {
panic!(
"cancelled refresh must not read result metadata: {}",
request.url().path()
);
});
let table = RemoteTable::new_with_sql_client(
client,
"cancel_view".into(),
vec![],
"cancel_view".into(),
Default::default(),
Some(sql_client),
);
let job = table.refresh_materialized_view_async(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(10), async {
while do_get_count.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
job.cancel().await.unwrap();
assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. })));
assert_eq!(cancel_count.load(Ordering::SeqCst), 1);
gate_refresh_release.notify_one();
shutdown_tx.send(()).unwrap();
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn submits_polls_fetches_cancels_and_reuses_client() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+239 -178
View File
@@ -8,6 +8,7 @@ use self::insert::{RemoteWriteExec, WriteOp};
use super::client::RequestResultExt;
use super::client::{HttpSend, RestfulLanceDbClient, Sender};
use super::db::ServerVersion;
use super::sql::SqlClient;
use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE, extract_job_id};
use crate::blob::BlobFile;
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
@@ -18,7 +19,7 @@ use crate::index::scalar::FtsQuery;
use crate::index::waiter::wait_for_index;
use crate::job::Job;
use crate::materialized_view::{
MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, ViewProjection,
MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult,
};
use crate::query::wal_fusion::PkFusionMemory; // WAL-PK-FUSION: delete.
use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest};
@@ -264,7 +265,6 @@ struct FreshnessJob<S: HttpSend> {
enum TrackedJobResult {
None,
RefreshColumn,
MaterializedView,
}
#[async_trait]
@@ -305,11 +305,6 @@ impl<S: HttpSend> crate::job::JobHandle for FreshnessJob<S> {
})
})
}),
TrackedJobResult::MaterializedView => result.value().and_then(|value| {
serde_json::from_value::<RefreshMaterializedViewResult>(value.clone())
.ok()
.map(|result| result.version)
}),
}
.filter(|version| *version != 0);
if let Some(version) = result_version {
@@ -347,6 +342,10 @@ fn compute_min_timestamp(
}
}
fn quote_sql_identifier(identifier: &str) -> String {
format!("\"{}\"", identifier.replace('"', "\"\""))
}
fn freshness_headers_snapshot(
freshness: &Mutex<FreshnessState>,
interval: Option<Duration>,
@@ -496,6 +495,7 @@ pub struct RemoteTable<S: HttpSend = Sender> {
namespace: Vec<String>,
identifier: String,
server_version: ServerVersion,
sql_client: Option<SqlClient>,
version: Arc<RwLock<Option<u64>>>,
location: RwLock<Option<String>>,
@@ -654,12 +654,13 @@ impl<S: HttpSend> RemoteTable<S> {
Ok(job_id)
}
pub fn new(
pub(super) fn new_with_sql_client(
client: RestfulLanceDbClient<S>,
name: String,
namespace: Vec<String>,
identifier: String,
server_version: ServerVersion,
sql_client: Option<SqlClient>,
) -> Self {
Self {
client,
@@ -667,6 +668,7 @@ impl<S: HttpSend> RemoteTable<S> {
namespace,
identifier,
server_version,
sql_client,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -701,6 +703,7 @@ impl<S: HttpSend> RemoteTable<S> {
namespace: self.namespace.clone(),
identifier: self.identifier.clone(),
server_version: self.server_version.clone(),
sql_client: self.sql_client.clone(),
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1177,11 +1180,14 @@ impl<S: HttpSend> RemoteTable<S> {
.values()
.iter()
.map(|v| {
serde_json::Value::Number(
serde_json::Number::from_f64(*v as f64).unwrap(),
)
serde_json::Number::from_f64(*v as f64)
.map(serde_json::Value::Number)
.ok_or_else(|| Error::InvalidInput {
message: "query vector must contain only finite values"
.into(),
})
})
.collect(),
.collect::<Result<Vec<_>>>()?,
))
}
_ => Err(Error::InvalidInput {
@@ -1619,6 +1625,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: version.map(ServerVersion).unwrap_or_default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1644,6 +1651,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: ServerVersion::default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1678,6 +1686,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: version.map(ServerVersion).unwrap_or_default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -2106,29 +2115,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
&self.identifier
}
async fn materialized_view_info(&self) -> Result<MaterializedViewInfo> {
#[derive(Deserialize)]
struct Projection {
output_column: String,
expression: String,
}
#[derive(Deserialize)]
struct DescribeMaterializedViewResponse {
source_table: String,
#[serde(default)]
source_namespace: Vec<String>,
#[serde(default)]
projections: Vec<Projection>,
#[serde(default)]
filter: Option<String>,
#[serde(default)]
limit: Option<u64>,
/// The defining query, once the server describes a view by it;
/// takes precedence over the structured fields.
#[serde(default)]
query: Option<String>,
#[serde(default)]
incarnation: Option<String>,
query: String,
}
let request = self.client.post(&format!(
@@ -2139,79 +2128,64 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
let response = self.check_table_response(&request_id, response).await?;
let response: DescribeMaterializedViewResponse =
response.json().await.err_to_http(request_id)?;
let definition = match response.query {
Some(query) => MaterializedViewDefinition::from_sql(&query)?,
None => MaterializedViewDefinition {
vector_source: None,
source_table: response.source_table,
source_namespace: response.source_namespace,
lateral: None,
projections: response
.projections
.into_iter()
.map(|projection| ViewProjection {
output: projection.output_column,
expression: projection.expression,
})
.collect(),
filter: response.filter,
group_by: Vec::new(),
limit: response.limit,
},
};
let parsed_definition = MaterializedViewDefinition::from_sql(&response.query).ok();
Ok(MaterializedViewInfo {
definition,
incarnation: response.incarnation,
definition_sql: response.query,
parsed_definition,
})
}
async fn refresh_materialized_view_async(
&self,
full: bool,
source_version: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<Job<RefreshMaterializedViewResult>> {
self.check_mutable().await?;
let mut body = serde_json::json!({ "full": full });
if let Some(source_version) = source_version {
body["source_version"] = source_version.into();
}
if let Some(expected_incarnation) = expected_incarnation {
body["expected_incarnation"] = expected_incarnation.into();
}
let request = self
.client
.post(&format!(
"/v1/materialized_view/{}/refresh",
self.identifier
))
.json(&body);
let freshness_request = self.snapshot_freshness_headers();
let (request_id, response) = self
.send_with_freshness(request, true, freshness_request)
.await?;
let response = self.check_table_response(&request_id, response).await?;
let status = response.status();
let body = response.text().await.err_to_http(request_id.clone())?;
if status != StatusCode::ACCEPTED {
return Err(Error::Http {
source: "materialized-view refresh must return 202 Accepted".into(),
request_id,
status_code: Some(status),
if source_version.is_some() {
return Err(Error::NotSupported {
message: "source-version pinning is available only for local materialized views"
.to_string(),
});
}
let job_id = extract_job_id(&body).ok_or_else(|| Error::Http {
source: "materialized-view refresh response did not contain a valid job_id".into(),
request_id,
status_code: Some(status),
let sql_client = self.sql_client.clone().ok_or_else(|| Error::NotSupported {
message: "SQL is unavailable for this remote table client".to_string(),
})?;
Ok(Job::new_typed(Box::new(FreshnessJob {
inner: RemoteJob::new(self.client.clone(), job_id),
freshness: self.freshness.clone(),
version: self.version.clone(),
tracked_result: TrackedJobResult::MaterializedView,
freshness_request,
})))
let statement = format!(
"REFRESH MATERIALIZED VIEW {}",
quote_sql_identifier(&self.name)
);
let namespace = self.namespace.clone();
let client = self.client.clone();
let name = self.name.clone();
let identifier = self.identifier.clone();
let server_version = self.server_version.clone();
let freshness = self.freshness.clone();
let freshness_request = self.snapshot_freshness_headers();
let job_sql_client = sql_client.clone();
Ok(
sql_client.submit_as_job(statement, namespace.clone(), move || async move {
let table = Self::new_with_sql_client(
client,
name,
namespace,
identifier,
server_version,
Some(job_sql_client),
);
table.checkout_latest().await?;
let version = table.version().await?;
let rows_written =
u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime {
message: "materialized-view row count exceeds u64".to_string(),
})?;
freshness_request.observe_version(&freshness, version);
Ok(RefreshMaterializedViewResult {
mode: crate::materialized_view::RefreshMode::Rebuild,
rows_written,
source_version: 0,
version,
})
}),
)
}
async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>> {
let description = self.describe().await?;
@@ -2301,11 +2275,14 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
Ok(Some(Arc::new(snapshot)))
}
async fn restore(&self) -> Result<()> {
let read_snapshot = self.snapshot_read_state().await;
let version = read_snapshot.version.ok_or_else(|| Error::InvalidInput {
message: "you must run checkout before running restore".to_string(),
})?;
let mut request = self
.client
.post(&format!("/v1/table/{}/restore/", self.identifier));
let read_snapshot = self.snapshot_read_state().await;
let mut body = serde_json::json!({ "version": read_snapshot.version });
let mut body = serde_json::json!({ "version": version });
self.apply_branch_body(&mut body);
request = request.json(&body);
@@ -3281,8 +3258,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
#[derive(Deserialize)]
struct LsmWriteSpecBody {
sharding: Sharding,
/// `null` selects every index the table has; `[]` selects none.
#[serde(default)]
maintained_indexes: Vec<String>,
maintained_indexes: Option<Vec<String>>,
#[serde(default)]
writer_config_defaults: std::collections::HashMap<String, String>,
}
@@ -4648,10 +4626,18 @@ mod tests {
}
#[rstest]
#[case(true)]
#[case(false)]
#[case::old_server("", 0, 0)]
#[case::rows_updated(r#"{"rows_updated": 5, "version": 43}"#, 5, 43)]
#[case::updated_rows(r#"{"updated_rows": 5, "version": 43}"#, 5, 43)]
#[case::zero_updated_rows(r#"{"updated_rows": 0, "version": 43}"#, 0, 43)]
#[case::missing_row_count(r#"{"version": 43}"#, 0, 43)]
#[tokio::test]
async fn test_update(#[case] old_server: bool) {
async fn test_update(
#[case] response_body: &'static str,
#[case] expected_rows_updated: u64,
#[case] expected_version: u64,
#[values(true, false)] filtered: bool,
) {
let table = Table::new_with_handler("my_table", move |request| {
if request.url().path() == "/v1/table/my_table/update/" {
assert_eq!(request.method(), "POST");
@@ -4676,32 +4662,29 @@ mod tests {
assert_eq!(col_name, "b");
assert_eq!(expression, "b - 1");
let only_if = value.get("predicate").unwrap().as_str().unwrap();
assert_eq!(only_if, "`B` > 10");
assert_eq!(
value.get("predicate").unwrap(),
&serde_json::json!(if filtered { Some("`B` > 10") } else { None })
);
}
if old_server {
http::Response::builder().status(200).body("").unwrap()
} else {
http::Response::builder()
.status(200)
.body(r#"{"rows_updated": 5, "version": 43}"#)
.unwrap()
}
http::Response::builder()
.status(200)
.body(response_body)
.unwrap()
} else {
panic!("Unexpected request path: {}", request.url().path());
}
});
let update = table
.update()
.column("a", "a + 1")
.column("b", "b - 1")
.only_if(r#""B" > 10"#);
let mut update = table.update().column("a", "a + 1").column("b", "b - 1");
if filtered {
update = update.only_if(r#""B" > 10"#);
}
let result = table.base_table().update(update).await.unwrap();
assert_eq!(result.version, if old_server { 0 } else { 43 });
assert_eq!(result.rows_updated, if old_server { 0 } else { 5 });
assert_eq!(result.version, expected_version);
assert_eq!(result.rows_updated, expected_rows_updated);
}
#[tokio::test]
@@ -5912,6 +5895,47 @@ mod tests {
assert_eq!(blobs.value(0), b"alpha");
}
#[rstest]
#[case(DEFAULT_SERVER_VERSION.clone())]
#[case(semver::Version::new(0, 2, 0))]
#[tokio::test]
async fn test_query_vector_non_finite(#[case] version: semver::Version) {
let table =
Table::new_with_handler_version("my_table", version, |_| -> http::Response<String> {
panic!("non-finite vectors must be rejected before sending a request")
});
for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
// Requests built without the query builder must also return an
// error, including a non-finite vector later in a batch.
for batched in [false, true] {
let mut request = table
.query()
.nearest_to(&[0.1, 0.2])
.unwrap()
.into_request();
if !batched {
request.query_vector.clear();
}
request
.query_vector
.push(Arc::new(arrow_array::Float32Array::from(vec![0.1, value])));
let result = table
.base_table()
.query(
&AnyQuery::VectorQuery(request),
QueryExecutionOptions::default(),
)
.await;
let Err(err) = result else {
panic!("non-finite query vector unexpectedly succeeded")
};
assert!(matches!(err, Error::InvalidInput { .. }));
assert!(err.to_string().contains("only finite values"));
}
}
}
#[tokio::test]
async fn test_query_vector_default_values() {
let expected_data = RecordBatch::try_new(
@@ -8186,6 +8210,43 @@ mod tests {
.unwrap();
}
#[tokio::test]
async fn test_restore_requires_checkout() {
let request_count = Arc::new(AtomicUsize::new(0));
let request_count_clone = request_count.clone();
let table = Table::new_with_handler("my_table", move |request| {
request_count_clone.fetch_add(1, Ordering::SeqCst);
assert_eq!(request_body_json(&request)["version"], 42);
let body = match request.url().path() {
"/v1/table/my_table/describe/" => r#"{"version":42,"schema":{"fields":[]}}"#,
"/v1/table/my_table/restore/" => r#"{"version":43}"#,
path => panic!("unexpected request path: {path}"),
};
http::Response::builder().status(200).body(body).unwrap()
});
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 0);
table.checkout(42).await.unwrap();
table.checkout_latest().await.unwrap();
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 1);
table.checkout(42).await.unwrap();
table.restore().await.unwrap();
assert_eq!(request_count.load(Ordering::SeqCst), 3);
let err = table.restore().await.unwrap_err();
assert!(matches!(err, Error::InvalidInput { message }
if message == "you must run checkout before running restore"));
assert_eq!(request_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn test_fails_if_checkout_version_doesnt_exist() {
let table = Table::new_with_handler("my_table", |request| {
@@ -9613,6 +9674,43 @@ mod tests {
}
}
/// Every selection reads back as the server reported it: every index
/// (`null`), none (`[]`), and a named list.
#[rstest::rstest]
#[case::every_index(serde_json::Value::Null, None)]
#[case::no_index(serde_json::json!([]), Some(vec![]))]
#[case::named(serde_json::json!(["id_idx"]), Some(vec!["id_idx".to_string()]))]
#[tokio::test]
async fn test_get_lsm_write_spec_round_trips_the_selection(
#[case] reported: serde_json::Value,
#[case] expected: Option<Vec<String>>,
) {
let table = Table::new_with_handler("my_table", move |_| {
let response = serde_json::json!({
"lsm_write_spec": {
"sharding": { "mode": "unsharded" },
"maintained_indexes": reported,
"writer_config_defaults": {},
}
});
http::Response::builder()
.status(200)
.body(response.to_string())
.unwrap()
});
let spec = table
.get_lsm_write_spec()
.await
.unwrap()
.expect("a spec should be reported");
assert_eq!(
spec.maintained_indexes().map(<[String]>::to_vec),
expected,
"the selection the server reported must survive the read"
);
}
#[tokio::test]
async fn test_get_lsm_write_spec_absent() {
let table = Table::new_with_handler("my_table", |request| {
@@ -10978,22 +11076,24 @@ mod tests {
}
});
table.checkout(1).await.unwrap();
// First schema call
let schema1 = table.schema().await.unwrap();
assert_eq!(call_count.load(Ordering::SeqCst), 1);
assert_eq!(call_count.load(Ordering::SeqCst), 2);
// Second schema call uses cache
let schema2 = table.schema().await.unwrap();
assert_eq!(Arc::as_ptr(&schema2), Arc::as_ptr(&schema1));
assert_eq!(call_count.load(Ordering::SeqCst), 1);
assert_eq!(call_count.load(Ordering::SeqCst), 2);
// Restore operation
let _ = table.restore().await;
table.restore().await.unwrap();
// Schema call after restore should re-fetch (cache invalidated)
let schema3 = table.schema().await.unwrap();
assert_ne!(Arc::as_ptr(&schema3), Arc::as_ptr(&schema1));
assert_eq!(call_count.load(Ordering::SeqCst), 2);
assert_eq!(call_count.load(Ordering::SeqCst), 3);
}
/// Test that centralized error handling invalidates cache on query errors
@@ -12976,54 +13076,16 @@ mod tests {
}
#[tokio::test]
async fn test_materialized_view_describe_and_refresh() {
async fn test_materialized_view_describe_and_refresh_requires_sql() {
const QUERY: &str =
"SELECT s.x, d.label FROM analytics.source s JOIN analytics.dim d ON s.id = d.id";
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
"/v1/materialized_view/my_table/describe" => http::Response::builder()
.status(200)
.body(
json!({
"name": "my_table",
"source_table": "source",
"source_namespace": ["analytics"],
"projections": [{
"output_column": "double_x",
"expression": "x * 2"
}],
"filter": "x > 0",
"limit": 10,
"inputs": ["x"],
"incarnation": "inc-1"
})
.to_string(),
)
.unwrap(),
"/v1/materialized_view/my_table/refresh" => {
assert_eq!(request.method(), "POST");
assert_eq!(
request_body_json(&request),
json!({
"full": true,
"source_version": 7,
"expected_incarnation": "inc-1"
})
);
http::Response::builder()
.status(202)
.body(json!({"job_id": "j1-mv-refresh"}).to_string())
.unwrap()
}
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(
json!({
"job_id": "j1-mv-refresh",
"job_state": "DONE",
"result": {
"mode": "rebuild",
"rows_written": 2,
"source_version": 7,
"version": 9
}
"query": QUERY
})
.to_string(),
)
@@ -13031,21 +13093,10 @@ mod tests {
path => panic!("unexpected request: {path}"),
});
let view = crate::MaterializedView::from_table(table).await.unwrap();
assert_eq!(view.definition().source_table, "source");
assert_eq!(view.definition().source_namespace, ["analytics"]);
assert_eq!(view.incarnation(), Some("inc-1"));
let result = view
.refresh()
.full(true)
.source_version(7)
.expect_incarnation("inc-1")
.execute()
.await
.unwrap();
assert_eq!(result.mode, crate::RefreshMode::Rebuild);
assert_eq!(result.rows_written, 2);
assert_eq!(result.version, 9);
assert_eq!(view.definition_sql(), QUERY);
assert!(view.definition().is_err());
let error = view.refresh().execute().await.unwrap_err();
assert!(error.to_string().contains("SQL is unavailable"));
}
#[tokio::test]
@@ -13863,8 +13914,17 @@ mod tests {
.status(200)
.body("{}".to_string())
.unwrap(),
"/v1/table/my_table/describe/" => {
assert_eq!(request_body_json(&request)["branch"], "exp");
assert_eq!(request_body_json(&request)["version"], 1);
http::Response::builder()
.status(200)
.body(r#"{"version":1,"schema":{"fields":[]}}"#.to_string())
.unwrap()
}
"/v1/table/my_table/restore/" => {
assert_eq!(request_body_json(&request)["branch"], "exp");
assert_eq!(request_body_json(&request)["version"], 1);
http::Response::builder()
.status(200)
.body(r#"{"version":1}"#.to_string())
@@ -13876,6 +13936,7 @@ mod tests {
.create_branch("exp", Ref::Version(None, None))
.await
.unwrap();
branch.checkout(1).await.unwrap();
branch.restore().await.unwrap();
}
+11 -9
View File
@@ -563,7 +563,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
fn id(&self) -> &str;
/// Get the arrow [Schema] of the table.
async fn schema(&self) -> Result<SchemaRef>;
/// Read this table's materialized-view definition and incarnation.
/// Read this table's materialized-view definition.
#[doc(hidden)]
async fn materialized_view_info(
&self,
@@ -578,9 +578,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
#[doc(hidden)]
async fn refresh_materialized_view_async(
&self,
_full: bool,
_source_version: Option<u64>,
_expected_incarnation: Option<&str>,
) -> Result<Job<crate::materialized_view::RefreshMaterializedViewResult>> {
Err(Error::NotSupported {
message: "remote materialized-view refresh is not supported on this table type".into(),
@@ -2036,12 +2034,16 @@ impl Table {
/// Read the [`LsmWriteSpec`] currently installed on this table.
///
/// Returns `Ok(None)` when the MemWAL LSM write path is not enabled (no
/// spec has been set, or it was removed with [`Table::unset_lsm_write_spec`]).
/// The returned spec mirrors what was passed to
/// [`Table::set_lsm_write_spec`], except that
/// [`LsmWriteSpec::maintained_indexes`] always reports the concrete list
/// resolved when the spec was set — `None` never round-trips.
/// `Ok(None)` means the LSM write path is not enabled at all — no spec has
/// been set, or one was removed with [`Table::unset_lsm_write_spec`]. That
/// is a different answer from a spec whose
/// [`LsmWriteSpec::maintained_indexes`] is `None`, which is an installed
/// spec selecting indexes automatically.
///
/// The spec read back is the one that was installed, selection included:
/// `None` maintains every supported index the table has now or gains
/// later, `[]` maintains none, and a non-empty list maintains exactly
/// those. All three round-trip.
///
/// # Example
///
+62 -7
View File
@@ -20,7 +20,8 @@ use datafusion_common::{DataFusionError, Result as DataFusionResult, Statistics}
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType, dml::InsertOp};
use datafusion_physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, stream::RecordBatchStreamAdapter,
DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, execution_plan::CardinalityEffect,
stream::RecordBatchStreamAdapter,
};
use futures::{TryFutureExt, TryStreamExt};
use lance::dataset::{WriteMode, WriteParams};
@@ -123,12 +124,11 @@ impl ExecutionPlan for MetadataEraserExec {
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
assert_eq!(children.len(), 1);
let new_properties = Self::compute_properties_from_input(&children[0], &self.schema);
Ok(Arc::new(Self {
input: children[0].clone(),
schema: self.schema.clone(),
properties: new_properties,
}))
Ok(Arc::new(Self::new(children[0].clone())))
}
fn cardinality_effect(&self) -> CardinalityEffect {
CardinalityEffect::Equal
}
fn execute(
@@ -668,4 +668,59 @@ pub mod tests {
assert!(matches!(partition_stats.num_rows, Precision::Exact(10)));
}
#[tokio::test]
async fn test_count_pushdown_through_metadata_eraser() {
let fixture = TestFixture::new().await;
let state = datafusion::execution::SessionStateBuilder::new()
.with_default_features()
.with_physical_optimizer_rule(Arc::new(lance::io::exec::count_pushdown::CountPushdown))
.build();
let ctx = SessionContext::new_with_state(state);
ctx.register_table("foo", fixture.adapter.clone()).unwrap();
let df = ctx
.sql("SELECT COUNT(*) FROM foo WHERE indexed > 2")
.await
.unwrap();
let plan = df.clone().create_physical_plan().await.unwrap();
let plan = datafusion_physical_plan::displayable(plan.as_ref())
.indent(true)
.to_string();
assert!(plan.contains("CountFromMask"), "{plan}");
let batches = df.collect().await.unwrap();
assert_eq!(
batches[0]
.column(0)
.as_primitive::<arrow_array::types::Int64Type>()
.value(0),
7
);
}
#[tokio::test]
async fn test_metadata_eraser_with_new_children_takes_child_schema() {
let fixture = TestFixture::new().await;
let ctx = SessionContext::new();
let scan = |adapter: Arc<BaseTableAdapter>, name: &str| {
LogicalPlanBuilder::scan(name, provider_as_source(adapter), None)
.unwrap()
.build()
.unwrap()
};
let foo = ctx
.state()
.create_physical_plan(&scan(fixture.adapter.clone(), "foo"))
.await
.unwrap();
let tbl2 = ctx
.state()
.create_physical_plan(&scan(fixture.adapter2.clone(), "tbl2"))
.await
.unwrap();
let rewired = foo.with_new_children(vec![tbl2.clone()]).unwrap();
assert_eq!(rewired.schema(), tbl2.schema());
}
}
@@ -6,11 +6,15 @@
use std::sync::Arc;
use arrow_schema::SchemaRef;
use arrow_array::{RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::prelude::{SessionConfig, SessionContext, col};
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::{Result, plan_err};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_common::{DataFusionError, Result, plan_err};
use datafusion_execution::{
SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder,
};
use datafusion_expr::{Expr, TableType};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_plan::{
@@ -25,7 +29,9 @@ use futures::{TryStreamExt, stream};
use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap};
use lance::{Dataset, index::DatasetIndexInternalExt};
use lance_core::datatypes::BlobHandling;
use lance_datafusion::exec::SessionContextExt;
use lance_index::metrics::NoOpMetricsCollector;
use roaring::RoaringTreemap;
use super::duplicate_pairs::{
DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput,
@@ -36,6 +42,84 @@ use crate::table::{NativeTableExt, datafusion::BaseTableAdapter};
const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MATERIALIZE_BATCH_ROWS: usize = 1024;
fn edge_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt64, false),
Field::new("b", DataType::UInt64, false),
]))
}
/// Sorting canonical edges makes the first unremoved endpoint the smallest
/// unassigned representative. The sort spills; only vertex masks stay in RAM.
async fn select_representatives(
stream: SendableRecordBatchStream,
memory: usize,
) -> crate::Result<(RoaringTreemap, RoaringTreemap)> {
let schema = edge_schema();
let batch_schema = schema.clone();
let canonical = stream
.map_ok(move |batch| {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
RecordBatch::try_new(
batch_schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)),
)),
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)),
)),
],
)
.map_err(DataFusionError::from)
})
.and_then(futures::future::ready);
let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical));
let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS);
config.options_mut().execution.sort_spill_reservation_bytes = memory / 2;
let ctx = SessionContext::new_with_config_rt(
config,
RuntimeEnvBuilder::new()
.with_memory_limit(memory, 1.0)
.build_arc()?,
);
let mut sorted = ctx
.read_one_shot(canonical)?
.sort_by(vec![col("a"), col("b")])?
.execute_stream()
.await?;
let mut removed = RoaringTreemap::new();
let mut representatives = RoaringTreemap::new();
while let Some(batch) = sorted.try_next().await? {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
for (&a, &b) in a.values().iter().zip(b.values()) {
if a != b && !removed.contains(a) && !removed.contains(b) {
representatives.insert(a);
removed.insert(b);
}
}
}
Ok((removed, representatives))
}
/// Select nonempty physical IVF scopes before dispatch, never by filtering the
/// resulting pairs. A count spans all segments; partition IDs are segment-local.
#[derive(Debug, Clone, Default)]
@@ -322,12 +406,9 @@ impl ExecutionPlan for VectorDedupExec {
let dataset = self.dataset.clone();
let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?;
let stream = stream::once(async move {
let (removed, _) = crate::materialized_view::vector_dedup::select_representatives(
pairs,
REPRESENTATIVE_SORT_MEMORY_BYTES,
)
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
let (removed, _) = select_representatives(pairs, REPRESENTATIVE_SORT_MEMORY_BYTES)
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
let mut scanner = dataset.scan();
scanner
.with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter(
+1 -1
View File
@@ -13,7 +13,7 @@ use crate::Result;
/// The result of an update operation
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct UpdateResult {
#[serde(default)]
#[serde(default, alias = "updated_rows")]
pub rows_updated: u64,
/// The commit version associated with the operation.
#[serde(default)]
+3 -205
View File
@@ -289,127 +289,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
);
let retry = ctx.sql(&sql).await?.collect().await?;
assert_eq!(pairs(&retry), pairs(&actual));
// The same table function can be maintained through the MV unit
// lifecycle. Staging attempts do not make any rows visible.
use lancedb::materialized_view::{
MaterializedViewDefinition, commit_partitioned_refresh, plan_partitioned_refresh,
prepare_definition, write_refresh_partition,
};
let view_name = format!("pairs_{}", threshold.to_bits());
let view = prepare_definition(&table, MaterializedViewDefinition::from_sql(&sql)?)
.await?
.create(&view_name)
.await?;
assert_eq!(view.table().schema().await?.fields().len(), 3);
let plan = plan_partitioned_refresh(view.table(), Some(version))
.await?
.unwrap();
assert_eq!(plan.units() as usize, tasks.len());
let plan = serde_json::from_slice(&serde_json::to_vec(&plan)?)?;
let first = write_refresh_partition(view.table(), 0, &plan).await?;
assert_eq!(view.table().count_rows(None).await?, 0);
assert!(
commit_partitioned_refresh(
view.table(),
&plan,
vec![first.clone(), first.clone()],
view.incarnation()
)
.await
.is_err()
);
// Lost receipt: its files remain unreferenced; the retry writes
// another attempt, of which only one receipt is ever committed.
drop(write_refresh_partition(view.table(), 0, &plan).await?);
let mut units = vec![first];
for unit in 1..plan.units() {
units.push(write_refresh_partition(view.table(), unit, &plan).await?);
}
assert_eq!(view.table().count_rows(None).await?, 0);
let outcome =
commit_partitioned_refresh(view.table(), &plan, units.clone(), view.incarnation())
.await?;
assert_eq!(outcome.rows_written as usize, pairs(&expected).len());
assert!(
commit_partitioned_refresh(view.table(), &plan, units, view.incarnation())
.await
.is_err()
);
use lancedb::query::ExecutableQuery;
let written = view
.table()
.query()
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
assert_eq!(pairs(&written), pairs(&expected));
// Repeating a refresh replaces the view rather than appending.
view.refresh().execute().await?;
assert_eq!(
view.table().count_rows(None).await? as usize,
pairs(&expected).len()
);
// The complete declarative pipeline keeps direct representatives
// and materializes original rows, including isolated source rows.
use lancedb::materialized_view::write_refresh_partition_with_inputs;
use lancedb::query::QueryBase;
let definition = MaterializedViewDefinition::from_sql(&format!(
"SELECT * FROM vector_dedup('source', {version}, 'vector', {threshold})"
))?;
let clean = prepare_definition(&table, definition)
.await?
.create(&format!("clean_{}", threshold.to_bits()))
.await?;
let clean_plan = plan_partitioned_refresh(clean.table(), Some(version))
.await?
.unwrap();
assert_eq!(clean_plan.dependency_units() as usize, tasks.len());
assert_eq!(clean_plan.units() as usize, tasks.len() + 2);
let clean_plan: lancedb::materialized_view::PartitionedRefreshPlan =
serde_json::from_slice(&serde_json::to_vec(&clean_plan)?)?;
assert!(
write_refresh_partition(clean.table(), tasks.len() as u32, &clean_plan)
.await
.is_err()
);
let mut receipts = Vec::new();
for unit in 0..clean_plan.units() {
let dependencies = &receipts[..receipts.len().min(tasks.len())];
if unit == 0 || unit == tasks.len() as u32 {
// Both selection and materialization can lose a receipt.
drop(
write_refresh_partition_with_inputs(
clean.table(),
unit,
&clean_plan,
dependencies,
)
.await?,
);
}
let receipt = write_refresh_partition_with_inputs(
clean.table(),
unit,
&clean_plan,
dependencies,
)
.await?;
receipts.push(serde_json::from_slice(&serde_json::to_vec(&receipt)?)?);
}
assert_eq!(clean.table().count_rows(None).await?, 0);
assert!(
commit_partitioned_refresh(
clean.table(),
&clean_plan,
receipts[..receipts.len() - 1].to_vec(),
clean.incarnation()
)
.await
.is_err()
);
commit_partitioned_refresh(clean.table(), &clean_plan, receipts, clean.incarnation())
.await?;
let mut edges = pairs(&expected)
.into_iter()
.map(|(a, b, _)| (a.min(b), a.max(b)))
@@ -449,27 +328,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
}
}
}
let kept = clean
.table()
.query()
.select(lancedb::query::Select::columns(&["id"]))
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut actual_ids = kept
.iter()
.flat_map(|b| {
b.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.iter()
.copied()
})
.collect::<Vec<_>>();
actual_ids.sort_unstable();
expected_ids.sort_unstable();
let mut query_ids = clean_query
.iter()
@@ -486,14 +344,10 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
query_ids.sort_unstable();
assert_eq!(query_ids, expected_ids);
assert_eq!(actual_ids, expected_ids);
if threshold == 1.0 {
assert!(actual_ids.contains(&0) && actual_ids.contains(&2));
assert!(!actual_ids.contains(&1), "A-B-C retains A and C");
assert!(query_ids.contains(&0) && query_ids.contains(&2));
assert!(!query_ids.contains(&1), "A-B-C retains A and C");
}
clean.refresh().execute().await?;
assert_eq!(clean.table().count_rows(None).await?, expected_ids.len());
assert_eq!(table.count_rows(None).await?, n * 2);
assert!(
DuplicatePairsExec::try_new(
ds.clone(),
@@ -527,10 +381,8 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
}
#[tokio::test]
async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() -> anyhow::Result<()>
{
async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Result<()> {
use arrow_array::{LargeBinaryArray, StringArray, StructArray, types::Float32Type};
use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition};
for stable in [false, true] {
let dir = tempfile::tempdir()?;
let conn = lancedb::connect(dir.path().to_str().unwrap())
@@ -607,7 +459,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
true,
)
.await?;
let version = ds.version().version;
use lancedb::table::datafusion::udtf::{
duplicate_pairs::DuplicatePairsOutput,
vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable},
@@ -692,59 +543,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]);
}
}
let definition = MaterializedViewDefinition::from_sql(&format!(
"SELECT * FROM vector_dedup('source', {version}, 'vector', 1)"
))?;
let clean = prepare_definition(&table, definition)
.await?
.create("clean")
.await?;
// A later source delete does not change the declared result snapshot.
table.delete("id = 0").await?;
let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?;
let explicit = adapter.vector_query_snapshot(Some(version)).await?;
let latest = adapter.vector_query_snapshot(None).await?;
assert_eq!(explicit.version().version, version);
assert!(latest.version().version > version);
assert_eq!(explicit.count_rows(None).await?, 5);
assert_eq!(latest.count_rows(None).await?, 4);
clean.refresh().execute().await?;
assert_eq!(clean.table().count_rows(None).await?, 4);
let output = Dataset::open(&clean.table().uri().await?).await?;
let batches = output
.scan()
.blob_handling(lance_core::datatypes::BlobHandling::AllBinary)
.try_into_stream()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut ids = Vec::new();
for batch in batches {
let row_ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let images = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
for i in 0..batch.num_rows() {
let id = row_ids.value(i);
assert_eq!(images.value(i), vec![id as u8; 4096]);
ids.push(id);
}
}
ids.sort_unstable();
assert_eq!(ids, [0, 2, 3, 4]);
assert_eq!(
table.count_rows(None).await?,
4,
"refresh must not mutate the source"
);
}
Ok(())
}