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

Merge main and align vendored encoder with Lance beta.15

Browse Source
This commit is contained in:
Gatefixer
2026-09-28 11:01:12 +00:00
parent 617ee19b68 0327618a62
commit 311306fe4e
31 changed files with 3522 additions and 136 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
+45 -43
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,7 +5372,7 @@ dependencies = [
[[package]]
name = "lance-encoding"
version = "13.0.0-beta.14"
version = "13.0.0-beta.15"
dependencies = [
"arrow-arith",
"arrow-array",
@@ -5404,8 +5405,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",
@@ -5437,8 +5438,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",
@@ -5452,8 +5453,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",
@@ -5523,8 +5524,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",
@@ -5542,12 +5543,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",
@@ -5589,8 +5591,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",
@@ -5604,8 +5606,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",
@@ -5619,8 +5621,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",
@@ -5673,8 +5675,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",
@@ -5688,8 +5690,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",
@@ -5729,8 +5731,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",
@@ -5743,8 +5745,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
@@ -14,20 +14,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 = { path = "vendor/lance-encoding", version = "=13.0.0-beta.14" }
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 = { path = "vendor/lance-encoding", version = "=13.0.0-beta.15" }
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/lancedb/rerankers/cross_encoder.py
+12 -1
View File
@@ -31,6 +31,12 @@ class CrossEncoderReranker(Reranker):
trust_remote_code : bool, default True
If True, will trust the remote code to be safe. If False, will not trust
the remote code and will not run it
batch_size : int, default 32
The number of query-passage pairs scored per forward pass.
**kwargs
Additional keyword arguments passed to
``sentence_transformers.CrossEncoder``, for example ``max_length`` or
``model_kwargs={"torch_dtype": "bfloat16"}``.
"""
def __init__(
@@ -40,6 +46,8 @@ class CrossEncoderReranker(Reranker):
device: Union[str, None] = None,
return_score="relevance",
trust_remote_code: bool = True,
batch_size: int = 32,
**kwargs,
):
super().__init__(return_score)
torch = attempt_import_or_raise("torch")
@@ -47,6 +55,8 @@ class CrossEncoderReranker(Reranker):
self.column = column
self.device = device
self.trust_remote_code = trust_remote_code
self.batch_size = batch_size
self.kwargs = kwargs
if self.device is None:
self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -61,6 +71,7 @@ class CrossEncoderReranker(Reranker):
self.model_name,
device=self.device,
trust_remote_code=self.trust_remote_code,
**self.kwargs,
)
return cross_encoder
@@ -71,7 +82,7 @@ class CrossEncoderReranker(Reranker):
return result_set
passages = result_set[self.column].to_pylist()
cross_inp = [[query, passage] for passage in passages]
cross_scores = self.model.predict(cross_inp)
cross_scores = self.model.predict(cross_inp, batch_size=self.batch_size)
result_set = result_set.append_column(
"_relevance_score", pa.array(cross_scores, type=pa.float32())
)
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/python/tests/test_rerankers.py
+8
View File
@@ -481,6 +481,14 @@ def test_cross_encoder_reranker(tmp_path):
_run_test_reranker(reranker, table, "single player experience", None, schema)
def test_cross_encoder_reranker_kwargs(tmp_path):
pytest.importorskip("sentence_transformers")
reranker = CrossEncoderReranker(batch_size=4, max_length=64)
assert reranker.model.max_length == 64
table, schema = get_test_table(tmp_path)
_run_test_reranker(reranker, table, "single player experience", None, schema)
def test_colbert_reranker(tmp_path):
pytest.importorskip("rerankers")
reranker = ColbertReranker()
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(())
}
vendor/lance-encoding/Cargo.toml
Vendored
+5 -5
View File
@@ -1,6 +1,6 @@
[package]
name = "lance-encoding"
version = "13.0.0-beta.14"
version = "13.0.0-beta.15"
edition = "2024"
license = "Apache-2.0"
repository = "https://github.com/lance-format/lance"
@@ -13,10 +13,10 @@ rust-version = "1.91.0"
[workspace]
[dependencies]
lance-arrow = { version = "=13.0.0-beta.14", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.14" }
lance-core = { version = "=13.0.0-beta.14", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.14" }
lance-bitpacking = { version = "=13.0.0-beta.14", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.14", optional = true }
fsst = { version = "=13.0.0-beta.14", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.14" }
lance-arrow = { version = "=13.0.0-beta.15", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.15" }
lance-core = { version = "=13.0.0-beta.15", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.15" }
lance-bitpacking = { version = "=13.0.0-beta.15", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.15", optional = true }
fsst = { version = "=13.0.0-beta.15", git = "https://github.com/lance-format/lance.git", tag = "v13.0.0-beta.15" }
arrow-arith = "58.0.0"
arrow-array = "58.0.0"
arrow-buffer = "58.0.0"
vendor/lance-encoding/PATCH.md
Vendored
+1 -1
View File
@@ -1,5 +1,5 @@
This is `rust/lance-encoding` from
[`lance-format/lance` v13.0.0-beta.14](https://github.com/lance-format/lance/tree/b80ac1d9f3c24b47bd8e297d03053ea35d216a9f/rust/lance-encoding),
[`lance-format/lance` v13.0.0-beta.15](https://github.com/lance-format/lance/tree/db53dcf367088b32ae3a7461407d43593759ae3e/rust/lance-encoding),
with the crate's build protos copied into `protos/`. Its manifest uses explicit
dependencies so it can be built outside the Lance workspace.