rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 14:52:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(sql): add native vector dedup and partitioned mv refresh (#4303)

Browse Source 
## Summary

Materialize a source-preserving dedup result from an indexed, pinned
snapshot with a short declaration. The default greedy policy keeps the
smallest unassigned snapshot row ID and removes only its directly
matching neighbors: For A < B < C, an A–B–C chain without an A–C edge
keeps A and C.

```python
from lancedb import vector_dedup

view = db.create_materialized_view(
    "images_clean",
    vector_dedup("images", column="phash", distance_threshold=4),
)
cleaned = view.table
```

The Python source expression captures the source version once at
creation and
waits for the materialized result by default. Use `dataset_version=...`
for an
explicit snapshot, `with_no_data=True` for declaration only, or
`create_materialized_view_async(...)` for a job handle. AsyncConnection
supports
the same expression. It uses the existing MV REST endpoints and registry
jobs,
without opening a SQL query connection. Source names are exact
root-namespace
table names. Additional select/where/limit options are rejected for this
source.

The SQL declaration remains available:

```sql
CREATE MATERIALIZED VIEW images_clean AS
SELECT * FROM vector_dedup('images', 7, 'phash', 4);
REFRESH MATERIALIZED VIEW images_clean;
```

This is the host-facing SQL entry point. The MV builder now accepts a
typed native vector source; the Rust SDK also supports the same
definition through `MaterializedViewDefinition::from_sql` and
`prepare_definition`; `view.refresh().execute()` executes it locally.
Hosts distribute its units using their existing job infrastructure. The
raw `vector_duplicate_pairs` source remains available separately when
callers want the pair relation.

## Implementation and contract

- Use actual Rust Lance `13.0.0-beta.15` (`db53dcf367`); the native
partition primitive computes pairs from index representations without
source-vector aggregation or a Python pair loop.
- Extend partitioned MV refresh with two dependent waves. Physical
segment/partition tasks stream pairs through a bounded external sort,
select direct representatives and write per-source-fragment bitmap
artifacts. Source-fragment tasks start after every selection receipt
completes and stream retained original rows, including blob payloads,
into uncommitted output fragments.
- Reuse normal blob write coercion and row-mask scans, including stable
row-ID translation after compaction. Null vectors, isolated rows and
source deletions are preserved correctly. Conflicting representative
decisions across overlapping index scopes fail before publication.
- Expose `dependency_units()` and `write_refresh_partition_with_inputs`
alongside the existing plan/commit API. Receipts bind
snapshot/configuration/policy/destination generation. Only a complete
accepted set can publish; unfinished units restart, completed units can
be reused. There is no pair-position checkpoint.
- Use MV format 4 for dedup results and format 3 for native pairs;
Python and TypeScript metadata readers recognize both. No source
deletion or source row-ID migration occurs.

Only same-segment/same-partition edges are covered. Quantized scores use
symmetric represented-vector distances. Result materialization requires
V2 data storage and is initially a native MV source, not an ordinary
SELECT table function. Compressed vertex masks stay in memory; edge
sorting spills to local disk. Unreferenced artifacts use normal Lance
retention/cleanup and must be retained for the job's recovery window.

## Validation

AWS Linux validation:

- Native pair/result integration matrix passed: multiple
segments/partitions, exact thresholds, empty output, cross-batch data,
single-process pair parity, source preservation, lost receipts and
publication guards.
- Blob/null-vector/pinned-deletion tests passed with both row-ID modes,
including compaction that changes physical row addresses.
- Representative policy, forced external-sort spill and
conflicting/corrupt artifact tests passed.
- Full MV regression suite: 153 passed, 2 pre-existing ignored. Rust
formatting and Clippy passed.

- Python MV suite: 47 passed; TypeScript MV suite: 11 passed. Native
bindings were rebuilt in project-local environments; formatting, lint
and TypeScript docs passed.
- Rust declaration doctests: 3 passed, including the complete local
dedup example.
- Companion host validation: 41 MV job tests passed (1 pre-existing
ignored), full workspace Clippy passed, and the complete process-level
CLI E2E suite passed (3 tests). Recovery exercises both selection and
materialization waves.
- A real CIFAR-10 pHash run used five actual registry-worker owners,
completed 32 selection + 24 materialization tasks, and retained 49,846
of 50,000 rows with identical original values and image payloads.

Host workflow performance measurements are recorded in the companion
host PR. Cross-platform CI is still running. The previous remote-test
bootstrap failed before Cargo tests with `401 Bad credentials`; the
workflow fix described below is being validated by CI. No deployment or
merge is performed.


## Python source validation

AWS tests cover automatic and explicit snapshots, physical/stable row
IDs, local
and remote creation, sync/job/async-connection APIs, delayed refresh,
strict
options and identifier quoting. The native Python extension was rebuilt
through
the project's uv/maturin backend. Ruff, Pyright and the full MkDocs
API-reference
build passed; both new public symbols are rendered.

A real 50,000-image FEv2 run disabled both Python SQL entry points and
supplied an
unreachable SQL endpoint. The new API still materialized the same 49,846
retained
rows and original image payloads. Three recorded creation/refresh jobs
completed
56 tasks each on all five registry-worker owners, and remained on source
snapshot
3 after source latest advanced to 4. No worker or server wire-protocol
changes
were needed for this client entry point.

Previous CI authentication failure:
https://github.com/lancedb/lancedb/actions/runs/36334679276/job/108663151265.

## Review and CI fixes

- Validate all JSON merge source fields before projection, rejecting
unknown and duplicate names at every struct level. Missing/reordered
fields and ordinary JSON merges remain supported. A checked-in 12-case
integration matrix covers root fields, structs, lists, large lists,
fixed-size lists and map values; rejected writes leave the table version
and row count unchanged.
- Match namespace connection MV method signatures to
`MaterializedViewSource` in all four sync/async methods, fixing the
Python 3.10 import failure from runtime override checks.
- Restore the remote-test script's encrypted-CLI path: use the Actions
token for artifact download and pass `SOPHON_READ_TOKEN` under the
variable the decryption script consumes. The previous workflow omitted
that variable and attempted the private release path, failing with HTTP
401. On the final head, remote CI passes: 1,464 Rust unit tests, the
integration targets and doctests pass (3 unit tests ignored).

AWS validation for these fixes: both review reproducers failed before
their fixes and pass afterward; 55 merge tests, 29 cast tests and the
standalone 7-case admission regression pass. Python 3.10.21 imports
successfully; 90 MV/namespace tests, 14 JSON tests and four real
namespace dedup cases pass. After `794b97a0`, the native binding was
rebuilt again and the JSON suite rerun. Rust formatting/Clippy, Ruff and
Pyright pass. Sophon's locked test compilation passes against the final
SDK source for all seven affected crates. Extended runners, logs and
benchmark artifacts remain local.

Gatekeeper recommends approval on `794b97a0`, and the `breaking-change`
label is set. Final CI: **49 checks passed, 3 expected skips, no
failures or pending checks**, including remote tests, Python 3.10,
minimum dependencies and Windows ARM Rust tests. All wheel builds pass
after retrying the Linux matrix affected by transient AlmaLinux
repository-metadata 404s. All code-review threads are resolved. The PR
remains open for maintainer review and merge.

---------

Co-authored-by: Yang Cen <yangcen@Yangs-Mac-mini.local>
Co-authored-by: Lance Release <lance-dev@lancedb.com>
This commit is contained in:
Yang Cen
2026-09-28 18:06:55 +08:00
committed by GitHub
co-authored by Yang Cen Lance Release
parent 139382b5f7
commit 0327618a62
27 changed files with 3497 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/rust.yml
+2 -1
View File
@@ -180,7 +180,8 @@ jobs:
env:
CC: clang-18
CXX: clang++-18
GH_TOKEN: ${{ secrets.SOPHON_READ_TOKEN }}
GH_TOKEN: ${{ github.token }}
SOPHON_READ_TOKEN: ${{ secrets.SOPHON_READ_TOKEN }}
steps:
- uses: actions/checkout@v6
with:
Cargo.lock
Generated
+46 -44
View File
@@ -3588,8 +3588,8 @@ checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c"
[[package]]
name = "fsst"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"rand 0.9.5",
@@ -5143,8 +5143,8 @@ checksum = "e037a2e1d8d5fdbd49b16a4ea09d5d6401c1f29eca5ff29d03d3824dba16256a"
[[package]]
name = "lance"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arc-swap",
"arrow",
@@ -5161,6 +5161,7 @@ dependencies = [
"async-trait",
"async_cell",
"aws-sdk-dynamodb",
"blake3",
"byteorder",
"bytes",
"chrono",
@@ -5216,8 +5217,8 @@ dependencies = [
[[package]]
name = "lance-arrow"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5239,7 +5240,7 @@ dependencies = [
[[package]]
name = "lance-arrow-scalar"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5253,7 +5254,7 @@ dependencies = [
[[package]]
name = "lance-arrow-stats"
version = "58.0.0"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5262,8 +5263,8 @@ dependencies = [
[[package]]
name = "lance-bitpacking"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrayref",
"crunchy",
@@ -5273,8 +5274,8 @@ dependencies = [
[[package]]
name = "lance-core"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5311,8 +5312,8 @@ dependencies = [
[[package]]
name = "lance-datafusion"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"arrow-array",
@@ -5343,8 +5344,8 @@ dependencies = [
[[package]]
name = "lance-datagen"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"arrow-array",
@@ -5361,8 +5362,8 @@ dependencies = [
[[package]]
name = "lance-derive"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"proc-macro2",
"quote",
@@ -5371,8 +5372,8 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5405,8 +5406,8 @@ dependencies = [
[[package]]
name = "lance-file"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5438,8 +5439,8 @@ dependencies = [
[[package]]
name = "lance-geo"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"datafusion",
"geo-traits",
@@ -5453,8 +5454,8 @@ dependencies = [
[[package]]
name = "lance-index"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arc-swap",
"arrow",
@@ -5524,8 +5525,8 @@ dependencies = [
[[package]]
name = "lance-index-core"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5543,12 +5544,13 @@ dependencies = [
"roaring",
"serde",
"serde_json",
"tokio",
]
[[package]]
name = "lance-io"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"arrow-array",
@@ -5590,8 +5592,8 @@ dependencies = [
[[package]]
name = "lance-linalg"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5605,8 +5607,8 @@ dependencies = [
[[package]]
name = "lance-namespace"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"async-trait",
@@ -5620,8 +5622,8 @@ dependencies = [
[[package]]
name = "lance-namespace-impls"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"arrow-ipc",
@@ -5674,8 +5676,8 @@ dependencies = [
[[package]]
name = "lance-select"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-buffer",
@@ -5689,8 +5691,8 @@ dependencies = [
[[package]]
name = "lance-table"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow",
"arrow-array",
@@ -5730,8 +5732,8 @@ dependencies = [
[[package]]
name = "lance-testing"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"arrow-array",
"arrow-schema",
@@ -5744,8 +5746,8 @@ dependencies = [
[[package]]
name = "lance-tokenizer"
version = "13.0.0-beta.14"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.14#b80ac1d9f3c24b47bd8e297d03053ea35d216a9f"
version = "13.0.0-beta.15"
source = "git+https://github.com/lance-format/lance.git?tag=v13.0.0-beta.15#db53dcf367088b32ae3a7461407d43593759ae3e"
dependencies = [
"frostem",
"icu_segmenter",
Cargo.toml
+14 -14
View File
@@ -13,20 +13,20 @@ categories = ["database-implementations"]
rust-version = "1.91.0"
[workspace.dependencies]
lance = { "version" = "=13.0.0-beta.14", default-features = false, "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=13.0.0-beta.14", default-features = false, "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=13.0.0-beta.14", default-features = false, "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=13.0.0-beta.14", "tag" = "v13.0.0-beta.14", "git" = "https://github.com/lance-format/lance.git" }
lance = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-core = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-datagen = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-file = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-io = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-index = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-linalg = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-namespace-impls = { "version" = "=13.0.0-beta.15", default-features = false, "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-table = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-testing = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-datafusion = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-encoding = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "git" = "https://github.com/lance-format/lance.git" }
lance-arrow = { "version" = "=13.0.0-beta.15", "tag" = "v13.0.0-beta.15", "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
docs/src/python/python.md
+27
View File
@@ -198,6 +198,33 @@ 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
nodejs/__test__/materialized_view.test.ts
+11 -1
View File
@@ -52,9 +52,19 @@ 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":3,"query":${JSON.stringify(query)}}`,
`{"format":5,"query":${JSON.stringify(query)}}`,
'{"kind":"select_v3","source_table":"people"}',
]) {
expect(() => read(newer)).toThrow(/cannot refresh/);
nodejs/lancedb/materialized_view.ts
+3 -2
View File
@@ -9,9 +9,10 @@ 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 2 being a query with `GROUP BY`, format 3 native pairs and format 4
* native dedup results.
*/
export const DEFINITION_FORMAT = 2;
export const DEFINITION_FORMAT = 4;
/**
* The query that defines a materialized view, as stored:
python/python/lancedb/__init__.py
+1
View File
@@ -46,6 +46,7 @@ from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewDefinition,
vector_dedup as vector_dedup,
)
from .view import ViewDescription as ViewDescription
from .table import AsyncTable, Table
python/python/lancedb/_lancedb.pyi
+4
View File
@@ -247,6 +247,8 @@ 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,
@@ -256,6 +258,8 @@ 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(
python/python/lancedb/db.py
+45 -10
View File
@@ -56,6 +56,8 @@ from .sql import QueryDescription
from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewSource,
VectorDedupSource,
SelectArg,
normalize_select,
)
@@ -555,7 +557,7 @@ class DBConnection(EnforceOverrides):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -569,17 +571,21 @@ class DBConnection(EnforceOverrides):
table. The view is a normal table: it can be queried, indexed and
searched, and it appears in ``table_names``.
The source table must have stable row ids (create it with the
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.
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
The name of the source table, in this database.
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.
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
@@ -602,7 +608,7 @@ class DBConnection(EnforceOverrides):
def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -1558,7 +1564,7 @@ class LanceDBConnection(DBConnection):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -1603,7 +1609,7 @@ class LanceDBConnection(DBConnection):
def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -2498,10 +2504,31 @@ class AsyncConnection(object):
await tbl.checkout(version)
return tbl
async def _materialized_view_source(
self,
source: MaterializedViewSource,
select: SelectArg,
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)
}
async def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -2512,6 +2539,9 @@ class AsyncConnection(object):
See
[DBConnection.create_materialized_view][lancedb.DBConnection.create_materialized_view].
"""
source, native_options = await self._materialized_view_source(
source, select, where, limit
)
inner = await self._inner.create_materialized_view(
name,
source,
@@ -2519,13 +2549,14 @@ class AsyncConnection(object):
filter=where,
limit=limit,
with_no_data=with_no_data,
**native_options,
)
return AsyncMaterializedView(AsyncTable(inner))
async def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -2536,6 +2567,9 @@ class AsyncConnection(object):
Wait for the returned job before opening or querying the view.
"""
source, native_options = await self._materialized_view_source(
source, select, where, limit
)
inner = await self._inner.create_materialized_view_async(
name,
source,
@@ -2543,6 +2577,7 @@ class AsyncConnection(object):
filter=where,
limit=limit,
with_no_data=with_no_data,
**native_options,
)
return AsyncJob(inner)
python/python/lancedb/materialized_view.py
+99 -2
View File
@@ -7,6 +7,7 @@ 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
@@ -20,6 +21,7 @@ if TYPE_CHECKING:
from .table import AsyncTable, LanceTable
DEFINITION_META_KEY = b"mv.definition"
_MAX_FLOAT32 = float.fromhex("0x1.fffffep+127")
SelectArg = Union[
str,
@@ -29,9 +31,104 @@ SelectArg = Union[
]
DEFINITION_FORMAT = 2
@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,
)
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
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
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
readers older than the format number."""
python/python/lancedb/namespace.py
+5 -4
View File
@@ -68,6 +68,7 @@ from lance_namespace import (
from lancedb.materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewSource,
SelectArg,
)
from lancedb.table import AsyncTable, LanceTable, Table
@@ -632,7 +633,7 @@ class LanceNamespaceDBConnection(DBConnection):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
@@ -662,7 +663,7 @@ class LanceNamespaceDBConnection(DBConnection):
def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
@@ -1242,7 +1243,7 @@ class AsyncLanceNamespaceDBConnection:
async def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
@@ -1265,7 +1266,7 @@ class AsyncLanceNamespaceDBConnection:
async def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: "SelectArg" = None,
where: Optional[str] = None,
python/python/lancedb/remote/db.py
+3 -3
View File
@@ -39,7 +39,7 @@ from ..functions import FunctionVersion, UdfDefinition
from ..job import AsyncJob, Job
from ..sql import Query as SqlQuery
from ..sql import QueryDescription
from ..materialized_view import MaterializedView, SelectArg
from ..materialized_view import MaterializedView, MaterializedViewSource, SelectArg
from ..secrets import EnvVarSecret, SecretInfo
from ..view import ViewDescription
@@ -713,7 +713,7 @@ class RemoteDBConnection(DBConnection):
def create_materialized_view(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
@@ -745,7 +745,7 @@ class RemoteDBConnection(DBConnection):
def create_materialized_view_async(
self,
name: str,
source: str,
source: MaterializedViewSource,
*,
select: SelectArg = None,
where: Optional[str] = None,
python/python/tests/test_materialized_views.py
+157 -1
View File
@@ -520,10 +520,166 @@ 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": 3, "query": query},
{"format": 5, "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,
},
)
]
python/src/connection.rs
+18 -2
View File
@@ -407,7 +407,8 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -416,10 +417,17 @@ 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);
}
@@ -435,7 +443,8 @@ impl Connection {
})
}
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
@@ -444,10 +453,17 @@ 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);
}
rust/lancedb/src/materialized_view.rs
+183 -28
View File
@@ -9,8 +9,12 @@
//! query this version cannot maintain reads back as unrefreshable, not as a
//! plain table. Queries, indexes and search work on the view unchanged.
mod duplicate_pairs;
mod grouped;
mod grouped_units;
mod partitioned;
mod vector_dedup;
pub use duplicate_pairs::{VectorSource, VectorSourceKind};
pub use grouped::IVF_PARTITION;
/// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index
/// partition, so the work can be spread over several processes. The caller
@@ -74,6 +78,10 @@ pub use grouped_units::{
GroupedRefreshPlan, WrittenUnit, commit_grouped_refresh, plan_grouped_refresh,
write_grouped_unit,
};
pub use partitioned::{
PartitionedRefreshPlan, WrittenPartition, commit_partitioned_refresh, plan_partitioned_refresh,
write_refresh_partition, write_refresh_partition_with_inputs,
};
mod query;
pub mod refresh;
@@ -146,18 +154,26 @@ const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions";
/// The newest layout this version reads under [`DEFINITION_META_KEY`]:
/// `{"format": N, "query": "<SQL>"}`, the query as
/// [`MaterializedViewDefinition::to_sql`] renders it. Format 2 is a query
/// with `GROUP BY`; every other query is written as format 1, so a reader
/// that predates grouping reports a grouped view as unrefreshable rather
/// than failing to parse it. A reader refuses a newer format rather than
/// [`MaterializedViewDefinition::to_sql`] renders it. Format 3 identifies a
/// native duplicate-pair source, format 4 a dedup result source, format 2 a query with `GROUP BY`, and
/// format 1 an ordinary query. A reader refuses a newer format rather than
/// guess at it. The layout also carries `"kind": "query"`, which readers
/// older than the format number report as an unrefreshable view instead of
/// failing to read the metadata.
pub const DEFINITION_FORMAT: u64 = 2;
pub const DEFINITION_FORMAT: u64 = 4;
/// The format `definition` is written in; see [`DEFINITION_FORMAT`].
fn format_of(definition: &MaterializedViewDefinition) -> u64 {
if definition.is_grouped() { 2 } else { 1 }
if let Some(source) = &definition.vector_source {
match source.kind {
VectorSourceKind::Pairs => 3,
VectorSourceKind::Dedup => 4,
}
} else if definition.is_grouped() {
2
} else {
1
}
}
/// Legacy `kind` tag of the structured layout written before
@@ -300,6 +316,8 @@ pub fn read_staging(metadata: &HashMap<String, String>) -> Result<Option<Staging
/// for the shape.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MaterializedViewDefinition {
/// A table-bound native vector source. It is never evaluated per source row.
pub vector_source: Option<Box<VectorSource>>,
/// Name of the source table, in the same database as the view.
pub source_table: String,
/// Namespace path holding the source table; empty is the root namespace.
@@ -498,6 +516,7 @@ pub fn read_definition(metadata: &HashMap<String, String>) -> Result<Option<Stor
)));
}
Ok(Some(StoredDefinition::Query(MaterializedViewDefinition {
vector_source: None,
source_table: legacy.source_table,
source_namespace: legacy.source_namespace,
lateral: None,
@@ -554,6 +573,9 @@ pub(crate) fn plan(
definition: &MaterializedViewDefinition,
staging: Option<&StagingBinding>,
) -> Result<Planned> {
if definition.vector_source.is_some() {
return duplicate_pairs::plan(source_schema, definition);
}
let filter = definition
.filter
.as_deref()
@@ -712,6 +734,7 @@ pub(crate) fn plan(
}
let definition = MaterializedViewDefinition {
vector_source: None,
source_table: definition.source_table.clone(),
source_namespace: definition.source_namespace.clone(),
lateral: definition.lateral.clone(),
@@ -1145,6 +1168,11 @@ impl PreparedDeclaration {
/// read: the column the view projects it to, if any, otherwise an
/// internal projection added here, named by [`input_column_name`].
pub fn input_column(&mut self, source_column: &str) -> Result<String> {
if self.definition.vector_source.is_some() {
return Err(Error::InvalidInput {
message: "native vector views have no per-source-row inputs".into(),
});
}
// A grouped view's rows are groups; no source row carries a value into one.
if self.definition.is_grouped() {
return Err(Error::InvalidInput {
@@ -1245,6 +1273,11 @@ impl PreparedDeclaration {
columns: Vec<(usize, ArrowField)>,
bindings: &[FunctionBinding],
) -> Result<Self> {
if self.definition.vector_source.is_some() {
return Err(Error::InvalidInput {
message: "computed columns on native vector views are not supported".into(),
});
}
let invalid = |message: String| Error::InvalidInput { message };
if columns.is_empty() {
return Err(invalid("at least one computed column is needed".into()));
@@ -1451,6 +1484,7 @@ pub async fn prepare_declaration(
limit: Option<u64>,
) -> Result<PreparedDeclaration> {
let definition = MaterializedViewDefinition {
vector_source: None,
source_table: source.name().to_string(),
source_namespace: source.namespace().to_vec(),
lateral: None,
@@ -1512,6 +1546,27 @@ pub async fn prepare_declaration(
/// # Ok(())
/// # }
/// ```
///
/// An indexed, pinned source can materialize retained original rows using the
/// default greedy direct-representative policy. A row is removed only when it
/// directly matches a retained row; an A-B-C chain keeps A and C. The source is
/// unchanged. Hosts distribute this through [`plan_partitioned_refresh`] and
/// [`write_refresh_partition_with_inputs`], with a barrier between selection
/// tasks and source-fragment materialization tasks.
///
/// ```
/// # #![recursion_limit = "256"]
/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition};
/// # async fn dedup(images: &lancedb::Table) -> Result<(), Box<dyn std::error::Error>> {
/// let version = images.version().await?;
/// let definition = MaterializedViewDefinition::from_sql(&format!(
/// "SELECT * FROM vector_dedup('images', {version}, 'phash', 4)"
/// ))?;
/// let clean = prepare_definition(images, definition).await?.create("images_clean").await?;
/// clean.refresh().execute().await?;
/// # Ok(())
/// # }
/// ```
pub async fn prepare_definition(
source: &Table,
definition: MaterializedViewDefinition,
@@ -1627,7 +1682,9 @@ async fn prepare_with(
),
});
}
if !native.dataset.get().await?.manifest.uses_stable_row_ids() {
if definition.vector_source.is_none()
&& !native.dataset.get().await?.manifest.uses_stable_row_ids()
{
return Err(Error::InvalidInput {
message: format!(
"materialized views require stable row ids on the source table; \
@@ -1653,7 +1710,22 @@ async fn prepare_with(
message: format!("view column name '{}' is reserved", reserved.output),
});
}
let source_schema = resolved.schema().await?;
let source_schema = if let Some(pairs) = &definition.vector_source {
let ds = native
.dataset
.get()
.await?
.checkout_version(pairs.dataset_version)
.await?;
crate::table::datafusion::udtf::duplicate_pairs::plan_duplicate_pairs(
Arc::new(ds.clone()),
&pairs.config()?,
)
.await?;
Arc::new(ArrowSchema::from(ds.schema()))
} else {
resolved.schema().await?
};
let source_metadata = source_schema.metadata().clone();
let Planned {
definition,
@@ -1670,11 +1742,13 @@ async fn prepare_with(
None => source_schema.clone(),
Some(unnest) => flattened_schema(&source_schema, &unnest)?,
};
fields.push(ArrowField::new(
SOURCE_ROW_ID_COLUMN,
DataType::UInt64,
false,
));
if definition.vector_source.is_none() {
fields.push(ArrowField::new(
SOURCE_ROW_ID_COLUMN,
DataType::UInt64,
false,
));
}
// Only column-describing metadata comes along: structural declarations
// describe how a table is written, and a view is written by refresh alone.
let mut metadata: HashMap<String, String> = HashMap::new();
@@ -1718,6 +1792,7 @@ pub struct CreateMaterializedViewBuilder {
namespace: Vec<String>,
source: String,
source_namespace: Vec<String>,
vector_source: Option<VectorSource>,
projections: Vec<(String, String)>,
filter: Option<String>,
limit: Option<u64>,
@@ -1732,6 +1807,7 @@ impl CreateMaterializedViewBuilder {
namespace: Vec::new(),
source,
source_namespace: Vec::new(),
vector_source: None,
projections: Vec::new(),
filter: None,
limit: None,
@@ -1752,6 +1828,38 @@ impl CreateMaterializedViewBuilder {
self
}
/// Read native pairs or retained original rows from a fixed indexed snapshot.
/// This source selects its full output schema and cannot be combined with
/// projections, filters or limits. Remote creation uses the ordinary MV
/// endpoint and its registry job; it does not execute a SQL query client-side.
pub fn vector_source(mut self, source: VectorSource) -> Self {
self.vector_source = Some(source);
self
}
fn native_definition(&self) -> Result<Option<MaterializedViewDefinition>> {
let Some(source) = &self.vector_source else {
return Ok(None);
};
source.config()?;
if !self.projections.is_empty() || self.filter.is_some() || self.limit.is_some() {
return Err(Error::InvalidInput {
message: "native vector sources cannot be combined with select, where or limit"
.into(),
});
}
Ok(Some(MaterializedViewDefinition {
vector_source: Some(Box::new(source.clone())),
source_table: self.source.clone(),
source_namespace: self.source_namespace.clone(),
projections: vec![ViewProjection::star()],
filter: None,
lateral: None,
group_by: Vec::new(),
limit: None,
}))
}
/// The view's columns, as `(name, SQL expression)` pairs. Not calling
/// this selects every source column, expanded at creation time.
pub fn select(
@@ -1784,7 +1892,10 @@ impl CreateMaterializedViewBuilder {
self
}
fn query(&self) -> String {
fn query(&self) -> Result<String> {
if let Some(definition) = self.native_definition()? {
return Ok(definition.to_sql());
}
fn quote(name: &str) -> String {
format!("\"{}\"", name.replace('"', "\"\""))
}
@@ -1813,12 +1924,13 @@ impl CreateMaterializedViewBuilder {
if let Some(limit) = self.limit {
query.push_str(&format!(" LIMIT {limit}"));
}
query
Ok(query)
}
/// Submit creation and initial population, returning a [`Job`] that
/// settles when the view is ready. The source must keep stable row ids --
/// they hold provenance across compaction, and cannot be enabled later.
/// settles when the view is ready. Ordinary sources must keep stable row
/// ids for provenance across compaction. Native vector sources also support
/// physical row IDs because their definitions pin the dataset version.
pub async fn execute_async(self) -> Result<Job> {
if self.connection.uri().starts_with("db://") {
return self
@@ -1827,7 +1939,7 @@ impl CreateMaterializedViewBuilder {
.create_materialized_view_async(CreateMaterializedViewRequest {
name: self.name.clone(),
namespace_path: self.namespace.clone(),
query: self.query(),
query: self.query()?,
with_no_data: self.with_no_data,
})
.await;
@@ -1855,19 +1967,24 @@ impl CreateMaterializedViewBuilder {
}
async fn execute_native(self) -> Result<MaterializedView> {
let native_definition = self.native_definition()?;
let source = self
.connection
.open_table(&self.source)
.namespace(self.source_namespace.clone())
.execute()
.await?;
let prepared = prepare_declaration(
&source,
(!self.projections.is_empty()).then_some(self.projections.as_slice()),
self.filter.as_deref(),
self.limit,
)
.await?;
let prepared = if let Some(definition) = native_definition {
prepare_definition(&source, definition).await?
} else {
prepare_declaration(
&source,
(!self.projections.is_empty()).then_some(self.projections.as_slice()),
self.filter.as_deref(),
self.limit,
)
.await?
};
let view = prepared.create_in(&self.namespace, &self.name).await?;
if !self.with_no_data {
view.refresh().execute().await?;
@@ -2186,6 +2303,7 @@ mod tests {
assert_eq!(
view.definition(),
&MaterializedViewDefinition {
vector_source: None,
source_table: "people".into(),
source_namespace: Vec::new(),
projections: vec![
@@ -2309,13 +2427,49 @@ mod tests {
.select([("double\"age", "age * 2")])
.only_if("age >= 18")
.limit(10)
.query();
.query()
.unwrap();
assert_eq!(
query,
"SELECT age * 2 AS \"double\"\"age\" FROM \"raw data\".\"odd\"\"source\" WHERE age >= 18 LIMIT 10"
);
}
#[tokio::test]
async fn native_vector_builder_preserves_source_and_rejects_extra_clauses() {
let conn = people_db().await;
let source = VectorSource {
kind: VectorSourceKind::Dedup,
dataset_version: 7,
column: "vector'field".into(),
distance_threshold: "4.0".into(),
};
let builder = || {
conn.create_materialized_view("clean", "odd.source'name")
.source_namespace(vec!["raw space".into()])
.vector_source(source.clone())
};
let query = builder().query().unwrap();
let definition = MaterializedViewDefinition::from_sql(&query).unwrap();
assert_eq!(definition.source_table, "odd.source'name");
assert_eq!(definition.source_namespace, ["raw space"]);
assert_eq!(definition.vector_source.as_deref(), Some(&source));
assert!(definition.selects_star());
assert!(builder().select([("id", "id")]).query().is_err());
assert!(builder().only_if("true").query().is_err());
assert!(builder().limit(1).query().is_err());
let invalid_source = VectorSource {
dataset_version: 0,
..source
};
assert!(
conn.create_materialized_view("bad", "people")
.vector_source(invalid_source)
.query()
.is_err()
);
}
/// No projection selects every source column, expanded now: the schema
/// captured at creation is the definition.
#[tokio::test]
@@ -3235,6 +3389,7 @@ mod tests {
fn definition(source_namespace: Vec<String>) -> MaterializedViewDefinition {
MaterializedViewDefinition {
vector_source: None,
source_table: "people".to_string(),
source_namespace,
lateral: None,
@@ -3321,8 +3476,8 @@ mod tests {
#[test]
fn a_newer_format_is_reported_not_guessed() {
assert_eq!(
read(r#"{"format":3,"query":"SELECT name FROM people"}"#).unwrap(),
Some(StoredDefinition::Newer { format: "3".into() })
read(r#"{"format":5,"query":"SELECT name FROM people"}"#).unwrap(),
Some(StoredDefinition::Newer { format: "5".into() })
);
assert_eq!(
read(r#"{"kind":"join"}"#).unwrap(),
rust/lancedb/src/materialized_view/duplicate_pairs.rs
+478
View File
@@ -0,0 +1,478 @@
// 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
}
rust/lancedb/src/materialized_view/grouped_units.rs
+42 -8
View File
@@ -53,6 +53,7 @@ const PROBE_ROWS: usize = 256;
/// How a grouped view can be refreshed in units.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct GroupedRefreshPlan {
/// Units `0..units`: unit `p` before the last holds index partition
/// `p`'s groups; the last holds the rows the index cannot place (null or
@@ -177,6 +178,7 @@ pub async fn plan_grouped_refresh(
/// accounted for has to be refusable. Travels to whoever commits, so it
/// serializes, but its fields are not a caller's to write.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenUnit {
unit: u32,
plan: GroupedRefreshPlan,
@@ -491,10 +493,38 @@ pub async fn commit_grouped_refresh(
plan: &GroupedRefreshPlan,
units: Vec<WrittenUnit>,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let grouped = open_planned(view, plan).await?;
let units = complete_set(view.name(), plan, units)?;
let rows_written = units.iter().map(|unit| unit.rows).sum();
let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect();
commit_fragments(
view,
plan,
&grouped.definition,
grouped.source.manifest.timestamp_nanos,
new_fragments,
rows_written,
expected_incarnation,
)
.await
}
/// The common all-or-nothing unit sink, shared by grouped and native views.
#[allow(clippy::too_many_arguments)]
pub(super) async fn commit_fragments(
view: &Table,
plan: &GroupedRefreshPlan,
definition: &MaterializedViewDefinition,
source_ts: u128,
new_fragments: Vec<Fragment>,
rows_written: u64,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let view_native = view.as_native().ok_or_else(|| Error::NotSupported {
message: "materialized views are supported only on local tables".into(),
})?;
view_native.dataset.ensure_mutable()?;
let lock = refresh_lock(view_native.dataset.get().await?.uri());
let _guard = lock.lock().await;
if let Some(expected) = expected_incarnation
@@ -508,10 +538,17 @@ pub async fn commit_grouped_refresh(
),
});
}
let grouped = open_planned(view, plan).await?;
let source_ts = grouped.source.manifest.timestamp_nanos;
let ds = Arc::new(grouped.view_ds);
view_native.dataset.reload().await?;
let ds = view_native.dataset.get().await?.as_ref().clone();
ensure_incarnation(&ds, Some(&plan.incarnation), view.name()).await?;
if super::read_definition(&ds.schema().metadata)?
!= Some(super::StoredDefinition::Query(definition.clone()))
{
return Err(Error::InvalidInput {
message: "view definition changed while its units were written".into(),
});
}
let ds = Arc::new(ds);
if ds.version().version != plan.view_version {
return Err(Error::Runtime {
message: format!(
@@ -523,11 +560,8 @@ pub async fn commit_grouped_refresh(
),
});
}
let units = complete_set(view.name(), plan, units)?;
let removed_fragment_ids: Vec<u64> = ds.get_fragments().iter().map(|f| f.id() as u64).collect();
// Fragment ids are assigned at commit; the units' fragments all carry 0.
let rows_written = units.iter().map(|unit| unit.rows).sum();
let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect();
let transaction = Transaction::new(
plan.view_version,
Operation::Update {
@@ -566,7 +600,7 @@ pub async fn commit_grouped_refresh(
committed,
plan.source_version,
source_ts,
Some(&grouped.definition),
Some(definition),
expected_incarnation,
)
.await?;
rust/lancedb/src/materialized_view/partitioned.rs
+192
View File
@@ -0,0 +1,192 @@
// 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
}
}
}
rust/lancedb/src/materialized_view/query.rs
+84 -1
View File
@@ -264,7 +264,36 @@ 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: \
@@ -416,6 +445,7 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let definition = MaterializedViewDefinition {
vector_source,
source_table,
source_namespace,
lateral,
@@ -424,6 +454,9 @@ 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)
}
@@ -483,7 +516,18 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
.chain(std::iter::once(&definition.source_table))
.map(|part| ident_sql(part))
.collect();
sql.push_str(&table.join("."));
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("."));
}
if let Some(lateral) = &definition.lateral {
match &lateral.source {
LateralSource::Unnest { column } => {
@@ -511,6 +555,45 @@ 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 [
rust/lancedb/src/materialized_view/refresh.rs
+23 -1
View File
@@ -117,6 +117,26 @@ pub(crate) async fn execute_refresh(
pinned: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
// Native table functions use the same partition writer and commit sink in
// local execution; Sophon drives these units through its job registry.
if let Some(super::StoredDefinition::Query(definition)) =
super::read_definition(view.schema().await?.metadata())?
&& definition.vector_source.is_some()
{
let plan = super::plan_partitioned_refresh(view, pinned)
.await?
.ok_or_else(|| Error::Runtime {
message: "native pair view has no partition plan".into(),
})?;
let mut units = Vec::with_capacity(plan.units() as usize);
for unit in 0..plan.units() {
let dependencies = &units[..units.len().min(plan.dependency_units() as usize)];
units.push(
super::write_refresh_partition_with_inputs(view, unit, &plan, dependencies).await?,
);
}
return super::commit_partitioned_refresh(view, &plan, units, expected_incarnation).await;
}
let view_native = view.as_native().ok_or_else(|| Error::NotSupported {
message: "materialized views are supported only on local tables".into(),
})?;
@@ -728,7 +748,7 @@ pub(super) async fn open_source(
message: "materialized views are supported only on local tables".into(),
})?;
let dataset = native.dataset.get().await?.as_ref().clone();
if !dataset.manifest.uses_stable_row_ids() {
if definition.vector_source.is_none() && !dataset.manifest.uses_stable_row_ids() {
return Err(Error::InvalidInput {
message: format!(
"source table '{}' does not have stable row ids; it is not the \
@@ -3674,6 +3694,7 @@ pub(crate) mod tests {
let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await;
let replacement = crate::materialized_view::MaterializedViewDefinition {
vector_source: None,
source_table: "src".into(),
source_namespace: Vec::new(),
projections: vec![
@@ -3715,6 +3736,7 @@ pub(crate) mod tests {
let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await;
let narrower = crate::materialized_view::MaterializedViewDefinition {
vector_source: None,
source_table: "src".into(),
source_namespace: Vec::new(),
projections: vec![crate::materialized_view::ViewProjection {
rust/lancedb/src/materialized_view/vector_dedup.rs
+636
View File
@@ -0,0 +1,636 @@
// 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.
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));
}
}
rust/lancedb/src/remote/table.rs
+1
View File
@@ -2140,6 +2140,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
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,
rust/lancedb/src/table/add_data.rs
+15 -3
View File
@@ -1019,14 +1019,23 @@ mod tests {
#[case::utf8(DataType::Utf8)]
#[case::large_utf8(DataType::LargeUtf8)]
#[tokio::test]
async fn test_add_arrow_json_into_lance_json_table(#[case] input_type: DataType) {
async fn test_add_arrow_json_into_lance_json_table(
#[case] input_type: DataType,
#[values(false, true)] table_has_serialized_metadata: bool,
#[values(false, true)] input_has_serialized_metadata: bool,
) {
use arrow_array::{Array, cast::AsArray};
use lance_arrow::ARROW_EXT_NAME_KEY;
use lance_arrow::json::{ARROW_JSON_EXT_NAME, JSON_EXT_NAME};
// Build a table whose "data" column is lance.json (LargeBinary +
// ARROW:extension:name = "lance.json").
let lance_json_field = lance_arrow::json::json_field("data", true);
let mut lance_json_field = lance_arrow::json::json_field("data", true);
if table_has_serialized_metadata {
let mut metadata = lance_json_field.metadata().clone();
metadata.insert(lance_arrow::ARROW_EXT_META_KEY.to_string(), String::new());
lance_json_field = lance_json_field.with_metadata(metadata);
}
let table_schema = Arc::new(Schema::new(vec![lance_json_field]));
let db = connect("memory://").execute().await.unwrap();
@@ -1050,10 +1059,13 @@ mod tests {
// Build an arrow.json input field (Utf8/LargeUtf8 + arrow.json extension).
// This is what PyArrow produces for pa.json_() arrays.
let arrow_json_metadata = std::collections::HashMap::from([(
let mut arrow_json_metadata = std::collections::HashMap::from([(
ARROW_EXT_NAME_KEY.to_string(),
ARROW_JSON_EXT_NAME.to_string(),
)]);
if input_has_serialized_metadata {
arrow_json_metadata.insert(lance_arrow::ARROW_EXT_META_KEY.to_string(), String::new());
}
let arrow_json_field =
Field::new("data", input_type.clone(), true).with_metadata(arrow_json_metadata);
let arrow_json_schema = Arc::new(Schema::new(vec![arrow_json_field]));
rust/lancedb/src/table/datafusion/udtf.rs
+2
View File
@@ -4,3 +4,5 @@
//! User-Defined Table Functions (UDTFs) for DataFusion integration
pub mod fts;
pub mod duplicate_pairs;
rust/lancedb/src/table/datafusion/udtf/duplicate_pairs.rs
+586
View File
@@ -0,0 +1,586 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Index-partition duplicate pairs. No source-vector scan or shuffle is needed.
use std::collections::BTreeSet;
use std::sync::Arc;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::catalog::TableFunctionImpl;
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::{DataFusionError, Result, ScalarValue, plan_err};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::{Expr, TableType};
use datafusion_physical_expr::EquivalenceProperties;
use datafusion_physical_plan::{
DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
execution_plan::{Boundedness, EmissionType},
metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet},
projection::ProjectionExec,
stream::RecordBatchStreamAdapter,
};
use futures::{StreamExt, TryStreamExt, stream};
use lance::{
Dataset,
index::{DatasetIndexExt, DatasetIndexInternalExt},
};
use lance_index::metrics::NoOpMetricsCollector;
use uuid::Uuid;
use super::super::BaseTableAdapter;
use crate::table::NativeTableExt;
// Admission estimates for Lance 13's bounded pairwise kernel. Keep these with
// the minimum dependency when the kernel's buffering contract changes.
/// Version of the native task descriptor and its execution semantics.
pub const DUPLICATE_PAIRS_OPERATOR_VERSION: u32 = 2;
const CODE_STAGING_BYTES: usize = 16 * 1024 * 1024;
const INDEX_READ_ROWS: usize = 8192;
const VECTOR_BATCH_ROWS: usize = 8192;
const VECTOR_BATCH_BYTES: usize = 16 * 1024 * 1024;
const MIN_VECTOR_BATCH_ROWS: usize = 32;
const ANCHOR_BLOCK_ROWS: usize = 32;
const SCORING_CONCURRENCY: usize = 4;
const MAX_COORDINATE_BYTES: usize = 8;
const STAGED_ROW_OVERHEAD_BYTES: usize = 64;
const PAIR_ROW_BYTES: usize = 20;
/// Configuration belongs to one immutable source snapshot. Threshold uses
/// Lance's float32 index distance, including symmetric reconstructed distances
/// for quantized indexes; it is not an exact source-embedding distance.
#[derive(Debug, Clone)]
pub struct DuplicatePairsConfig {
pub dataset_version: u64,
pub column: String,
pub distance_threshold: f32,
}
/// One worker input, never a row of source vectors. The complete task identity
/// also includes the dataset URI/version, column, threshold and operator version.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DuplicatePairTask {
pub segment_id: Uuid,
pub partition_id: usize,
}
/// Resolve an authorized source table and preserve its catalog identity.
pub trait DuplicatePairsResolver: std::fmt::Debug + Send + Sync {
fn resolve(
&self,
table: &str,
config: DuplicatePairsConfig,
tasks_only: bool,
) -> Result<Arc<dyn TableProvider>>;
}
/// `vector_duplicate_pairs('table', version, 'column', threshold)`.
/// Registration is explicit so the host controls catalog access and execution.
#[derive(Debug)]
pub struct DuplicatePairsTableFunction {
resolver: Arc<dyn DuplicatePairsResolver>,
tasks_only: bool,
}
impl DuplicatePairsTableFunction {
pub fn new(resolver: Arc<dyn DuplicatePairsResolver>) -> Self {
Self {
resolver,
tasks_only: false,
}
}
}
impl DuplicatePairsTableFunction {
/// Register as `vector_duplicate_pair_tasks` to inspect or persist exactly
/// the descriptors used by the pair operator, without producing any pairs.
pub fn task_manifest(resolver: Arc<dyn DuplicatePairsResolver>) -> Self {
Self {
resolver,
tasks_only: true,
}
}
}
impl TableFunctionImpl for DuplicatePairsTableFunction {
fn call(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
if args.len() != 4 {
return plan_err!(
"vector_duplicate_pairs requires (table, dataset_version, column, distance_threshold)"
);
}
let string = |expr: &Expr| match expr {
Expr::Literal(ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)), _) => {
Ok(s.clone())
}
_ => plan_err!("table and column must be string literals"),
};
let dataset_version = match &args[1] {
Expr::Literal(ScalarValue::Int64(Some(v)), _) if *v > 0 => *v as u64,
Expr::Literal(ScalarValue::UInt64(Some(v)), _) if *v > 0 => *v,
_ => return plan_err!("dataset_version must be a positive integer literal"),
};
let distance_threshold = match &args[3] {
Expr::Literal(value, _) if value.data_type().is_numeric() => {
value.cast_to(&DataType::Float32)?
}
_ => return plan_err!("distance_threshold must be a finite numeric literal"),
};
let ScalarValue::Float32(Some(distance_threshold)) = distance_threshold else {
return plan_err!("distance_threshold cannot be null");
};
if !distance_threshold.is_finite() {
return plan_err!("distance_threshold must be finite");
}
self.resolver.resolve(
&string(&args[0])?,
DuplicatePairsConfig {
dataset_version,
column: string(&args[2])?,
distance_threshold,
},
self.tasks_only,
)
}
}
/// Three non-null columns; pair direction is storage traversal order, not row-ID order.
pub fn duplicate_pair_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("row_id_a", DataType::UInt64, false),
Field::new("row_id_b", DataType::UInt64, false),
Field::new("distance", DataType::Float32, false),
]))
}
#[derive(Debug)]
pub struct DuplicatePairsTable {
source: Arc<dyn crate::table::BaseTable>,
config: DuplicatePairsConfig,
tasks_only: bool,
}
impl BaseTableAdapter {
/// Bind the builtin to this authorized table; only metadata is read while
/// planning. The source snapshot must remain retained until execution ends.
pub fn duplicate_pairs(
&self,
config: DuplicatePairsConfig,
tasks_only: bool,
) -> DuplicatePairsTable {
DuplicatePairsTable {
source: self.table.clone(),
config,
tasks_only,
}
}
}
#[async_trait]
impl TableProvider for DuplicatePairsTable {
fn schema(&self) -> SchemaRef {
if self.tasks_only {
duplicate_pair_task_schema()
} else {
duplicate_pair_schema()
}
}
fn table_type(&self) -> TableType {
TableType::View
}
async fn scan(
&self,
_state: &dyn Session,
projection: Option<&Vec<usize>>,
_filters: &[Expr],
_limit: Option<usize>,
) -> Result<Arc<dyn ExecutionPlan>> {
let native = self.source.as_native().ok_or_else(|| {
DataFusionError::Plan("duplicate pairs requires a native indexed table".into())
})?;
let current = native
.dataset
.get()
.await
.map_err(|e| DataFusionError::External(Box::new(e)))?;
let dataset = Arc::new(
current
.checkout_version(self.config.dataset_version)
.await?,
);
let tasks = plan_duplicate_pairs(dataset.clone(), &self.config).await?;
if self.tasks_only {
let batch = duplicate_pair_task_batch(&dataset, &self.config, &tasks)?;
let table =
datafusion::datasource::MemTable::try_new(batch.schema(), vec![vec![batch]])?;
return table.scan(_state, projection, _filters, _limit).await;
}
let plan: Arc<dyn ExecutionPlan> = Arc::new(DuplicatePairsExec::try_new(
dataset,
self.config.clone(),
tasks,
)?);
match projection {
Some(indices) => {
let schema = plan.schema();
let exprs = indices
.iter()
.map(|&i| {
(
Arc::new(datafusion_physical_expr::expressions::Column::new(
schema.field(i).name(),
i,
)) as _,
schema.field(i).name().clone(),
)
})
.collect::<Vec<_>>();
Ok(Arc::new(ProjectionExec::try_new(exprs, plan)?))
}
None => Ok(plan),
}
}
}
/// Enumerate segment-local IVF partitions using metadata only. Live-row masks
/// and index codes are opened by the scoped worker, never by a source-vector
/// scan on the coordinator. The high-level pair API is only a test oracle.
pub async fn plan_duplicate_pairs(
dataset: Arc<Dataset>,
config: &DuplicatePairsConfig,
) -> Result<Vec<DuplicatePairTask>> {
if dataset.version().version != config.dataset_version || !config.distance_threshold.is_finite()
{
return plan_err!("duplicate pairs source snapshot or threshold is invalid");
}
let field = dataset.schema().field_id(&config.column)?;
let indices = dataset.load_indices().await?;
let segments = indices
.iter()
.filter(|meta| {
meta.keyed_fields() == [field]
&& meta.index_details.as_ref().map_or_else(
|| {
meta.files.as_ref().is_some_and(|files| {
files
.iter()
.any(|file| file.path == lance_index::INDEX_FILE_NAME)
})
},
|details| details.type_url.ends_with("VectorIndexDetails"),
)
})
.collect::<Vec<_>>();
let names = segments
.iter()
.map(|meta| meta.name.as_str())
.collect::<BTreeSet<_>>();
if names.len() != 1 {
return plan_err!(
"column '{}' requires exactly one current-format vector index, found {}",
config.column,
names.len()
);
}
let name = names.into_iter().next().expect("one index name");
if !dataset.unindexed_fragments(name).await?.is_empty() {
return plan_err!(
"vector index on column '{}' does not cover all fragments; optimize the index first",
config.column
);
}
let mut tasks = Vec::new();
for meta in segments {
if dataset.fragments().iter().any(|fragment| {
meta.fragment_bitmap
.as_ref()
.is_none_or(|ids| ids.contains(fragment.id as u32))
&& fragment.overlays.iter().any(|overlay| {
overlay.committed_version > meta.dataset_version
&& overlay.data_file.fields.contains(&field)
})
}) {
return plan_err!(
"column '{}' has stale index values in segment {}; rebuild the index",
config.column,
meta.uuid
);
}
let index = dataset
.open_vector_index(&config.column, &meta.uuid, &NoOpMetricsCollector)
.await?;
if !index.supports_pairwise_vectors() {
return plan_err!(
"segment {} requires a current-format pairwise vector index; rebuild the index",
meta.uuid
);
}
for partition_id in 0..index.ivf_model().num_partitions() {
tasks.push(DuplicatePairTask {
segment_id: meta.uuid,
partition_id,
});
}
}
if tasks.is_empty() {
return plan_err!("no vector index partitions found");
}
Ok(tasks)
}
/// Streaming native operator over partition descriptors, never source-vector
/// rows. Assigned descriptors are consumed sequentially; distributed hosts send
/// one descriptor per worker request. No Python or per-vector ANN calls.
/// Distributed merging deliberately advertises no global row ordering.
#[derive(Debug)]
pub struct DuplicatePairsExec {
dataset: Arc<Dataset>,
config: DuplicatePairsConfig,
tasks: Vec<DuplicatePairTask>,
properties: Arc<PlanProperties>,
metrics: ExecutionPlanMetricsSet,
}
impl DuplicatePairsExec {
pub fn try_new(
dataset: Arc<Dataset>,
config: DuplicatePairsConfig,
tasks: Vec<DuplicatePairTask>,
) -> Result<Self> {
if dataset.version().version != config.dataset_version
|| !config.distance_threshold.is_finite()
|| tasks.is_empty()
{
return plan_err!("invalid duplicate-pair snapshot, threshold, or empty task list");
}
let unique = tasks
.iter()
.map(|t| (t.segment_id, t.partition_id))
.collect::<BTreeSet<_>>();
if unique.len() != tasks.len() {
return plan_err!("duplicate partition tasks");
}
let properties = Arc::new(PlanProperties::new(
EquivalenceProperties::new(duplicate_pair_schema()),
Partitioning::UnknownPartitioning(1),
EmissionType::Incremental,
Boundedness::Bounded,
));
Ok(Self {
dataset,
config,
tasks,
properties,
metrics: ExecutionPlanMetricsSet::new(),
})
}
pub fn dataset(&self) -> &Arc<Dataset> {
&self.dataset
}
pub fn config(&self) -> &DuplicatePairsConfig {
&self.config
}
pub fn tasks(&self) -> &[DuplicatePairTask] {
&self.tasks
}
}
impl DisplayAs for DuplicatePairsExec {
fn fmt_as(&self, _: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"DuplicatePairsExec: source_version={}, column={}, threshold={}, tasks={}",
self.config.dataset_version,
self.config.column,
self.config.distance_threshold,
self.tasks.len()
)
}
}
impl ExecutionPlan for DuplicatePairsExec {
fn name(&self) -> &str {
"DuplicatePairsExec"
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.properties
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> Result<Arc<dyn ExecutionPlan>> {
if !children.is_empty() {
return plan_err!("DuplicatePairsExec has no children");
}
Ok(self)
}
fn metrics(&self) -> Option<MetricsSet> {
Some(self.metrics.clone_inner())
}
fn execute(
&self,
partition: usize,
context: Arc<TaskContext>,
) -> Result<SendableRecordBatchStream> {
if partition != 0 {
return plan_err!("duplicate-pair output partition out of bounds");
}
let tasks = self.tasks.clone();
let dataset = self.dataset.clone();
let config = self.config.clone();
let output_rows = MetricBuilder::new(&self.metrics).counter("output_rows", partition);
let completed =
MetricBuilder::new(&self.metrics).counter("completed_partitions", partition);
// Beta.14 stages native codes and scores bounded tiles concurrently.
// Keep the host staging budget explicit rather than inheriting Lance's
// larger default or multiplying CPU-pool concurrency across PE tasks.
let dimension = match dataset
.schema()
.field(&config.column)
.map(|f| f.data_type())
{
Some(DataType::FixedSizeList(_, dimension)) if dimension > 0 => dimension as usize,
_ => return plan_err!("duplicate pairs requires a fixed-size vector column"),
};
let batch_bytes = VECTOR_BATCH_BYTES.max(
MIN_VECTOR_BATCH_ROWS * (dimension * MAX_COORDINATE_BYTES + STAGED_ROW_OVERHEAD_BYTES),
);
// Allow staging/preparation overlap, at least three live spill batches,
// source input, and in-flight output plus the batch being encoded.
// Quantizer models and caches remain outside this admission estimate.
let budget = 2 * CODE_STAGING_BYTES
+ 3 * batch_bytes
+ INDEX_READ_ROWS * (dimension * MAX_COORDINATE_BYTES + STAGED_ROW_OVERHEAD_BYTES)
+ (SCORING_CONCURRENCY + 1) * ANCHOR_BLOCK_ROWS * VECTOR_BATCH_ROWS * PAIR_ROW_BYTES;
let reservation =
datafusion_execution::memory_pool::MemoryConsumer::new("DuplicatePairsExec")
.register(context.memory_pool());
reservation.try_grow(budget)?;
let reservation = Arc::new(reservation);
// Sequential local work prevents one query from reserving a kernel for
// every IVF partition at once. Distributed hosts split these descriptors
// into one-task PE requests before execution.
let stream = stream::iter(tasks)
.then(move |task| {
let dataset = dataset.clone();
let config = config.clone();
let reservation = reservation.clone();
let output_rows = output_rows.clone();
let completed = completed.clone();
async move {
let reader = lance::index::vector::dedup::find_duplicate_pairs_in_partition_with_options(
dataset,
&config.column,
task.segment_id,
task.partition_id,
config.distance_threshold,
lance::index::vector::dedup::DuplicatePairsOptions::default()
.with_memory_limit(CODE_STAGING_BYTES)
.with_max_concurrency(SCORING_CONCURRENCY),
)
.await?;
let measured = stream::try_unfold(
(reader, reservation),
move |(mut reader, reservation)| {
let output_rows = output_rows.clone();
let completed = completed.clone();
async move {
match reader.try_next().await? {
Some(batch) => {
output_rows.add(batch.num_rows());
Ok::<_, DataFusionError>(Some((
batch,
(reader, reservation),
)))
}
None => {
completed.add(1);
Ok(None)
}
}
}
},
);
Ok::<_, DataFusionError>(Box::pin(RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
measured,
)) as SendableRecordBatchStream)
}
})
.try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
stream,
)))
}
}
/// Inspectable immutable task manifest. Credentials are never included.
pub fn duplicate_pair_task_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("task_id", DataType::Utf8, false),
Field::new("dataset_uri", DataType::Utf8, false),
Field::new("dataset_version", DataType::UInt64, false),
Field::new("column", DataType::Utf8, false),
Field::new("segment_id", DataType::Utf8, false),
Field::new("partition_id", DataType::UInt64, false),
Field::new("distance_threshold", DataType::Float32, false),
Field::new("operator_version", DataType::UInt32, false),
]))
}
pub fn duplicate_pair_task_batch(
dataset: &Dataset,
config: &DuplicatePairsConfig,
tasks: &[DuplicatePairTask],
) -> Result<arrow_array::RecordBatch> {
use arrow_array::{Float32Array, StringArray, UInt32Array, UInt64Array};
use sha2::{Digest, Sha256};
let ids = tasks
.iter()
.map(|task| {
// Length-delimited strings prevent ambiguous concatenation; threshold
// bits preserve the exact float32 value sent to the native primitive.
let bytes = serde_json::to_vec(&(
DUPLICATE_PAIRS_OPERATOR_VERSION,
dataset.uri(),
config.dataset_version,
&config.column,
task.segment_id.to_string(),
task.partition_id,
config.distance_threshold.to_bits(),
))
.expect("task fields are serializable");
format!("{:x}", Sha256::digest(bytes))
})
.collect::<Vec<_>>();
let n = tasks.len();
Ok(arrow_array::RecordBatch::try_new(
duplicate_pair_task_schema(),
vec![
Arc::new(StringArray::from(ids)),
Arc::new(StringArray::from(vec![dataset.uri(); n])),
Arc::new(UInt64Array::from(vec![config.dataset_version; n])),
Arc::new(StringArray::from(vec![config.column.as_str(); n])),
Arc::new(StringArray::from(
tasks
.iter()
.map(|t| t.segment_id.to_string())
.collect::<Vec<_>>(),
)),
Arc::new(UInt64Array::from(
tasks
.iter()
.map(|t| t.partition_id as u64)
.collect::<Vec<_>>(),
)),
Arc::new(Float32Array::from(vec![config.distance_threshold; n])),
Arc::new(UInt32Array::from(vec![DUPLICATE_PAIRS_OPERATOR_VERSION; n])),
],
)?)
}
rust/lancedb/src/table/merge.rs
+259 -5
View File
@@ -1,15 +1,18 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use arrow_array::RecordBatchReader;
use arrow_schema::{DataType, Fields};
use futures::future::Either;
use futures::{FutureExt, TryFutureExt};
use lance::dataset::{
MergeInsertBuilder as LanceMergeInsertBuilder, WhenMatched, WhenNotMatchedBySource,
};
use lance_datafusion::utils::StreamingWriteSource;
use serde::{Deserialize, Serialize};
use crate::error::{Error, Result};
@@ -248,6 +251,59 @@ fn canonicalize_merge_filter(filter: Option<MergeFilter>) -> Result<Option<Merge
.transpose()
}
// The JSON projection iterates target fields, so reject unknown source fields
// before it can erase them. Missing and reordered fields remain valid merge
// inputs; type compatibility is still checked by the cast and Lance writer.
fn validate_merge_source_fields(input: &Fields, target: &Fields, parent: &str) -> Result<()> {
let mut names = HashSet::with_capacity(input.len());
for field in input {
let path = if parent.is_empty() {
field.name().clone()
} else {
format!("{parent}.{}", field.name())
};
if !names.insert(field.name()) {
return Err(Error::InvalidInput {
message: format!("merge source field '{path}' is specified more than once"),
});
}
let target_field = target
.iter()
.find(|candidate| candidate.name() == field.name())
.ok_or_else(|| Error::InvalidInput {
message: format!("merge source field '{path}' is not present in the table schema"),
})?;
validate_merge_source_type(field.data_type(), target_field.data_type(), &path)?;
}
Ok(())
}
fn validate_merge_source_type(input: &DataType, target: &DataType, path: &str) -> Result<()> {
match (input, target) {
(DataType::Struct(input), DataType::Struct(target)) => {
validate_merge_source_fields(input, target, path)
}
(
DataType::List(input)
| DataType::LargeList(input)
| DataType::FixedSizeList(input, _)
| DataType::ListView(input)
| DataType::LargeListView(input),
DataType::List(target)
| DataType::LargeList(target)
| DataType::FixedSizeList(target, _)
| DataType::ListView(target)
| DataType::LargeListView(target),
)
| (DataType::Map(input, _), DataType::Map(target, _)) => {
validate_merge_source_type(input.data_type(), target.data_type(), path)
}
(DataType::Dictionary(_, input), _) => validate_merge_source_type(input, target, path),
(_, DataType::Dictionary(_, target)) => validate_merge_source_type(input, target, path),
_ => Ok(()),
}
}
/// Internal implementation of the merge insert logic
///
/// This logic was moved from NativeTable::merge_insert to keep table.rs clean.
@@ -279,6 +335,27 @@ pub(crate) async fn execute_merge_insert(
}
let dataset = table.dataset.get().await?;
let schema = arrow_schema::Schema::from(dataset.schema());
// JSON source fields must carry the stored extension metadata just as on
// append. Keep arrow.json text labelled until Lance encodes it as JSONB.
let source = if schema
.fields()
.iter()
.any(|field| lance_arrow::json::has_json_fields(field))
{
validate_merge_source_fields(new_data.schema().fields(), schema.fields(), "")?;
let plan = Arc::new(super::datafusion::scannable_exec::ScannableExec::new(
Box::new(new_data),
None,
));
let plan = super::datafusion::cast::cast_to_table_schema(plan, &schema)?;
datafusion_physical_plan::execute_stream(
plan,
Arc::new(datafusion_execution::TaskContext::default()),
)?
} else {
new_data.into_stream()
};
let mut builder = LanceMergeInsertBuilder::try_new(dataset.clone(), params.on)?;
match (
params.when_matched_update_all,
@@ -313,10 +390,7 @@ pub(crate) async fn execute_merge_insert(
builder.use_index(params.use_index);
let future = if let Some(timeout) = params.timeout {
let future = builder
.retry_timeout(timeout)
.try_build()?
.execute_reader(new_data);
let future = builder.retry_timeout(timeout).try_build()?.execute(source);
Either::Left(tokio::time::timeout(timeout, future).map(|res| match res {
Ok(Ok((new_dataset, stats))) => Ok((new_dataset, stats)),
Ok(Err(e)) => Err(e.into()),
@@ -326,7 +400,7 @@ pub(crate) async fn execute_merge_insert(
}))
} else {
let job = builder.try_build()?;
Either::Right(job.execute_reader(new_data).map_err(|e| e.into()))
Either::Right(job.execute(source).map_err(|e| e.into()))
};
let (new_dataset, stats) = future.await?;
let version = new_dataset.manifest().version;
@@ -353,6 +427,186 @@ mod tests {
use crate::connect;
#[rstest::rstest]
#[case(None)]
#[case(Some("struct"))]
#[case(Some("list"))]
#[case(Some("large_list"))]
#[case(Some("fixed_size_list"))]
#[case(Some("map"))]
#[tokio::test]
async fn merge_json_rejects_ambiguous_source_fields(
#[case] nested: Option<&str>,
#[values(false, true)] duplicate: bool,
) {
let mut target = vec![
Field::new("id", DataType::Int32, false),
lance_arrow::json::json_field("payload", true),
];
let mut input = vec![
Field::new("id", DataType::Int32, false),
Field::new("payload", DataType::Utf8, true),
];
let mut columns: Vec<Arc<dyn arrow_array::Array>> = vec![
Arc::new(Int32Array::from(vec![1])),
Arc::new(StringArray::from(vec![r#"{"x":1}"#])),
];
if let Some(nested) = nested {
let wrap = |children: Vec<Field>| {
let structure = DataType::Struct(children.into());
let item = Arc::new(Field::new("item", structure.clone(), true));
match nested {
"struct" => structure,
"list" => DataType::List(item),
"large_list" => DataType::LargeList(item),
"fixed_size_list" => DataType::FixedSizeList(item, 2),
"map" => DataType::Map(
Arc::new(Field::new(
"entries",
DataType::Struct(
vec![
Field::new("key", DataType::Utf8, false),
Field::new("value", structure, true),
]
.into(),
),
false,
)),
false,
),
_ => unreachable!(),
}
};
let known = Field::new("value", DataType::Int32, true);
let unexpected = if duplicate { "value" } else { "extra" };
let target_type = wrap(vec![known.clone()]);
let input_type = wrap(vec![known, Field::new(unexpected, DataType::Int32, true)]);
columns.push(arrow_array::new_null_array(&input_type, 1));
target.push(Field::new("details", target_type, true));
input.push(Field::new("details", input_type, true));
} else {
let unexpected = if duplicate { "id" } else { "extra" };
input.push(Field::new(unexpected, DataType::Int32, true));
columns.push(Arc::new(Int32Array::from(vec![99])));
}
let db = connect("memory://").execute().await.unwrap();
let table = db
.create_empty_table("ambiguous_json", Arc::new(Schema::new(target)))
.execute()
.await
.unwrap();
let version = table.version().await.unwrap();
let schema = Arc::new(Schema::new(input));
let batch = RecordBatch::try_new(schema.clone(), columns).unwrap();
let mut merge = table.merge_insert(&["id"]);
merge.when_not_matched_insert_all();
let error = merge
.execute(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)))
.await
.expect_err("JSON alignment must not drop source fields");
let expected = if duplicate {
"more than once"
} else {
"not present"
};
assert!(error.to_string().contains(expected), "{error}");
assert_eq!(table.version().await.unwrap(), version);
assert_eq!(table.count_rows(None).await.unwrap(), 0);
}
#[tokio::test]
async fn merge_json_preserves_empty_extension_metadata() {
use crate::query::{ExecutableQuery, QueryBase, Select};
use futures::TryStreamExt;
let mut stored_json = lance_arrow::json::json_field("payload", true);
let mut metadata = stored_json.metadata().clone();
metadata.insert(lance_arrow::ARROW_EXT_META_KEY.into(), String::new());
stored_json = stored_json.with_metadata(metadata);
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
stored_json,
]));
let db = connect("memory://").execute().await.unwrap();
let table = db
.create_empty_table("json_merge_metadata", schema)
.execute()
.await
.unwrap();
let input_schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("payload", DataType::Utf8, true).with_metadata(
std::collections::HashMap::from([(
lance_arrow::ARROW_EXT_NAME_KEY.into(),
lance_arrow::json::ARROW_JSON_EXT_NAME.into(),
)]),
),
]));
let seed = RecordBatch::try_new(
input_schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1])),
Arc::new(StringArray::from(vec![r#"{"old":true}"#])),
],
)
.unwrap();
table.add(seed).execute().await.unwrap();
let batch = RecordBatch::try_new(
input_schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec![
r#"{"updated":true}"#,
r#"{"inserted":true}"#,
])),
],
)
.unwrap();
let mut merge = table.merge_insert(&["id"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
let result = merge
.execute(Box::new(RecordBatchIterator::new(
vec![Ok(batch)],
input_schema,
)))
.await
.unwrap();
assert_eq!(result.num_updated_rows, 1);
assert_eq!(result.num_inserted_rows, 1);
let output = table
.query()
.select(Select::columns(&["payload"]))
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
let values = output
.iter()
.flat_map(|b| {
b.column(0)
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.iter()
})
.flatten()
.collect::<Vec<_>>();
assert_eq!(values.len(), 2);
assert!(
values
.iter()
.any(|v| serde_json::from_str::<serde_json::Value>(v).unwrap()["updated"] == true)
);
assert!(
values
.iter()
.any(|v| serde_json::from_str::<serde_json::Value>(v).unwrap()["inserted"] == true)
);
}
fn merge_insert_test_batches(offset: i32, age: i32) -> Box<dyn RecordBatchReader + Send> {
let schema = Arc::new(Schema::new(vec![
Field::new("i", DataType::Int32, false),
rust/lancedb/tests/duplicate_pairs.rs
+561
View File
@@ -0,0 +1,561 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
use arrow_array::{
Array, FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
UInt64Array,
};
use arrow_schema::{DataType, Field, Schema};
use datafusion::prelude::SessionContext;
use datafusion_catalog::TableProvider;
use datafusion_common::Result;
use futures::TryStreamExt;
use lance::{
Dataset,
dataset::WriteParams,
index::{DatasetIndexExt, vector::VectorIndexParams},
};
use lance_index::{IndexType, vector::ivf::IvfBuildParams};
use lance_linalg::distance::MetricType;
use lancedb::table::datafusion::{
BaseTableAdapter,
udtf::duplicate_pairs::{
DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsResolver,
DuplicatePairsTableFunction, duplicate_pair_task_batch, plan_duplicate_pairs,
},
};
use std::sync::Arc;
#[derive(Debug)]
struct Resolver(Arc<BaseTableAdapter>);
impl DuplicatePairsResolver for Resolver {
fn resolve(
&self,
table: &str,
config: DuplicatePairsConfig,
tasks_only: bool,
) -> Result<Arc<dyn TableProvider>> {
assert_eq!(table, "source");
Ok(Arc::new(self.0.duplicate_pairs(config, tasks_only)))
}
}
fn pairs(batches: &[RecordBatch]) -> Vec<(u64, u64, u32)> {
let mut result = Vec::new();
for batch in batches {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
let d = batch
.column(2)
.as_any()
.downcast_ref::<Float32Array>()
.unwrap();
for i in 0..batch.num_rows() {
result.push((a.value(i), b.value(i), d.value(i).to_bits()));
}
}
result.sort_unstable();
result
}
#[tokio::test]
async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result<()> {
for partitions in [1, 8] {
let dir = tempfile::tempdir()?;
let uri = dir
.path()
.join("source.lance")
.to_str()
.unwrap()
.to_string();
let n = 1025;
let vectors = (0..n * 2)
.flat_map(|i| {
let i = i % n;
let x = if i == n - 1 {
0.0
} else if i < 3 {
i as f32
} else {
i as f32 * 10.0
};
[x, 0.0]
})
.collect::<Vec<_>>();
let vectors = FixedSizeListArray::try_new(
Arc::new(Field::new("item", DataType::Float32, true)),
2,
Arc::new(Float32Array::from(vectors)),
None,
)?;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("vector", vectors.data_type().clone(), false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from_iter_values(0..(n * 2) as i64)),
Arc::new(vectors),
],
)?;
let mut ds = Dataset::write(
RecordBatchIterator::new(vec![Ok(batch)], schema),
&uri,
Some(WriteParams {
max_rows_per_file: n,
..Default::default()
}),
)
.await?;
let fragments = ds
.fragments()
.iter()
.map(|f| f.id as u32)
.collect::<Vec<_>>();
assert_eq!(fragments.len(), 2);
let mut segments = Vec::new();
for (i, fragment) in fragments.into_iter().enumerate() {
let centroids = FixedSizeListArray::try_new(
Arc::new(Field::new("item", DataType::Float32, true)),
2,
Arc::new(Float32Array::from(
(0..partitions)
.flat_map(|p| [p as f32 * 1400.0 + i as f32 * 0.125, 0.0])
.collect::<Vec<_>>(),
)),
None,
)?;
let ivf = IvfBuildParams {
centroids: Some(Arc::new(centroids)),
..IvfBuildParams::new(partitions)
};
let params = VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf);
segments.push(
ds.create_index_builder(&["vector"], IndexType::Vector, &params)
.name("vector_idx".into())
.fragments(vec![fragment])
.execute_uncommitted()
.await?,
);
}
ds.commit_existing_index_segments("vector_idx", "vector", segments)
.await?;
let version = ds.version().version;
let ds = Arc::new(ds);
let conn = lancedb::connect(dir.path().to_str().unwrap())
.execute()
.await?;
let table = conn.open_table("source").execute().await?;
let resolver = Arc::new(Resolver(Arc::new(
BaseTableAdapter::try_new(table.base_table().clone()).await?,
)));
// More partitions than this pool could hold concurrently: local
// execution must consume scoped readers sequentially.
let runtime = Arc::new(
datafusion_execution::runtime_env::RuntimeEnvBuilder::new()
.with_memory_limit(128 * 1024 * 1024, 1.0)
.build()?,
);
let ctx =
SessionContext::new_with_config_rt(datafusion::prelude::SessionConfig::new(), runtime);
ctx.register_udtf(
"vector_duplicate_pairs",
Arc::new(DuplicatePairsTableFunction::new(resolver.clone())),
);
ctx.register_udtf(
"vector_duplicate_pair_tasks",
Arc::new(DuplicatePairsTableFunction::task_manifest(resolver)),
);
for threshold in [-1.0, 0.0, 1.0] {
let config = DuplicatePairsConfig {
dataset_version: version,
column: "vector".into(),
distance_threshold: threshold,
};
let tasks = plan_duplicate_pairs(ds.clone(), &config).await?;
assert_eq!(tasks.len(), partitions * 2);
let manifest = duplicate_pair_task_batch(&ds, &config, &tasks)?;
let sql_manifest = ctx.sql(&format!("SELECT * FROM vector_duplicate_pair_tasks('source', {version}, 'vector', {threshold})")).await?.collect().await?;
assert_eq!(sql_manifest, vec![manifest]);
let expected =
lance::index::vector::dedup::find_duplicate_pairs(ds.clone(), "vector", threshold)
.await?
.try_collect::<Vec<_>>()
.await?;
let sql = format!(
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'vector', {threshold})"
);
let actual = ctx.sql(&sql).await?.collect().await?;
assert_eq!(pairs(&actual), pairs(&expected));
assert_eq!(
pairs(&actual).len(),
match threshold {
-1.0 => 0,
0.0 => 2,
_ => 8,
}
);
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)))
.collect::<Vec<_>>();
edges.sort_unstable();
let mut removed = std::collections::BTreeSet::new();
for (a, b) in edges {
if !removed.contains(&a) {
removed.insert(b);
}
}
let source_ids = ds
.scan()
.project(&["id"])?
.with_row_id()
.try_into_stream()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut expected_ids = Vec::new();
for batch in source_ids {
let ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let rowids = batch
.column_by_name("_rowid")
.unwrap()
.as_any()
.downcast_ref::<UInt64Array>()
.unwrap();
for i in 0..batch.num_rows() {
if !removed.contains(&rowids.value(i)) {
expected_ids.push(ids.value(i));
}
}
}
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();
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");
}
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(),
config.clone(),
vec![tasks[0].clone(), tasks[0].clone()]
)
.is_err()
);
}
// Latest no longer contains B, but pinned queries must retain the same pairs.
table.delete("id = 1").await?;
let actual = ctx
.sql(&format!(
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'vector', 1.0)"
))
.await?
.collect()
.await?;
assert_eq!(pairs(&actual).len(), 8);
assert!(
ctx.sql(&format!(
"SELECT * FROM vector_duplicate_pairs('source', {version}, 'missing', 1.0)"
))
.await?
.collect()
.await
.is_err()
);
}
Ok(())
}
#[tokio::test]
async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() -> 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())
.execute()
.await?;
let vectors = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
[
Some(vec![Some(0.0), Some(0.0)]),
Some(vec![Some(1.0), Some(0.0)]),
Some(vec![Some(2.0), Some(0.0)]),
None,
Some(vec![Some(30.0), Some(0.0)]),
Some(vec![Some(40.0), Some(0.0)]),
],
2,
);
let blob = lancedb::blob::blob("image", true);
let DataType::Struct(fields) = blob.data_type().clone() else {
unreachable!()
};
let images = StructArray::new(
fields,
vec![
Arc::new(LargeBinaryArray::from_iter_values(
(0..6).map(|id| vec![id; 4096]),
)),
Arc::new(StringArray::from(vec![None::<&str>; 6])),
],
None,
);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("vector", vectors.data_type().clone(), true),
blob,
])),
vec![
Arc::new(Int64Array::from_iter_values(0..6)),
Arc::new(vectors),
Arc::new(images),
],
)?;
let table = conn
.create_table("source", batch)
.storage_option("new_table_enable_stable_row_ids", stable.to_string())
.execute()
.await?;
table.delete("id = 5").await?;
table
.optimize(lancedb::table::OptimizeAction::Compact {
options: lancedb::table::CompactionOptions::default(),
remap_options: None,
})
.await?;
let uri = table.uri().await?;
let mut ds = Dataset::open(&uri).await?;
assert!(
ds.fragments().iter().all(|f| f.id > 0),
"compaction must change physical row addresses"
);
let centroids = FixedSizeListArray::from_iter_primitive::<Float32Type, _, _>(
[Some(vec![Some(0.0), Some(0.0)])],
2,
);
let ivf = IvfBuildParams {
centroids: Some(Arc::new(centroids)),
..IvfBuildParams::new(1)
};
ds.create_index(
&["vector"],
IndexType::Vector,
Some("vector_idx".into()),
&VectorIndexParams::with_ivf_flat_params(MetricType::L2, ivf),
true,
)
.await?;
let version = ds.version().version;
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?;
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(())
}