mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-14 07:42:28 +00:00
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:
+1
-1
@@ -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*)\\.
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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">
|
||||
|
||||
[](https://www.reveriesummit.com/?utm_medium=social&utm_source=github&utm_campaign=reverie-2026)
|
||||
[](https://lancedb.com)
|
||||
[](https://lancedb.com/)
|
||||
[](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 |
@@ -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>
|
||||
```
|
||||
|
||||
|
||||
@@ -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`<[`ListDatabasesResponse`](../interfaces/ListDatabasesResponse.md)>
|
||||
[`DatabaseNames`](DatabaseNames.md)
|
||||
|
||||
#### Example
|
||||
|
||||
```ts
|
||||
for await (const name of catalog.listDatabases({ pageLimit: 20 })) {
|
||||
console.log(name);
|
||||
}
|
||||
```
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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`<`string`>
|
||||
|
||||
## Methods
|
||||
|
||||
### \[asyncIterator\]()
|
||||
|
||||
```ts
|
||||
asyncIterator: AsyncIterableIterator<string>
|
||||
```
|
||||
|
||||
#### Returns
|
||||
|
||||
`AsyncIterableIterator`<`string`>
|
||||
|
||||
#### Implementation of
|
||||
|
||||
`AsyncIterableIterator.[asyncIterator]`
|
||||
|
||||
***
|
||||
|
||||
### next()
|
||||
|
||||
```ts
|
||||
next(): Promise<IteratorResult<string, any>>
|
||||
```
|
||||
|
||||
Fetch the next name, requesting another page only when needed.
|
||||
|
||||
#### Returns
|
||||
|
||||
`Promise`<`IteratorResult`<`string`, `any`>>
|
||||
|
||||
#### 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`
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
```
|
||||
|
||||
@@ -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;
|
||||
```
|
||||
@@ -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".
|
||||
|
||||
***
|
||||
|
||||
|
||||
@@ -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
@@ -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 |
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -646,6 +646,6 @@ export async function connectNamespace(
|
||||
export {
|
||||
Catalog,
|
||||
CatalogOptions,
|
||||
ListDatabasesResponse,
|
||||
DatabaseNames,
|
||||
connectCatalog,
|
||||
} from "./catalog";
|
||||
|
||||
@@ -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
@@ -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,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,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,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,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,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,6 +1,6 @@
|
||||
{
|
||||
"name": "@lancedb/lancedb-win32-arm64-msvc",
|
||||
"version": "0.40.0-beta.12",
|
||||
"version": "0.40.0-beta.13",
|
||||
"os": [
|
||||
"win32"
|
||||
],
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)))
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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 == []
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
@@ -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
@@ -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()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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
@@ -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,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
@@ -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
@@ -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(())
|
||||
}
|
||||
|
||||
+208
-2527
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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
|
||||
@@ -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
@@ -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());
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
///
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user