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

refactor: always fully refresh materialized views (#4373)

Browse Source 
Simplify experimental materialized views to one output-only storage
layout and one full-rebuild refresh implementation.

## Current layout

Local materialized views currently maintain incremental-refresh state
alongside the materialized values:

- an internal `__source_row_id` column for source provenance;
- `mv.definition` plus `mv.incarnation`;
- source version/timestamp, view version, and last-refresh watermarks;
- optional `mv.staging` state for table-function sources; and
- stable-row-ID requirements on the source table.

Refresh classifies source changes, appends or evicts rows when possible,
and falls back to a rebuild for unsupported deltas. Remote
create/refresh also use separate REST mutation paths.

## New layout

An MV is an ordinary Lance backing table containing only:

- the defining query's output columns;
- its current data fragments; and
- `mv.definition`, encoded as a versioned SQL envelope such as
`{"kind":"query","format":1,"query":"SELECT ..."}`.

There is no provenance column, incarnation, source/view watermark,
staging metadata, or stable-row-ID requirement. Local refresh may still
pin a source version for reproducibility, but it always recomputes the
complete result.

Full refresh streams the new result into uncommitted fragments and then
publishes a single transaction that removes every old fragment and
installs the new set. A shared conflict token prevents concurrent
replacements from both committing. This fragment-replacement primitive
is also used by Sophon's SQL worker.

## New remote create workflow

```text
Connection::create_materialized_view(...)
        -> builder renders the defining SELECT
        -> remote Database receives a CreateMaterializedViewRequest
        -> submit CREATE MATERIALIZED VIEW ... AS ... through Flight SQL
        -> retain the SQL Query handle while the Job is running
        -> consume the durable SQL-job result
        -> open the created table and read mv.definition
        -> return MaterializedView
```

Remote refresh similarly submits `REFRESH MATERIALIZED VIEW` through
Flight SQL. Cancelling either wrapper forwards cancellation to the
retained SQL query before aborting the local task, so the server
operation is not silently left running.

## API changes

- remove incremental/no-op modes and the `full` switch;
- remove distributed grouped/partitioned refresh APIs and native
vector-source MV APIs;
- keep full-refresh results across Rust, Python, and TypeScript;
- preserve list/describe/drop as object-management operations while
routing create/refresh mutations through SQL.

The storage contract is experimental and has no compatibility shim;
views using the removed layout must be recreated.
This commit is contained in:
Jack Ye
2026-10-05 21:33:34 -07:00
committed by GitHub
parent d864e7198f
commit 7e8f0d4982
33 changed files with 1030 additions and 11498 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/js/classes/Connection.md
-4
View File
@@ -183,10 +183,6 @@ Define a materialized view named `name` over the table `source`.
The view is populated before creation returns. Set `withNoData` to create
only its definition and empty backing table. The view is a normal table:
it can be queried, indexed and searched, and it appears in `tableNames`.
The source table must have stable row ids (create it with
the `newTableEnableStableRowIds` storage option); they keep the view's
provenance valid across source compactions and cannot be enabled after
a table exists.
#### Parameters
docs/src/js/classes/MaterializedView.md
+3 -10
View File
@@ -65,21 +65,14 @@ refresh(options?): Promise<RefreshMaterializedViewResult>
Recompute the view from its source.
The refresh is incremental when the source's changes can be reconciled
into the view -- rows added, changed or removed since the last one --
and otherwise rebuilds. `full` forces a rebuild; `sourceVersion`
refreshes to that source version instead of the latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser throws
rather than writing them a second time.
Every refresh rebuilds the complete view. For local views,
`sourceVersion` refreshes to that source version instead of the latest;
remote SQL refreshes do not support source-version pinning.
#### Parameters
* **options?**
* **options.full?**: `boolean`
* **options.sourceVersion?**: `number`
#### Returns
docs/src/js/interfaces/RefreshMaterializedViewResult.md
+1 -1
View File
@@ -14,7 +14,7 @@
mode: string;
```
How the view was brought up to date: "rebuild", "incremental" or "no_op".
How the view was brought up to date: "rebuild".
***
docs/src/python/python.md
-27
View File
@@ -199,33 +199,6 @@ listing a storage directory.
::: lancedb.materialized_view.MaterializedViewDefinition
::: lancedb.vector_dedup
::: lancedb.materialized_view.VectorDedupSource
For an existing indexed source, declare and materialize a dedup result without
writing SQL. Creation captures the source version once and waits for the complete
result; remote connections use the service's MV jobs.
```python
from lancedb import vector_dedup
view = db.create_materialized_view(
"images_clean",
vector_dedup("images", column="phash", distance_threshold=4),
)
cleaned = view.table
```
To submit without waiting, call `db.create_materialized_view_async(...)` with the
same source expression and wait on the returned job. To declare only, pass
`with_no_data=True`, then use `view.refresh_async().wait()` when ready. Pass
`dataset_version=...` to `vector_dedup` for an explicit snapshot. The default
policy retains direct representatives (A-B-C without A-C keeps A and C), preserves
the source, and materializes its original columns. Only within-segment/partition
pairs are covered, using the index's distance representation.
## Views
::: lancedb.view.ViewDescription
nodejs/__test__/materialized_view.test.ts
+14 -31
View File
@@ -16,15 +16,11 @@ describe("materialized views", () => {
beforeEach(async () => {
tmpDir = tmp.dirSync({ unsafeCleanup: true });
db = await connect(tmpDir.name);
await db.createTable(
"people",
[
{ name: "ada", age: 36 },
{ name: "kid", age: 7 },
{ name: "grace", age: 85 },
],
{ storageOptions: { newTableEnableStableRowIds: "true" } },
);
await db.createTable("people", [
{ name: "ada", age: 36 },
{ name: "kid", age: 7 },
{ name: "grace", age: 85 },
]);
});
afterEach(() => tmpDir.removeCallback());
@@ -52,19 +48,9 @@ describe("materialized views", () => {
"SELECT * FROM people",
);
for (const [format, fn] of [
[3, "vector_duplicate_pairs"],
[4, "vector_dedup"],
]) {
const native = `SELECT * FROM ${fn}('images', 3, 'phash', 4)`;
expect(read(JSON.stringify({ format, query: native })).query).toBe(
native,
);
}
// A newer writer's layout is reported, never guessed at.
for (const newer of [
`{"format":5,"query":${JSON.stringify(query)}}`,
`{"format":3,"query":${JSON.stringify(query)}}`,
'{"kind":"select_v3","source_table":"people"}',
]) {
expect(() => read(newer)).toThrow(/cannot refresh/);
@@ -106,7 +92,7 @@ describe("materialized views", () => {
);
});
it("refreshes incrementally after an append", async () => {
it("fully refreshes after an append", async () => {
const view = await db.createMaterializedView("copy", "people", {
withNoData: true,
});
@@ -115,11 +101,11 @@ describe("materialized views", () => {
const people = await db.openTable("people");
await people.add([{ name: "alan", age: 41 }]);
const result = await view.refresh();
expect(result.mode).toBe("incremental");
expect(Number(result.rowsWritten)).toBe(1);
expect(result.mode).toBe("rebuild");
expect(Number(result.rowsWritten)).toBe(4);
expect(await view.table().countRows()).toBe(4);
expect((await view.refresh()).mode).toBe("no_op");
expect((await view.refresh()).mode).toBe("rebuild");
});
it("lists views and rejects non-views", async () => {
@@ -172,9 +158,7 @@ describe("materialized views", () => {
});
it("quotes bare select names", async () => {
await db.createTable("odd_names", [{ "order item": "widget" }], {
storageOptions: { newTableEnableStableRowIds: "true" },
});
await db.createTable("odd_names", [{ "order item": "widget" }]);
const view = await db.createMaterializedView("quoted", "odd_names", {
select: ["order item"],
withNoData: true,
@@ -183,10 +167,9 @@ describe("materialized views", () => {
expect(Number(result.rowsWritten)).toBe(1);
});
it("requires stable row ids on the source", async () => {
it("does not require stable row ids on the source", async () => {
await db.createTable("plain", [{ x: 1 }]);
await expect(db.createMaterializedView("v", "plain")).rejects.toThrow(
"stable row ids",
);
const view = await db.createMaterializedView("v", "plain");
expect(await view.table().countRows()).toBe(1);
});
});
nodejs/lancedb/connection.ts
-4
View File
@@ -325,10 +325,6 @@ export abstract class Connection {
* The view is populated before creation returns. Set `withNoData` to create
* only its definition and empty backing table. The view is a normal table:
* it can be queried, indexed and searched, and it appears in `tableNames`.
* The source table must have stable row ids (create it with
* the `newTableEnableStableRowIds` storage option); they keep the view's
* provenance valid across source compactions and cannot be enabled after
* a table exists.
*/
abstract createMaterializedView(
name: string,
nodejs/lancedb/materialized_view.ts
+6 -16
View File
@@ -9,10 +9,9 @@ export const DEFINITION_META_KEY = "mv.definition";
/**
* The newest stored layout this version reads: `{"format": N, "query": "<SQL>"}`,
* format 2 being a query with `GROUP BY`, format 3 native pairs and format 4
* native dedup results.
* format 2 being a query with `GROUP BY`.
*/
export const DEFINITION_FORMAT = 4;
export const DEFINITION_FORMAT = 2;
/**
* The query that defines a materialized view, as stored:
@@ -188,23 +187,14 @@ export class MaterializedView {
/**
* Recompute the view from its source.
*
* The refresh is incremental when the source's changes can be reconciled
* into the view -- rows added, changed or removed since the last one --
* and otherwise rebuilds. `full` forces a rebuild; `sourceVersion`
* refreshes to that source version instead of the latest.
*
* Concurrent refreshes of one view do not duplicate its rows. Two that
* plan the same source rows conflict on commit, and the loser throws
* rather than writing them a second time.
* Every refresh rebuilds the complete view. For local views,
* `sourceVersion` refreshes to that source version instead of the latest;
* remote SQL refreshes do not support source-version pinning.
*/
async refresh(options?: {
full?: boolean;
sourceVersion?: number;
}): Promise<RefreshMaterializedViewResult> {
validateNonNegativeInteger(options?.sourceVersion, "sourceVersion");
return await this.inner.refreshMaterializedView(
options?.full,
options?.sourceVersion,
);
return await this.inner.refreshMaterializedView(options?.sourceVersion);
}
}
nodejs/lancedb/table.ts
+1 -3
View File
@@ -685,7 +685,6 @@ export abstract class Table {
* @ignore
*/
abstract refreshMaterializedView(
full?: boolean,
sourceVersion?: number,
): Promise<RefreshMaterializedViewResult>;
@@ -1420,10 +1419,9 @@ export class LocalTable extends Table {
}
async refreshMaterializedView(
full?: boolean,
sourceVersion?: number,
): Promise<RefreshMaterializedViewResult> {
return await this.inner.refreshMaterializedView(full, sourceVersion);
return await this.inner.refreshMaterializedView(sourceVersion);
}
async materializedViewDefinition(): Promise<string> {
nodejs/src/table.rs
+5 -10
View File
@@ -469,13 +469,12 @@ impl Table {
#[napi(catch_unwind)]
pub async fn refresh_materialized_view(
&self,
full: Option<bool>,
source_version: Option<i64>,
) -> napi::Result<RefreshMaterializedViewResult> {
let view = lancedb::MaterializedView::from_table(self.inner_ref()?.clone())
.await
.default_error()?;
let mut builder = view.refresh().full(full.unwrap_or(false));
let mut builder = view.refresh();
if let Some(version) = source_version {
let version = u64::try_from(version).map_err(|_| {
napi::Error::from_reason("sourceVersion must be a non-negative integer")
@@ -492,11 +491,9 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.default_error()?;
view.definition().to_json().map_err(|err| {
napi::Error::from_reason(format!(
"failed to serialize materialized-view definition: {err}"
))
})
Ok(lancedb::materialized_view::definition_metadata_from_sql(
view.definition_sql(),
))
}
#[napi(catch_unwind)]
@@ -1608,7 +1605,7 @@ impl From<lancedb::function::FunctionErrors> for FunctionErrors {
#[napi(object)]
pub struct RefreshMaterializedViewResult {
/// How the view was brought up to date: "rebuild", "incremental" or "no_op".
/// How the view was brought up to date: "rebuild".
pub mode: String,
pub rows_written: i64,
pub source_version: i64,
@@ -1619,8 +1616,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
fn from(value: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match value.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
python/python/lancedb/__init__.py
-1
View File
@@ -46,7 +46,6 @@ from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewDefinition,
vector_dedup as vector_dedup,
)
from .view import ViewDescription as ViewDescription
from .table import AsyncTable, Table
python/python/lancedb/_lancedb.pyi
+2 -6
View File
@@ -255,8 +255,6 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
*,
vector_source_json: Optional[str] = None,
) -> Table: ...
async def create_materialized_view_async(
self,
@@ -266,8 +264,6 @@ class Connection(object):
filter: Optional[str] = None,
limit: Optional[int] = None,
with_no_data: bool = False,
*,
vector_source_json: Optional[str] = None,
) -> Job: ...
async def list_materialized_views(self) -> List[str]: ...
async def drop_materialized_view(
@@ -500,10 +496,10 @@ class Table:
limit: Optional[int] = None,
) -> FunctionErrors: ...
async def refresh_materialized_view(
self, full: bool = False, source_version: Optional[int] = None
self, source_version: Optional[int] = None
) -> RefreshMaterializedViewResult: ...
async def refresh_materialized_view_async(
self, full: bool = False, source_version: Optional[int] = None
self, source_version: Optional[int] = None
) -> Job: ...
async def materialized_view_definition(self) -> str: ...
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
python/python/lancedb/db.py
+4 -29
View File
@@ -57,7 +57,6 @@ from .materialized_view import (
AsyncMaterializedView,
MaterializedView,
MaterializedViewSource,
VectorDedupSource,
SelectArg,
normalize_select,
)
@@ -613,21 +612,12 @@ class DBConnection(EnforceOverrides):
table. The view is a normal table: it can be queried, indexed and
searched, and it appears in ``table_names``.
Ordinary source tables must have stable row ids (create them with the
``new_table_enable_stable_row_ids`` storage option): they keep the
view's provenance valid across source compactions, and cannot be
enabled after a table exists. Native dedup sources also support physical
row IDs interpreted against their pinned snapshot.
Parameters
----------
name: str
The name of the view.
source: str or VectorDedupSource
The name of the source table, or an indexed dedup source declared
with [vector_dedup][lancedb.vector_dedup]. Dedup sources capture a
fixed snapshot and preserve the source table. They accept no
additional select, where or limit options.
source: str
The name of the source table.
select: list or dict, optional
The view's columns: column names, ``(alias, SQL expression)``
pairs, or a dict of the same. Omitting it selects every source
@@ -1654,10 +1644,7 @@ class LanceDBConnection(DBConnection):
Examples
--------
>>> import lancedb
>>> db = lancedb.connect(
... "./.lancedb",
... storage_options={"new_table_enable_stable_row_ids": "true"},
... )
>>> db = lancedb.connect("./.lancedb")
>>> data = [{"name": "ada", "age": 36}, {"name": "kid", "age": 7}]
>>> table = db.create_table("people", data)
>>> view = db.create_materialized_view(
@@ -2625,19 +2612,7 @@ class AsyncConnection(object):
where: Optional[str],
limit: Optional[int],
) -> Tuple[str, Dict[str, str]]:
if not isinstance(source, VectorDedupSource):
return source, {}
if select is not None or where is not None or limit is not None:
raise ValueError(
"vector_dedup cannot be combined with select, where or limit"
)
version = source.dataset_version
if version is None:
table = await self.open_table(source.source)
version = await table.version()
return source.source, {
"vector_source_json": source._native_source_json(version)
}
return source, {}
async def create_materialized_view(
self,
python/python/lancedb/materialized_view.py
+18 -122
View File
@@ -6,11 +6,11 @@ maintained by refresh. See ``DBConnection.create_materialized_view``."""
from __future__ import annotations
import json
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
import json
from dataclasses import dataclass
from .background_loop import LOOP
from .job import AsyncJob, Job, _typed_job
@@ -21,7 +21,6 @@ if TYPE_CHECKING:
from .table import AsyncTable, LanceTable
DEFINITION_META_KEY = b"mv.definition"
_MAX_FLOAT32 = float.fromhex("0x1.fffffep+127")
SelectArg = Union[
str,
@@ -31,104 +30,12 @@ SelectArg = Union[
]
@dataclass(frozen=True)
class VectorDedupSource:
"""A declarative source of retained original rows from an indexed snapshot.
Construct with [vector_dedup][lancedb.vector_dedup] and pass it as the
``source`` of ``create_materialized_view`` or
``create_materialized_view_async``. An omitted ``dataset_version`` is
captured once when creation is submitted, not when this object is made.
Refreshes keep that captured version. The source table is never modified.
"""
source: str
column: str
distance_threshold: float
dataset_version: Optional[int] = None
def __post_init__(self):
for name in ("source", "column"):
value = getattr(self, name)
if not isinstance(value, str) or not value:
raise ValueError(f"{name} must be a non-empty string")
if self.dataset_version is not None and (
type(self.dataset_version) is not int
or not 0 < self.dataset_version < 2**64
):
raise ValueError("dataset_version must be a positive uint64 integer")
value = self.distance_threshold
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError("distance_threshold must be a number")
try:
value = float(value)
except OverflowError as exc:
raise ValueError("distance_threshold must fit a finite float32") from exc
if not math.isfinite(value) or abs(value) > _MAX_FLOAT32:
raise ValueError("distance_threshold must fit a finite float32")
object.__setattr__(self, "distance_threshold", value)
def _native_source_json(self, version: int) -> str:
return json.dumps(
{
"kind": "Dedup",
"dataset_version": version,
"column": self.column,
"distance_threshold": repr(self.distance_threshold),
},
allow_nan=False,
)
MaterializedViewSource = str
def vector_dedup(
source: str,
*,
column: str,
distance_threshold: float,
dataset_version: Optional[int] = None,
) -> VectorDedupSource:
"""Declare an indexed dedup source for a materialized view.
The default policy keeps the smallest unassigned snapshot row ID and removes
only its direct qualifying neighbors: an A-B-C chain without an A-C edge
retains A and C. Rows with no pair, including null vectors, survive.
Only same-segment/same-IVF-partition pairs are considered; quantized distances
use reconstructed index vectors. The index must cover the selected snapshot.
``source`` is an exact table name in the connection's root namespace, just
as for ``create_materialized_view``. ``dataset_version=None`` captures the
source's version when creation is submitted. Pass an explicit positive
version to reproduce a previous snapshot. The result contains original
source columns and rows. Do not combine this source with ``select``,
``where`` or ``limit``; query the resulting table for further transforms.
Creation follows the usual MV contract: by default it waits for the result.
Use ``with_no_data=True`` to declare only, or
``create_materialized_view_async`` to get a job handle. On Cloud/Enterprise,
execution uses server-side MV registry jobs; local connections run in process.
Examples
--------
Given a connected ``db`` and an indexed ``images.phash`` column::
from lancedb import vector_dedup
view = db.create_materialized_view(
"images_clean",
vector_dedup("images", column="phash", distance_threshold=4),
)
cleaned = view.table
"""
return VectorDedupSource(source, column, distance_threshold, dataset_version)
MaterializedViewSource = Union[str, VectorDedupSource]
DEFINITION_FORMAT = 4
DEFINITION_FORMAT = 2
"""The newest stored layout this version reads: ``{"format": N, "query": "<SQL>"}``,
format 2 being a query with ``GROUP BY``, format 3 native pairs and format 4
native dedup results. A ``kind`` key beside it is for
format 2 being a query with ``GROUP BY``. A ``kind`` key beside it is for
readers older than the format number."""
@@ -272,38 +179,31 @@ class AsyncMaterializedView:
return _definition_from_json(raw)
async def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source.
The refresh is incremental when the source's changes can be
reconciled into the view -- rows added, changed or removed since the
last one -- and otherwise rebuilds. ``full=True`` forces a rebuild;
Every refresh rebuilds the complete view. For local views,
``source_version`` refreshes to that source version instead of the
latest.
Concurrent refreshes of one view do not duplicate its rows. Two that
plan the same source rows conflict on commit, and the loser raises
rather than writing them a second time.
latest; remote SQL refreshes do not support source-version pinning.
"""
return await self._table._inner.refresh_materialized_view(
full=full, source_version=source_version
source_version=source_version
)
async def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "AsyncJob[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
The job may already be complete for a local view. On LanceDB Cloud
and Enterprise, its ``id`` is the server job identifier returned by
the refresh endpoint.
The job may already be complete for a local view. Remote refreshes are
submitted through the SQL service.
"""
from ._lancedb import RefreshMaterializedViewResult
return _typed_job(
await self._table._inner.refresh_materialized_view_async(
full=full, source_version=source_version
source_version=source_version
),
RefreshMaterializedViewResult.from_json,
)
@@ -335,22 +235,18 @@ class MaterializedView:
return LOOP.run(self._async.definition())
def refresh(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "RefreshMaterializedViewResult":
"""Recompute the view from its source. See
[AsyncMaterializedView.refresh][lancedb.materialized_view.AsyncMaterializedView.refresh]."""
return LOOP.run(self._async.refresh(full=full, source_version=source_version))
return LOOP.run(self._async.refresh(source_version=source_version))
def refresh_async(
self, *, full: bool = False, source_version: Optional[int] = None
self, *, source_version: Optional[int] = None
) -> "Job[RefreshMaterializedViewResult]":
"""Submit a refresh and return its job without waiting.
See
[AsyncMaterializedView.refresh_async][lancedb.materialized_view.AsyncMaterializedView.refresh_async].
"""
return Job(
LOOP.run(
self._async.refresh_async(full=full, source_version=source_version)
)
)
return Job(LOOP.run(self._async.refresh_async(source_version=source_version)))
python/python/tests/test_materialized_views.py
+17 -191
View File
@@ -97,24 +97,6 @@ def test_remote_list_uses_namespace_route():
assert requests == ["/v1/namespace/$/materialized_view/list"]
def test_remote_create_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
client_config={"retry_config": {"retries": 0}},
)
job = db.create_materialized_view_async("adults", "people", where="age >= 18")
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/adults/create",
{"query": 'SELECT * FROM "people" WHERE age >= 18', "with_no_data": False},
)
]
def test_remote_drop_async_returns_server_job():
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
@@ -221,29 +203,29 @@ def test_definition_round_trips(tmp_path):
)
def test_incremental_refresh_after_append(tmp_path):
def test_full_refresh_after_append(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").add([{"name": "alan", "age": 41}])
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 4
assert view.table.count_rows() == 4
assert view.refresh().mode == "no_op"
assert view.refresh().mode == "rebuild"
def test_incremental_refresh_after_update(tmp_path):
def test_full_refresh_after_update(tmp_path):
db = make_db(tmp_path)
view = db.create_materialized_view("copy", "people", with_no_data=True)
view.refresh()
db.open_table("people").update(where="name = 'kid'", values={"age": 8})
result = view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 3
rows = view.table.search().to_list()
assert sorted(row["age"] for row in rows) == [8, 36, 85]
@@ -331,8 +313,8 @@ async def test_async_create_and_refresh_jobs(tmp_path):
@pytest.mark.asyncio
async def test_async_incremental(tmp_path):
db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS)
async def test_async_refresh_rebuilds(tmp_path):
db = await lancedb.connect_async(tmp_path)
await db.create_table("people", [{"name": "ada", "age": 36}])
view = await db.create_materialized_view("copy", "people", with_no_data=True)
await view.refresh()
@@ -340,15 +322,15 @@ async def test_async_incremental(tmp_path):
table = await db.open_table("people")
await table.add([{"name": "alan", "age": 41}])
result = await view.refresh()
assert result.mode == "incremental"
assert result.rows_written == 1
assert result.mode == "rebuild"
assert result.rows_written == 2
def test_source_requires_stable_row_ids(tmp_path):
def test_source_does_not_require_stable_row_ids(tmp_path):
db = lancedb.connect(tmp_path)
db.create_table("plain", [{"x": 1}])
with pytest.raises(Exception, match="stable row ids"):
db.create_materialized_view("v", "plain")
view = db.create_materialized_view("v", "plain")
assert view.table.count_rows() == 1
def test_bare_select_names_are_quoted(tmp_path):
@@ -373,7 +355,7 @@ def test_scalar_select_is_one_column(tmp_path):
view = db.create_materialized_view("just_name", "people", select="name")
view.refresh()
rows = view.table.search().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
assert sorted(row["name"] for row in rows) == ["ada", "grace", "kid"]
@@ -384,7 +366,7 @@ async def test_async_scalar_select_is_one_column(tmp_path):
view = await db.create_materialized_view("just_name", "people", select="name")
await view.refresh()
rows = await view.table.query().to_list()
assert set(rows[0]) - {"__source_row_id"} == {"name"}
assert set(rows[0]) == {"name"}
def test_limit_above_i64_max_is_refused(tmp_path):
@@ -520,166 +502,10 @@ def test_stored_queries_and_legacy_layouts_are_read():
"SELECT * FROM people"
)
for fmt, function in [(3, "vector_duplicate_pairs"), (4, "vector_dedup")]:
native = f"SELECT * FROM {function}('images', 3, 'phash', 4)"
assert read({"format": fmt, "query": native}).query == native
# A newer writer's layout is reported, never guessed at.
for newer in (
{"format": 5, "query": query},
{"format": 3, "query": query},
{"kind": "select_v3", "source_table": "people"},
):
with pytest.raises(NotImplementedError, match="cannot refresh"):
read(newer)
def indexed_dedup_source(db, name="images"):
import pyarrow as pa
from lancedb.index import IvfFlat
vectors = pa.FixedSizeListArray.from_arrays(
pa.array(
[v for i in range(256) for v in (float(i if i < 3 else i * 10), 0.0)],
type=pa.float32(),
),
2,
)
table = db.create_table(name, pa.table({"id": range(256), "vector": vectors}))
table.create_index("vector", config=IvfFlat(num_partitions=1))
return table
@pytest.mark.parametrize("stable", [False, True])
def test_vector_dedup_source_captures_snapshot_and_preserves_source(tmp_path, stable):
db = lancedb.connect(
tmp_path,
storage_options={"new_table_enable_stable_row_ids": str(stable).lower()},
)
source = indexed_dedup_source(db, "images.with.dot")
version = source.version
expression = lancedb.vector_dedup(
source.name, column="vector", distance_threshold=1
)
view = db.create_materialized_view("clean", expression)
assert "vector_dedup(" in view.definition.query
assert f", {version}, 'vector', 1.0)" in view.definition.query
assert view.table.count_rows() == 255
assert [r["id"] for r in view.table.search().where("id < 3").to_list()] == [0, 2]
assert source.version == version
assert source.count_rows() == 256
source.delete("id = 0")
view.refresh_async().wait()
assert view.table.search().where("id = 0").to_list() == [
{"id": 0, "vector": [0.0, 0.0]}
]
assert expression.dataset_version is None # reusable immutable declaration
def test_vector_dedup_create_job_and_explicit_snapshot(tmp_path):
db = lancedb.connect(tmp_path)
source = indexed_dedup_source(db)
version = source.version
source.delete("id = 0")
job = db.create_materialized_view_async(
"clean",
lancedb.vector_dedup(
"images", column="vector", distance_threshold=1, dataset_version=version
),
)
assert job.wait() is None
view = db.open_materialized_view("clean")
assert view.table.count_rows() == 255
assert len(view.table.search().where("id = 0").to_list()) == 1
@pytest.mark.asyncio
async def test_async_vector_dedup_declaration_and_creation(tmp_path):
sync_db = lancedb.connect(tmp_path)
source = indexed_dedup_source(sync_db)
db = await lancedb.connect_async(tmp_path)
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
view = await db.create_materialized_view("clean", expression, with_no_data=True)
assert await view.table.count_rows() == 0
# The snapshot is captured at declaration even if refresh comes later.
source.delete("id = 0")
job = await view.refresh_async()
await job.wait()
assert await view.table.count_rows() == 255
another = await db.create_materialized_view_async(
"later", expression, with_no_data=True
)
await another.wait()
later = await db.open_materialized_view("later")
assert (await later.definition()).query != (await view.definition()).query
assert await later.table.count_rows() == 0
@pytest.mark.parametrize(
"options", [{"select": ["id"]}, {"select": []}, {"where": "id > 0"}, {"limit": 1}]
)
def test_vector_dedup_rejects_extra_clauses_before_opening_source(tmp_path, options):
db = lancedb.connect(tmp_path)
expression = lancedb.vector_dedup("missing", column="vector", distance_threshold=1)
with pytest.raises(ValueError, match="cannot be combined"):
db.create_materialized_view("bad", expression, **options)
assert db.list_tables().tables == []
@pytest.mark.parametrize("version", [0, -1, True, 1.5, "2", 2**64])
def test_vector_dedup_rejects_invalid_snapshot(version):
with pytest.raises(ValueError, match="dataset_version"):
lancedb.vector_dedup(
"images", column="vector", distance_threshold=1, dataset_version=version
)
@pytest.mark.parametrize(
"threshold", [float("nan"), float("inf"), -float("inf"), 1e100, 10**1000, True, "4"]
)
def test_vector_dedup_rejects_invalid_threshold(threshold):
with pytest.raises((ValueError, TypeError), match="distance_threshold"):
lancedb.vector_dedup("images", column="vector", distance_threshold=threshold)
def test_vector_dedup_descriptor_is_frozen():
from dataclasses import FrozenInstanceError
expression = lancedb.vector_dedup("images", column="vector", distance_threshold=1)
with pytest.raises(FrozenInstanceError):
expression.column = "other"
@pytest.mark.parametrize("with_no_data", [False, True])
def test_remote_vector_dedup_uses_mv_endpoint_without_sql_client(with_no_data):
with mock_remote_materialized_view_create() as (host, requests):
db = lancedb.connect(
"db://dev",
api_key="fake",
host_override=host,
sql_host_override="grpc://127.0.0.1:1",
client_config={"retry_config": {"retries": 0}},
)
job = db.create_materialized_view_async(
"clean",
lancedb.vector_dedup(
"images.with'dot",
column="vector'field",
distance_threshold=4,
dataset_version=7,
),
with_no_data=with_no_data,
)
assert job.id == "mv-create-123"
assert requests == [
(
"/v1/materialized_view/clean/create",
{
"query": (
"SELECT * FROM vector_dedup('`images.with''dot`', "
"7, 'vector''field', 4.0)"
),
"with_no_data": with_no_data,
},
)
]
python/src/connection.rs
+2 -18
View File
@@ -407,8 +407,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view(
self_: PyRef<'_, Self>,
name: String,
@@ -417,17 +416,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
@@ -443,8 +435,7 @@ impl Connection {
})
}
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, *, vector_source_json=None))]
#[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false))]
pub fn create_materialized_view_async(
self_: PyRef<'_, Self>,
name: String,
@@ -453,17 +444,10 @@ impl Connection {
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
vector_source_json: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.get_inner()?.clone();
future_into_py(self_.py(), async move {
let mut builder = inner.create_materialized_view(name, source);
if let Some(source) = vector_source_json {
builder = builder.vector_source(
serde_json::from_str(&source)
.map_err(|e| PyValueError::new_err(e.to_string()))?,
);
}
if let Some(projections) = projections {
builder = builder.select(projections);
}
python/src/table.rs
+7 -13
View File
@@ -592,8 +592,6 @@ impl From<lancedb::RefreshMaterializedViewResult> for RefreshMaterializedViewRes
fn from(result: lancedb::RefreshMaterializedViewResult) -> Self {
let mode = match result.mode {
lancedb::RefreshMode::Rebuild => "rebuild",
lancedb::RefreshMode::Incremental => "incremental",
lancedb::RefreshMode::NoOp => "no_op",
};
Self {
mode: mode.to_string(),
@@ -1817,10 +1815,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1828,7 +1825,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1837,10 +1834,9 @@ impl Table {
})
}
#[pyo3(signature = (full=false, source_version=None))]
#[pyo3(signature = (source_version=None))]
pub fn refresh_materialized_view_async(
self_: PyRef<'_, Self>,
full: bool,
source_version: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
let inner = self_.inner_ref()?.clone();
@@ -1848,7 +1844,7 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
let mut builder = view.refresh().full(full);
let mut builder = view.refresh();
if let Some(version) = source_version {
builder = builder.source_version(version);
}
@@ -1863,11 +1859,9 @@ impl Table {
let view = lancedb::MaterializedView::from_table(inner)
.await
.infer_error()?;
view.definition().to_json().map_err(|err| {
PyRuntimeError::new_err(format!(
"failed to serialize materialized-view definition: {err}"
))
})
Ok(lancedb::materialized_view::definition_metadata_from_sql(
view.definition_sql(),
))
})
}
rust/lancedb/src/job.rs
+50 -1
View File
@@ -3,6 +3,8 @@
//! Handles to operations a server may run asynchronously.
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use arrow_array::RecordBatch;
@@ -250,6 +252,18 @@ where
pub(crate) fn spawned(task: JoinHandle<Result<T>>) -> Self {
Self::new_typed(Box::new(SpawnedJob::new(task)))
}
pub(crate) fn spawned_with_cancellation<F, Fut>(task: JoinHandle<Result<T>>, cancel: F) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<()>> + Send + 'static,
{
let cancellation: Cancellation = Arc::new(move || Box::pin(cancel()));
Self::new_typed(Box::new(SpawnedJob::new_with_cancellation(
task,
cancellation,
)))
}
}
impl<T> Job<T>
@@ -501,10 +515,28 @@ impl Outcome {
struct SpawnedJob {
outcome: watch::Receiver<Option<Outcome>>,
abort: AbortHandle,
cancellation: Option<Cancellation>,
}
type CancellationFuture = Pin<Box<dyn Future<Output = Result<()>> + Send>>;
type Cancellation = Arc<dyn Fn() -> CancellationFuture + Send + Sync>;
impl SpawnedJob {
fn new<T>(task: JoinHandle<Result<T>>) -> Self
where
T: Serialize + Send + 'static,
{
Self::new_inner(task, None)
}
fn new_with_cancellation<T>(task: JoinHandle<Result<T>>, cancellation: Cancellation) -> Self
where
T: Serialize + Send + 'static,
{
Self::new_inner(task, Some(cancellation))
}
fn new_inner<T>(task: JoinHandle<Result<T>>, cancellation: Option<Cancellation>) -> Self
where
T: Serialize + Send + 'static,
{
@@ -526,7 +558,11 @@ impl SpawnedJob {
};
let _ = tx.send(Some(outcome));
});
Self { outcome, abort }
Self {
outcome,
abort,
cancellation,
}
}
}
@@ -556,6 +592,19 @@ impl JobHandle for SpawnedJob {
}
async fn cancel(&self) -> Result<()> {
if let Some(cancellation) = &self.cancellation {
let mut outcome = self.outcome.clone();
tokio::select! {
biased;
settled = outcome.wait_for(Option::is_some) => {
settled.map_err(|_| Error::Runtime {
message: "job outcome was dropped before cancellation completed".to_string(),
})?;
return Ok(());
}
result = cancellation() => result?,
}
}
self.abort.abort();
Ok(())
}
rust/lancedb/src/materialized_view.rs
+208 -2527
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/materialized_view/differential.rs
-730
View File
@@ -1,730 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Differential refresh testing.
//!
//! The refresh contract is a property: after any sequence of source
//! mutations, a view maintained by default (incremental-where-possible)
//! refreshes equals the definition evaluated against the source directly,
//! and so does a forced rebuild. The oracle is an independent read of the
//! source -- plain column scan, filter applied in Rust -- so it shares
//! nothing with the refresh path it checks.
//!
//! The oracle runs after every step, not just at the end: a later mutation
//! that forces a rebuild would silently heal an incremental error, and those
//! transient errors are exactly the bugs this exists to catch.
use arrow_array::{Float32Array, Int32Array, RecordBatch};
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
use futures::{StreamExt, TryStreamExt};
use lance::dataset::NewColumnTransform;
use std::sync::Arc;
use super::MaterializedView;
use super::refresh::RefreshMode;
use crate::connect;
use crate::connection::Connection;
use crate::query::{ExecutableQuery, QueryBase, Select};
use crate::table::{CompactionOptions, OptimizeAction, Table};
/// One source mutation, one per correctness-relevant class.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SrcOp {
/// Fresh non-colliding ids of both parities, so every other op has
/// view-resident rows to act on: the only op that should refresh
/// incrementally.
AppendNew,
/// Deletion in surviving fragments must break the pure-append check.
DeleteEven,
/// An in-place update; on the filtered shape it crosses the predicate,
/// so rows must leave the view.
UpdateOddScore,
/// Fragment rewrite/renumber must break the pure-append check.
Compact,
/// A column the view does not read must NOT force a rebuild.
AddColumn,
/// merge_insert commits an Update whose by-source arm deletes rows, so a
/// classifier that reads Update as "changed only" loses those deletions.
MergeDropLargest,
/// merge_insert that both changes existing rows and inserts new ones in
/// one transaction.
MergeUpsert,
}
const ALL_OPS: [SrcOp; 7] = [
SrcOp::AppendNew,
SrcOp::DeleteEven,
SrcOp::UpdateOddScore,
SrcOp::Compact,
SrcOp::AddColumn,
SrcOp::MergeDropLargest,
SrcOp::MergeUpsert,
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Shape {
/// SELECT id, score.
Identity,
/// SELECT id, score WHERE score > 50: additionally sensitive to rows
/// crossing the predicate.
Filtered,
/// SELECT id, score LIMIT 4. Which rows are held depends on the order
/// they were first materialized, so the oracle checks containment and
/// the cap rather than equality.
Limited,
}
impl Shape {
fn filter(&self) -> Option<&'static str> {
match self {
Self::Identity | Self::Limited => None,
Self::Filtered => Some("score > 50"),
}
}
fn matches(&self, score: f32) -> bool {
match self {
Self::Identity | Self::Limited => true,
Self::Filtered => score > 50.0,
}
}
fn limit(&self) -> Option<usize> {
match self {
Self::Limited => Some(4),
_ => None,
}
}
}
struct Case {
conn: Connection,
source: Table,
view: MaterializedView,
shape: Shape,
next_id: i32,
added_columns: u32,
}
fn rows_batch(ids: &[i32]) -> RecordBatch {
let scores: Vec<f32> = ids.iter().map(|id| (*id * 10) as f32).collect();
RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, true),
ArrowField::new("score", DataType::Float32, true),
])),
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(Float32Array::from(scores)),
],
)
.unwrap()
}
fn merge_batch(ids: &[i32]) -> RecordBatch {
let scores: Vec<f32> = ids.iter().map(|id| (*id * 10 + 5) as f32).collect();
RecordBatch::try_new(
Arc::new(ArrowSchema::new(vec![
ArrowField::new("id", DataType::Int32, true),
ArrowField::new("score", DataType::Float32, true),
])),
vec![
Arc::new(Int32Array::from(ids.to_vec())),
Arc::new(Float32Array::from(scores)),
],
)
.unwrap()
}
impl Case {
async fn new(shape: Shape) -> Self {
let conn = connect("memory://").execute().await.unwrap();
let source = conn
.create_table("src", rows_batch(&[1, 2, 3, 4]))
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
let mut builder = conn
.create_materialized_view("view", "src")
.select([("id", "id"), ("score", "score")]);
if let Some(filter) = shape.filter() {
builder = builder.only_if(filter);
}
if let Some(limit) = shape.limit() {
builder = builder.limit(limit as u64);
}
let view = builder.execute().await.unwrap();
Self {
conn,
source,
view,
shape,
next_id: 100,
added_columns: 0,
}
}
async fn apply(&mut self, op: SrcOp) {
match op {
SrcOp::AppendNew => {
// Mixed parity: the middle id is odd, so UpdateOddScore always
// has a filter-matching appended row to evict.
let ids = vec![self.next_id, self.next_id + 101, self.next_id + 202];
self.next_id += 303;
self.source.add(rows_batch(&ids)).execute().await.unwrap();
}
SrcOp::DeleteEven => {
self.source.delete("id % 2 = 0").await.unwrap();
}
SrcOp::UpdateOddScore => {
self.source
.update()
.column("score", "-1.0")
.only_if("id % 2 = 1")
.execute()
.await
.unwrap();
}
SrcOp::Compact => {
self.source
.optimize(OptimizeAction::Compact {
options: CompactionOptions::default(),
remap_options: None,
})
.await
.unwrap();
}
SrcOp::MergeDropLargest => {
let mut ids = self.source_ids().await;
ids.sort_unstable();
ids.pop();
if ids.is_empty() {
return;
}
let batch = rows_batch(&ids);
let reader =
arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
let mut merge = self.source.merge_insert(&["id"]);
merge.when_not_matched_by_source_delete(None);
merge.execute(Box::new(reader)).await.unwrap();
}
SrcOp::MergeUpsert => {
let mut ids = self.source_ids().await;
ids.sort_unstable();
// One row that exists (updated in place) and one that does not.
let existing = ids.first().copied().unwrap_or(self.next_id);
let fresh = self.next_id;
self.next_id += 1;
let batch = merge_batch(&[existing, fresh]);
let reader =
arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema());
let mut merge = self.source.merge_insert(&["id"]);
merge
.when_matched_update_all(None)
.when_not_matched_insert_all();
merge.execute(Box::new(reader)).await.unwrap();
}
SrcOp::AddColumn => {
self.added_columns += 1;
let field = ArrowField::new(
format!("extra_{}", self.added_columns),
DataType::Int32,
true,
);
self.source
.add_columns()
.transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(
vec![field],
))))
.execute()
.await
.unwrap();
}
}
}
async fn source_ids(&self) -> Vec<i32> {
read_rows(
self.source
.query()
.select(Select::columns(&["id", "score"])),
)
.await
.into_iter()
.map(|(id, _)| id)
.collect()
}
/// The definition's result, read independently of the refresh path:
/// plain column scan, filter applied here, sorted.
async fn oracle(&self) -> Vec<(i32, i32)> {
let mut rows = read_rows(
self.source
.query()
.select(Select::columns(&["id", "score"])),
)
.await
.into_iter()
.filter(|(_, score)| self.shape.matches(*score as f32))
.collect::<Vec<_>>();
rows.sort_unstable();
rows
}
async fn view_rows(&self) -> Vec<(i32, i32)> {
let mut rows = read_rows(
self.view
.table()
.query()
.select(Select::columns(&["id", "score"])),
)
.await;
rows.sort_unstable();
rows
}
async fn check(&self, label: &str) -> Result<(), String> {
let expected = self.oracle().await;
let actual = self.view_rows().await;
let Some(cap) = self.shape.limit() else {
if expected != actual {
return Err(format!(
"{label}: view diverged from oracle\n expected: {expected:?}\n actual: {actual:?}"
));
}
return Ok(());
};
// A capped view holds some subset of the definition's result, never
// more than the cap, and never the same row twice.
if actual.len() > cap {
return Err(format!(
"{label}: view holds {} rows, over its cap of {cap}: {actual:?}",
actual.len()
));
}
let mut unique = actual.clone();
unique.dedup();
if unique.len() != actual.len() {
return Err(format!("{label}: view holds a row twice: {actual:?}"));
}
if let Some(stray) = actual.iter().find(|row| !expected.contains(row)) {
return Err(format!(
"{label}: view holds {stray:?}, which the definition does not select: {expected:?}"
));
}
// Below the cap the view must be complete, or a row was lost.
if actual.len() < cap.min(expected.len()) {
return Err(format!(
"{label}: view holds {} of {} selectable rows under a cap of {cap}: {actual:?}",
actual.len(),
expected.len()
));
}
Ok(())
}
}
async fn read_rows(query: impl ExecutableQuery) -> Vec<(i32, i32)> {
let batches = query
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
batches
.iter()
.flat_map(|batch| {
let ids = batch["id"].as_any().downcast_ref::<Int32Array>().unwrap();
let scores = batch["score"]
.as_any()
.downcast_ref::<Float32Array>()
.unwrap();
// Scores are integer-valued by construction; compare exactly.
(0..batch.num_rows())
.map(|i| (ids.value(i), scores.value(i) as i32))
.collect::<Vec<_>>()
})
.collect()
}
/// Drive one mutation sequence: refresh + oracle-check after every step,
/// then a forced rebuild checked against the same oracle.
async fn run_sequence(ops: &[SrcOp], shape: Shape) -> Result<(), String> {
let label = format!("{shape:?} {ops:?}");
let mut case = Case::new(shape).await;
case.view
.refresh()
.execute()
.await
.map_err(|e| format!("{label}: initial refresh failed: {e}"))?;
case.check(&format!("{label} (initial)")).await?;
for (step, op) in ops.iter().enumerate() {
case.apply(*op).await;
case.view
.refresh()
.execute()
.await
.map_err(|e| format!("{label}: refresh at step {step} failed: {e}"))?;
case.check(&format!("{label} (step {step}, {op:?})"))
.await?;
}
case.view
.refresh()
.full(true)
.execute()
.await
.map_err(|e| format!("{label}: final full refresh failed: {e}"))?;
case.check(&format!("{label} (final rebuild)")).await?;
// Silence the unused-connection lint without dropping it mid-case.
let _ = &case.conn;
Ok(())
}
/// Every op sequence up to `max_len`.
fn all_sequences(max_len: u32) -> Vec<Vec<SrcOp>> {
let mut sequences = Vec::new();
for len in 1..=max_len {
for mut index in 0..ALL_OPS.len().pow(len) {
let mut ops = Vec::with_capacity(len as usize);
for _ in 0..len {
ops.push(ALL_OPS[index % ALL_OPS.len()]);
index /= ALL_OPS.len();
}
sequences.push(ops);
}
}
sequences
}
async fn run_exhaustive(max_len: u32) {
let mut cases = Vec::new();
for shape in [Shape::Identity, Shape::Filtered, Shape::Limited] {
for ops in all_sequences(max_len) {
cases.push((ops, shape));
}
}
let failures: Vec<String> = futures::stream::iter(cases)
.map(|(ops, shape)| async move { run_sequence(&ops, shape).await.err() })
.buffer_unordered(8)
.filter_map(|failure| async move { failure })
.collect()
.await;
assert!(
failures.is_empty(),
"{} sequences diverged; first: {}",
failures.len(),
failures[0]
);
}
#[tokio::test(flavor = "multi_thread")]
async fn differential_exhaustive() {
run_exhaustive(3).await;
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "longer sweep; run manually"]
async fn differential_exhaustive_deep() {
run_exhaustive(4).await;
}
/// Named interleavings that double as repro handles. The mode assertions pin
/// the classifier, which value comparison alone cannot: a wrongly rebuilt
/// view still matches the oracle.
#[tokio::test(flavor = "multi_thread")]
async fn differential_named_regressions() {
// An append is the one op that must stay incremental.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AppendNew).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
case.check("append stays incremental").await.unwrap();
// A column the view does not read must not force a rebuild.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AddColumn).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 0);
// Compaction rearranges rows without changing them: the watermark
// advances and nothing rebuilds.
let mut case = Case::new(Shape::Identity).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::AppendNew).await;
case.view.refresh().execute().await.unwrap();
case.apply(SrcOp::Compact).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 0);
case.check("compaction alone").await.unwrap();
// Fragment bookkeeping stays coherent across the compaction: the next
// append is separable and computed alone.
case.apply(SrcOp::AppendNew).await;
let result = case.view.refresh().execute().await.unwrap();
assert_eq!(result.mode, RefreshMode::Incremental);
assert_eq!(result.rows_written, 3);
case.check("compact then append").await.unwrap();
// A row updated to no longer match the filter must leave the view --
// and the fixture must prove the eviction happened, not merely that the
// end state matches: an update that never touched a view-resident row
// would also "match".
let mut case = Case::new(Shape::Filtered).await;
case.apply(SrcOp::AppendNew).await;
case.view.refresh().execute().await.unwrap();
let before = case.view_rows().await.len();
case.apply(SrcOp::UpdateOddScore).await;
case.view.refresh().execute().await.unwrap();
let after = case.view_rows().await.len();
assert!(
after < before,
"no view-resident row was evicted ({before} -> {after}); the fixture \
no longer exercises the filtered-update transition"
);
case.check("update crosses the filter").await.unwrap();
}
// ---------------------------------------------------------------------------
// Concurrency
// ---------------------------------------------------------------------------
//
// The sequential cases above cannot observe a cross-process race: the
// per-view refresh lock is process-local, so a second refresh in this
// process queues behind the first. What is missing is not more op
// sequences but a second process. These cases add one, and assert the same
// property the harness always asserts -- the view holds each row once.
/// Rows the definition selects from the source: every id but the first,
/// read straight from the source, sharing nothing with the refresh path.
async fn concurrency_oracle(conn: &Connection) -> Vec<i32> {
let batches: Vec<RecordBatch> = conn
.open_table("src")
.execute()
.await
.unwrap()
.query()
.select(Select::columns(&["id"]))
.execute()
.await
.unwrap()
.try_collect()
.await
.unwrap();
let mut ids = Vec::new();
for batch in &batches {
let column = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..batch.num_rows() {
if column.value(i) > 1 {
ids.push(column.value(i));
}
}
}
ids.sort_unstable();
ids
}
/// The view's ids, sorted.
async fn concurrency_view_ids(conn: &Connection) -> Vec<i32> {
let batches: Vec<RecordBatch> = conn
.open_table("mv")
.execute()
.await
.unwrap()
.query()
.select(Select::columns(&["id"]))
.execute()
.await
.unwrap()
.try_collect()
.await
.unwrap();
let mut ids = Vec::new();
for batch in &batches {
let column = batch
.column(0)
.as_any()
.downcast_ref::<Int32Array>()
.unwrap();
for i in 0..batch.num_rows() {
ids.push(column.value(i));
}
}
ids.sort_unstable();
ids
}
/// One refresh of the view at `MV_RACE_DIR`, in its own process.
///
/// Setup happens before the start barrier so warm-up does not stagger the
/// two processes. What makes the race certain rather than likely is the
/// second barrier inside `refresh()` itself, which holds every participant
/// between staging and commit.
#[tokio::test]
#[ignore = "spawned as a child process by the concurrency cases"]
async fn cross_process_refresh_child() {
let Ok(dir) = std::env::var("MV_RACE_DIR") else {
return;
};
let dir = std::path::PathBuf::from(dir);
let tag = std::env::var("MV_RACE_TAG").unwrap();
let conn = connect(dir.to_str().unwrap()).execute().await.unwrap();
let table = conn.open_table("mv").execute().await.unwrap();
let _ = table.schema().await.unwrap();
let _ = table.count_rows(None).await.unwrap();
let source = conn.open_table("src").execute().await.unwrap();
let _ = source.count_rows(None).await.unwrap();
let view = MaterializedView::from_table(table).await.unwrap();
std::fs::write(dir.join(format!("ready-{tag}")), b"1").unwrap();
while !dir.join("START").exists() {
std::thread::sleep(std::time::Duration::from_millis(2));
}
let outcome = match view.refresh().execute().await {
Ok(result) => format!("committed rows={}", result.rows_written),
Err(err) if is_commit_conflict(&err) => "conflicted".to_string(),
Err(err) => format!("failed {err}"),
};
std::fs::write(dir.join(format!("outcome-{tag}")), outcome).unwrap();
}
/// Whether a refresh lost its commit to a concurrent one, as opposed to
/// failing for any other reason.
fn is_commit_conflict(err: &crate::Error) -> bool {
let text = err.to_string();
text.contains("Retryable commit conflict") || text.contains("preempted by concurrent")
}
/// Two processes refreshing one view concurrently must leave the view
/// equal to the oracle: each selected row present exactly once.
///
/// Both plan the same incremental delta from one watermark. A refresh is
/// meant to land on the generation it planned or leave nothing behind, so
/// at most one of them may write.
#[tokio::test]
async fn concurrent_refreshes_hold_each_row_once() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().to_str().unwrap().to_string();
let conn = connect(&path).execute().await.unwrap();
conn.create_table("src", rows_batch(&[1, 2, 3, 4]))
.write_options(crate::materialized_view::tests::stable_row_ids())
.execute()
.await
.unwrap();
let view = conn
.create_materialized_view("mv", "src")
.select([("id", "id"), ("score", "score")])
.only_if("id > 1")
.execute()
.await
.unwrap();
// Seed the watermark so the racing refreshes are both incremental.
view.refresh().execute().await.unwrap();
// Large enough that a refresh is real work rather than a formality.
let ids: Vec<i32> = (100..200_100).collect();
conn.open_table("src")
.execute()
.await
.unwrap()
.add(rows_batch(&ids))
.execute()
.await
.unwrap();
let tags = ["a", "b"];
let exe = std::env::current_exe().unwrap();
let children: Vec<std::process::Child> = tags
.iter()
.map(|tag| {
std::process::Command::new(&exe)
.args([
"--exact",
"materialized_view::differential::cross_process_refresh_child",
"--ignored",
"--nocapture",
])
.env("MV_RACE_DIR", dir.path())
.env("MV_RACE_SYNC", dir.path())
.env("MV_RACE_PEERS", "2")
.env("MV_RACE_TAG", tag)
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.spawn()
.unwrap()
})
.collect();
let deadline = std::time::Instant::now() + std::time::Duration::from_secs(180);
while tags
.iter()
.any(|tag| !dir.path().join(format!("ready-{tag}")).exists())
{
assert!(
std::time::Instant::now() < deadline,
"children never became ready"
);
std::thread::sleep(std::time::Duration::from_millis(10));
}
std::fs::write(dir.path().join("START"), b"1").unwrap();
for (tag, mut child) in tags.iter().zip(children) {
let status = loop {
match child.try_wait().unwrap() {
Some(status) => break status,
None if std::time::Instant::now() >= deadline => {
child.kill().unwrap();
panic!("child {tag} never finished");
}
None => std::thread::sleep(std::time::Duration::from_millis(10)),
}
};
assert!(status.success(), "child {tag} exited {status}");
}
// Both refreshes reached the commit boundary before either committed --
// the in-refresh barrier guarantees it -- so exactly one may win.
let outcomes: Vec<String> = tags
.iter()
.map(|tag| {
std::fs::read_to_string(dir.path().join(format!("outcome-{tag}")))
.unwrap_or_else(|_| panic!("child {tag} recorded no outcome"))
})
.collect();
for tag in tags {
assert!(
dir.path().join(format!("planned-{tag}")).exists(),
"child {tag} never reached the commit boundary, so nothing was synchronized"
);
}
let committed = outcomes.iter().filter(|o| o.contains("committed")).count();
let conflicted = outcomes.iter().filter(|o| o.contains("conflicted")).count();
assert_eq!(
(committed, conflicted),
(1, 1),
"exactly one refresh may win the generation both planned: {outcomes:?}"
);
let expected = concurrency_oracle(&conn).await;
let actual = concurrency_view_ids(&conn).await;
assert_eq!(
actual.len(),
expected.len(),
"the view holds {} rows, the oracle {}: a losing refresh left rows behind",
actual.len(),
expected.len()
);
assert_eq!(actual, expected, "the view does not match the oracle");
}
rust/lancedb/src/materialized_view/duplicate_pairs.rs
-478
View File
@@ -1,478 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Native pair generation as partitioned materialized-view refresh work.
use std::sync::Arc;
use arrow_schema::{Schema, SchemaRef};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_sql::sqlparser::ast::{
Expr, FunctionArg, FunctionArgExpr, TableFunctionArgs, Value,
};
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::transaction::Operation;
use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams};
use lance_table::format::Fragment;
use serde::{Deserialize, Serialize};
use super::grouped_units::GroupedRefreshPlan;
use super::refresh::{
RefreshMaterializedViewResult, ensure_incarnation, ensure_no_mem_wal, open_source,
};
use super::{MaterializedViewDefinition, Planned, ViewProjection};
use crate::table::Table;
use crate::table::datafusion::udtf::duplicate_pairs::{
DUPLICATE_PAIRS_OPERATOR_VERSION, DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec,
duplicate_pair_schema, plan_duplicate_pairs,
};
use crate::{Error, Result};
pub(super) const FUNCTION_NAME: &str = "vector_duplicate_pairs";
pub(super) const DEDUP_FUNCTION_NAME: &str = "vector_dedup";
/// The result of a pinned native vector source.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum VectorSourceKind {
/// The qualifying pair relation.
Pairs,
/// Original source rows retained by greedy direct representatives.
Dedup,
}
/// A native table-function source, pinned for the lifetime of the view.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct VectorSource {
pub kind: VectorSourceKind,
pub dataset_version: u64,
pub column: String,
/// Validated numeric SQL spelling; execution and task identity use f32 bits.
pub distance_threshold: String,
}
impl VectorSource {
pub(super) fn function_name(&self) -> &'static str {
match self.kind {
VectorSourceKind::Pairs => FUNCTION_NAME,
VectorSourceKind::Dedup => DEDUP_FUNCTION_NAME,
}
}
pub(crate) fn config(&self) -> Result<DuplicatePairsConfig> {
let threshold: f32 = self
.distance_threshold
.parse()
.map_err(|_| invalid("invalid pair threshold"))?;
if self.dataset_version == 0 || self.column.is_empty() || !threshold.is_finite() {
return Err(invalid(
"duplicate pairs require a positive snapshot, a column and a finite threshold",
));
}
Ok(DuplicatePairsConfig {
dataset_version: self.dataset_version,
column: self.column.clone(),
distance_threshold: threshold,
})
}
}
fn invalid(message: impl Into<String>) -> Error {
Error::InvalidInput {
message: message.into(),
}
}
pub(super) fn source(
args: &TableFunctionArgs,
kind: VectorSourceKind,
) -> Result<(Vec<String>, String, VectorSource)> {
if args.settings.is_some() {
return Err(invalid("duplicate pairs do not accept function settings"));
}
let exprs = args
.args
.iter()
.map(|arg| match arg {
FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => Ok(expr),
_ => Err(invalid(
"duplicate pairs require four positional literal arguments",
)),
})
.collect::<Result<Vec<_>>>()?;
let [table, version, column, threshold] = exprs.as_slice() else {
return Err(invalid(
"native vector sources require (table, dataset_version, column, distance_threshold)",
));
};
let string = |expr: &Expr| match expr {
Expr::Value(value) => match &value.value {
Value::SingleQuotedString(value) => Ok(value.clone()),
_ => Err(invalid("table and column must be string literals")),
},
_ => Err(invalid("table and column must be string literals")),
};
let table = string(table)?;
let plain = MaterializedViewDefinition::from_sql(&format!("SELECT * FROM {table}"))?;
if plain.vector_source.is_some()
|| plain.lateral.is_some()
|| plain.filter.is_some()
|| plain.is_grouped()
|| plain.limit.is_some()
{
return Err(invalid("duplicate pairs must name one source table"));
}
let dataset_version = match version {
Expr::Value(value) => match &value.value {
Value::Number(number, _) => number
.parse()
.map_err(|_| invalid("invalid source version"))?,
_ => return Err(invalid("source version must be an integer literal")),
},
_ => return Err(invalid("source version must be an integer literal")),
};
let config = VectorSource {
kind,
dataset_version,
column: string(column)?,
distance_threshold: threshold.to_string(),
};
config.config()?;
Ok((plain.source_namespace, plain.source_table, config))
}
pub(super) fn check_shape(definition: &MaterializedViewDefinition) -> Result<()> {
if definition
.vector_source
.as_ref()
.is_some_and(|s| s.kind == VectorSourceKind::Dedup)
{
if !definition.selects_star()
|| definition.filter.is_some()
|| definition.limit.is_some()
|| definition.lateral.is_some()
|| definition.is_grouped()
{
return Err(invalid(
"a vector_dedup view must select * without other clauses",
));
}
return Ok(());
}
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect::<Vec<_>>();
if (!definition.selects_star() && definition.projections != projections)
|| definition.filter.is_some()
|| definition.limit.is_some()
|| definition.lateral.is_some()
|| definition.is_grouped()
{
return Err(invalid(
"a native pair view selects *, or row_id_a, row_id_b, distance, without other clauses",
));
}
Ok(())
}
pub(super) fn plan(
source_schema: SchemaRef,
definition: &MaterializedViewDefinition,
) -> Result<Planned> {
check_shape(definition)?;
let config = definition.vector_source.as_ref().expect("native source");
config.config()?;
source_schema.field_with_name(&config.column)?;
if config.kind == VectorSourceKind::Dedup {
return Ok(Planned {
definition: definition.clone(),
fields: source_schema
.fields()
.iter()
.map(|f| super::without_declarations(f))
.collect(),
lineage: source_schema
.fields()
.iter()
.map(|f| (f.name().clone(), vec![f.name().clone()]))
.collect(),
inputs: source_schema
.fields()
.iter()
.map(|f| f.name().clone())
.collect(),
});
}
let schema = duplicate_pair_schema();
let projections = schema
.fields()
.iter()
.map(|field| ViewProjection {
output: field.name().clone(),
expression: field.name().clone(),
})
.collect();
let mut definition = definition.clone();
let inputs = vec![config.column.clone()];
definition.projections = projections;
Ok(Planned {
definition,
fields: schema.fields().iter().map(|f| f.as_ref().clone()).collect(),
lineage: Default::default(),
inputs,
})
}
/// A saved native refresh plan, including every physical partition identity.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct DuplicatePairsRefreshPlan {
pub snapshot: GroupedRefreshPlan,
pub source_uri: String,
pub definition_sql: String,
pub operator_version: u32,
pub threshold_bits: u32,
pub tasks: Vec<DuplicatePairTask>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenPairsPartition {
pub(super) unit: u32,
pub(super) plan_id: String,
pub(super) task: DuplicatePairTask,
pub(super) fragments: Vec<Fragment>,
pub(super) rows: u64,
}
impl DuplicatePairsRefreshPlan {
pub(super) fn identity(&self) -> Result<String> {
use sha2::{Digest, Sha256};
let encoded = serde_json::to_vec(&(
&self.snapshot,
&self.source_uri,
&self.definition_sql,
self.operator_version,
self.threshold_bits,
))
.map_err(|e| invalid(e.to_string()))?;
Ok(format!("{:x}", Sha256::digest(encoded)))
}
}
async fn open(view: &Table) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let native = view
.as_native()
.ok_or_else(|| invalid("native pair refresh requires a local view"))?;
native.dataset.ensure_mutable()?;
native.dataset.reload().await?;
let view_ds = native.dataset.get().await?.as_ref().clone();
let Some(super::StoredDefinition::Query(definition)) =
super::read_definition(&view_ds.schema().metadata)?
else {
return Err(invalid("view does not carry a supported definition"));
};
let config = definition
.vector_source
.as_ref()
.ok_or_else(|| invalid("not a native pair view"))?
.config()?;
let source = Arc::new(
open_source(view, &definition, None)
.await?
.checkout_version(config.dataset_version)
.await?,
);
ensure_no_mem_wal(&source, "source table", &definition.source_table).await?;
let planned = plan(Arc::new(Schema::from(source.schema())), &definition)?;
let expected = Schema::new(planned.fields);
let physical = Schema::from(view_ds.schema());
if physical.fields().len() != expected.fields().len()
|| physical
.fields()
.iter()
.zip(expected.fields())
.any(|(a, b)| {
a.name() != b.name()
|| a.data_type() != b.data_type()
|| a.is_nullable() != b.is_nullable()
})
{
return Err(invalid("native pair view schema changed since creation"));
}
Ok((view_ds, planned.definition, source))
}
pub(super) async fn plan_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<DuplicatePairsRefreshPlan> {
let (view_ds, definition, source) = open(view).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if pinned.is_some_and(|v| v != config.dataset_version) {
return Err(invalid(
"refresh source version differs from the native pair view's fixed snapshot",
));
}
let tasks = plan_duplicate_pairs(source.clone(), &config).await?;
let units = u32::try_from(tasks.len()).map_err(|_| invalid("too many index partitions"))?;
let incarnation = view_ds
.schema()
.metadata
.get(super::INCARNATION_META_KEY)
.cloned()
.ok_or_else(|| invalid("native pair view has no incarnation"))?;
Ok(DuplicatePairsRefreshPlan {
snapshot: GroupedRefreshPlan {
units,
source_version: config.dataset_version,
view_version: view_ds.version().version,
incarnation,
},
source_uri: source.uri().to_string(),
definition_sql: definition.to_sql(),
operator_version: DUPLICATE_PAIRS_OPERATOR_VERSION,
threshold_bits: config.distance_threshold.to_bits(),
tasks,
})
}
pub(super) async fn open_planned(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
) -> Result<(Dataset, MaterializedViewDefinition, Arc<Dataset>)> {
let (view_ds, definition, source) = open(view).await?;
ensure_incarnation(&view_ds, Some(&plan.snapshot.incarnation), view.name()).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if source.uri() != plan.source_uri
|| source.version().version != plan.snapshot.source_version
|| definition.to_sql() != plan.definition_sql
|| config.distance_threshold.to_bits() != plan.threshold_bits
|| plan.operator_version != DUPLICATE_PAIRS_OPERATOR_VERSION
|| plan.tasks.len() != plan.snapshot.units as usize
{
return Err(invalid(
"native pair task belongs to a different snapshot, configuration or operator",
));
}
Ok((view_ds, definition, source))
}
pub(super) async fn write_unit(
view: &Table,
unit: u32,
plan: &DuplicatePairsRefreshPlan,
) -> Result<WrittenPairsPartition> {
let (view_ds, definition, source) = open_planned(view, plan).await?;
let task = plan
.tasks
.get(unit as usize)
.ok_or_else(|| invalid("pair partition is outside the plan"))?
.clone();
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
let exec = DuplicatePairsExec::try_new(source, config, vec![task.clone()])?;
let stream = exec.execute(0, Arc::new(TaskContext::default()))?;
let rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let count = rows.clone();
let stream = stream.inspect_ok(move |batch| {
count.fetch_add(
batch.num_rows() as u64,
std::sync::atomic::Ordering::Relaxed,
);
});
let stream: SendableRecordBatchStream = Box::pin(
datafusion_physical_plan::stream::RecordBatchStreamAdapter::new(
duplicate_pair_schema(),
stream,
),
);
let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds)))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append { fragments } = txn.operation else {
return Err(invalid(
"native pair partition did not stage append fragments",
));
};
Ok(WrittenPairsPartition {
unit,
plan_id: plan.identity()?,
task,
fragments,
rows: rows.load(std::sync::atomic::Ordering::Relaxed),
})
}
pub(super) async fn commit(
view: &Table,
plan: &DuplicatePairsRefreshPlan,
units: Vec<WrittenPairsPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let (_, definition, source) = open_planned(view, plan).await?;
let config = definition
.vector_source
.as_ref()
.expect("native")
.config()?;
if plan_duplicate_pairs(source.clone(), &config).await? != plan.tasks {
return Err(invalid(
"pair plan does not cover the pinned index partitions exactly",
));
}
let identity = plan.identity()?;
let mut ordered = vec![None; plan.snapshot.units as usize];
for written in units {
if written.plan_id != identity
|| plan.tasks.get(written.unit as usize) != Some(&written.task)
{
return Err(invalid("pair partition receipt belongs to another plan"));
}
let slot = ordered
.get_mut(written.unit as usize)
.ok_or_else(|| invalid("unknown pair partition receipt"))?;
if slot.is_some() {
return Err(invalid("duplicate pair partition receipt"));
}
*slot = Some(written);
}
let units = ordered
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid("missing pair partition receipt"))?;
let rows = units.iter().map(|u| u.rows).sum();
let fragments = units.into_iter().flat_map(|u| u.fragments).collect();
super::grouped_units::commit_fragments(
view,
&plan.snapshot,
&definition,
source.manifest.timestamp_nanos,
fragments,
rows,
expected,
)
.await
}
rust/lancedb/src/materialized_view/grouped.rs
+13 -122
View File
@@ -12,7 +12,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use arrow_array::cast::AsArray;
use arrow_array::{Array, FixedSizeListArray, UInt32Array};
use arrow_buffer::NullBuffer;
use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef};
use arrow_schema::{DataType, Schema as ArrowSchema, SchemaRef};
use datafusion::catalog::default_table_source::provider_as_source;
use datafusion::datasource::empty::EmptyTable;
use datafusion::execution::FunctionRegistry;
@@ -36,20 +36,16 @@ use datafusion_sql::sqlparser::parser::Parser;
use futures::StreamExt;
use lance::Dataset;
use lance::index::{DatasetIndexExt, DatasetIndexInternalExt};
use lance_core::ROW_ID;
use lance_datafusion::exec::SessionContextExt;
use lance_index::metrics::NoOpMetricsCollector;
use lance_index::vector::VectorIndex;
use lance_index::vector::ivf::{IvfTransformer, new_ivf_transformer};
use lance_linalg::distance::DistanceType;
use lance_linalg::kernels::normalize_fsl;
use roaring::RoaringBitmap;
use uuid::Uuid;
use super::refresh::to_view_batch;
use super::{
MaterializedViewDefinition, Planned, SOURCE_ROW_ID_COLUMN, ViewProjection, ensure_immutable,
plan_filter, query, without_declarations,
MaterializedViewDefinition, Planned, ViewProjection, plan_filter, query, without_declarations,
};
use crate::{Error, Result};
@@ -68,50 +64,6 @@ pub(super) struct IvfBinding {
pub(super) transformer: Arc<IvfTransformer>,
/// A cosine index assigns L2 over unit vectors, as lance's index path does.
pub(super) normalize: bool,
/// Partitions of the model; the ids `ivf_partition` returns are below it.
pub(super) partitions: u32,
}
/// The segments of the one IVF index on a column, opened, and the fragments
/// they cover between them; every other fragment's rows are unindexed.
/// `covered` is `None` when a segment does not say which fragments it holds
/// (`fragment_bitmap: None`, which lance defines as unknown rather than
/// empty): its rows would be read once by the partition reader and again as
/// unindexed, so no caller may split this index.
pub(super) struct IvfSegments {
pub(super) segments: Vec<IvfSegment>,
pub(super) covered: Option<RoaringBitmap>,
}
/// One opened segment and the fragments it holds rows for. Its postings
/// still name rows of fragments it has since lost -- a column rewrite
/// attaches a new file and takes the fragment out of this bitmap -- so a
/// reader of its partitions keeps only the rows it still owns.
pub(super) struct IvfSegment {
pub(super) index: Arc<dyn VectorIndex>,
pub(super) fragments: Option<RoaringBitmap>,
}
impl IvfSegments {
pub(super) fn is_unindexed(&self, fragment_id: u64) -> Option<bool> {
Some(!self.covered.as_ref()?.contains(fragment_id as u32))
}
}
/// Coverage after folding in one more segment. Lance defines a segment
/// without a fragment bitmap as unknown coverage, not empty, and one such
/// segment is enough to lose the coverage of the whole index.
fn coverage(
covered: Option<RoaringBitmap>,
bitmap: Option<&RoaringBitmap>,
) -> Option<RoaringBitmap> {
match (covered, bitmap) {
(Some(mut covered), Some(bitmap)) => {
covered |= bitmap;
Some(covered)
}
_ => None,
}
}
/// `bindings` maps a source column to its index; planning runs unbound.
@@ -278,16 +230,6 @@ pub(super) async fn bind(
}
async fn bind_column(source: &Dataset, column: &str) -> Result<IvfBinding> {
bind_column_with_segments(source, column)
.await
.map(|(binding, _)| binding)
}
/// The binding plus the opened segments it was read from.
pub(super) async fn bind_column_with_segments(
source: &Dataset,
column: &str,
) -> Result<(IvfBinding, IvfSegments)> {
let field = source
.schema()
.field(column)
@@ -298,8 +240,6 @@ pub(super) async fn bind_column_with_segments(
// on its own fragments carries its own centroids, and grouping every row
// by one segment's model would bucket the other segments' rows wrongly.
let mut found: Option<(String, Uuid, DistanceType, FixedSizeListArray)> = None;
let mut segments = Vec::new();
let mut covered = Some(RoaringBitmap::new());
for index in source.load_indices().await?.iter() {
if index.fields != [field.id] {
continue;
@@ -314,12 +254,6 @@ pub(super) async fn bind_column_with_segments(
continue;
};
let metric = vector.metric_type();
covered = coverage(covered, index.fragment_bitmap.as_ref());
let fragments = index.fragment_bitmap.clone();
segments.push(IvfSegment {
index: vector.clone(),
fragments,
});
match &found {
None => found = Some((index.name.clone(), index.uuid, metric, centroids)),
Some((name, _, seen_metric, seen)) if *name == index.name => {
@@ -369,32 +303,20 @@ pub(super) async fn bind_column_with_segments(
// Lance's index path assigns cosine by L2 over unit vectors.
let normalize = metric == DistanceType::Cosine;
let distance = if normalize { DistanceType::L2 } else { metric };
let partitions = u32::try_from(centroids.len()).map_err(|_| Error::InvalidInput {
message: format!("{IVF_PARTITION}({column}): too many partitions"),
})?;
Ok((
IvfBinding {
index: uuid,
transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])),
normalize,
partitions,
},
IvfSegments { segments, covered },
))
Ok(IvfBinding {
index: uuid,
transformer: Arc::new(new_ivf_transformer(centroids, distance, vec![])),
normalize,
})
}
pub(super) fn empty_source(source_schema: &SchemaRef) -> Arc<dyn TableSource> {
let mut fields = source_schema.fields().to_vec();
fields.push(Arc::new(ArrowField::new(ROW_ID, DataType::UInt64, false)));
provider_as_source(Arc::new(EmptyTable::new(Arc::new(ArrowSchema::new(
fields,
)))))
provider_as_source(Arc::new(EmptyTable::new(source_schema.clone())))
}
/// The query DataFusion runs: the view's projections plus the group's
/// smallest source row id, which stands as the row's provenance. `filtered`
/// puts the view's predicate in the query, for a caller whose rows did not
/// come from a lance scan that already applied it.
/// The query DataFusion runs. `filtered` puts the view's predicate in the
/// query for a caller whose rows did not come from a Lance scan that already
/// applied it.
fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String {
let items: Vec<String> = definition
.projections
@@ -406,7 +328,7 @@ fn sql(definition: &MaterializedViewDefinition, filtered: bool) -> String {
_ => String::new(),
};
format!(
"SELECT {}, min({ROW_ID}) AS {ROW_ID} FROM {SOURCE}{predicate} GROUP BY {}",
"SELECT {} FROM {SOURCE}{predicate} GROUP BY {}",
items.join(", "),
definition.group_by.join(", ")
)
@@ -499,11 +421,6 @@ pub(super) fn plan(
};
let mut projections: Vec<ViewProjection> = Vec::with_capacity(definition.projections.len());
for p in &definition.projections {
if p.output == SOURCE_ROW_ID_COLUMN || p.output == ROW_ID {
return Err(Error::InvalidInput {
message: format!("view column name '{}' is reserved", p.output),
});
}
if projections.iter().any(|seen| seen.output == p.output) {
return Err(Error::ColumnAlreadyExists {
name: p.output.clone(),
@@ -549,14 +466,11 @@ pub(super) fn plan(
message: format!("invalid grouped view: {e}"),
})?;
for expr in &exprs {
ensure_immutable(expr, |message| Error::InvalidInput {
message: format!("invalid grouped view: {message}"),
})?;
inputs.extend(
expr.column_refs()
.into_iter()
.map(|c| c.name.clone())
.filter(|name| name != ROW_ID && source_schema.field_with_name(name).is_ok()),
.filter(|name| source_schema.field_with_name(name).is_ok()),
);
}
inputs.sort();
@@ -566,7 +480,6 @@ pub(super) fn plan(
.schema()
.fields()
.iter()
.filter(|f| f.name() != ROW_ID)
.map(|f| without_declarations(f))
.collect();
Ok(Planned {
@@ -586,7 +499,6 @@ pub(super) async fn stream(
rows_written: Arc<AtomicU64>,
) -> Result<SendableRecordBatchStream> {
let mut scanner = source.scan();
scanner.with_row_id();
if let Some(filter) = &definition.filter {
scanner.filter(filter)?;
}
@@ -611,24 +523,3 @@ pub(super) async fn stream(
});
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped)))
}
#[cfg(test)]
mod tests {
use super::*;
/// A segment that does not name the fragments it holds leaves the whole
/// index's coverage unknown: reading "everything it does not cover" would
/// read its rows a second time.
#[test]
fn one_segment_without_a_bitmap_loses_the_coverage() {
let first = RoaringBitmap::from_iter([0u32, 1]);
let second = RoaringBitmap::from_iter([2u32]);
let both = coverage(
coverage(Some(RoaringBitmap::new()), Some(&first)),
Some(&second),
);
assert_eq!(both, Some(RoaringBitmap::from_iter([0u32, 1, 2])));
assert_eq!(coverage(both, None), None);
assert_eq!(coverage(None, Some(&first)), None);
}
}
rust/lancedb/src/materialized_view/grouped_units.rs
-1368
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/materialized_view/partitioned.rs
-192
View File
@@ -1,192 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Shared plan/write/commit surface for materialized views with independent units.
use serde::{Deserialize, Serialize};
use super::duplicate_pairs::{
self, DuplicatePairsRefreshPlan, VectorSourceKind, WrittenPairsPartition,
};
use super::vector_dedup::{self, VectorDedupPlan, WrittenDedupUnit};
use super::{GroupedRefreshPlan, RefreshMaterializedViewResult, StoredDefinition, WrittenUnit};
use crate::table::Table;
use crate::{Error, Result};
/// Serializable work description consumed by a host's existing job workers.
/// The grouped variant preserves the previous saved-plan representation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PartitionedRefreshPlan {
VectorDedup(VectorDedupPlan),
DuplicatePairs(DuplicatePairsRefreshPlan),
Grouped(GroupedRefreshPlan),
}
impl PartitionedRefreshPlan {
/// The shared source snapshot and destination generation fence.
pub fn snapshot(&self) -> &GroupedRefreshPlan {
match self {
Self::VectorDedup(plan) => &plan.snapshot,
Self::DuplicatePairs(plan) => &plan.snapshot,
Self::Grouped(plan) => plan,
}
}
/// Selection tasks that must complete before remaining tasks can start.
pub fn dependency_units(&self) -> u32 {
match self {
Self::VectorDedup(plan) => plan.dependencies(),
_ => 0,
}
}
/// Number of independent units in this immutable plan.
pub fn units(&self) -> u32 {
self.snapshot().units
}
}
/// One completed unit's uncommitted fragments. Failed attempts produce no receipt.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum WrittenPartition {
VectorDedup(WrittenDedupUnit),
DuplicatePairs(WrittenPairsPartition),
Grouped(WrittenUnit),
}
impl WrittenPartition {
/// Number of pair or grouped rows staged by this unit.
pub fn rows(&self) -> u64 {
match self {
Self::VectorDedup(unit) => unit.rows,
Self::DuplicatePairs(unit) => unit.rows,
Self::Grouped(unit) => unit.rows(),
}
}
/// Uncommitted output files, retained only after task completion.
pub fn fragments(&self) -> &[lance_table::format::Fragment] {
match self {
Self::VectorDedup(unit) => &unit.fragments,
Self::DuplicatePairs(unit) => &unit.fragments,
Self::Grouped(unit) => unit.fragments(),
}
}
}
/// Plan native pair or indexed grouped work; return `None` for ordinary views.
/// `pinned`, when supplied, must match a native view's declared source version.
pub async fn plan_partitioned_refresh(
view: &Table,
pinned: Option<u64>,
) -> Result<Option<PartitionedRefreshPlan>> {
if let Some(StoredDefinition::Query(definition)) =
super::read_definition(view.schema().await?.metadata())?
&& definition.vector_source.is_some()
{
if definition
.vector_source
.as_ref()
.is_some_and(|s| s.kind == VectorSourceKind::Dedup)
{
return vector_dedup::plan_refresh(view, pinned)
.await
.map(|p| Some(PartitionedRefreshPlan::VectorDedup(p)));
}
return duplicate_pairs::plan_refresh(view, pinned)
.await
.map(|p| Some(PartitionedRefreshPlan::DuplicatePairs(p)));
}
super::plan_grouped_refresh(view, pinned)
.await
.map(|plan| plan.map(PartitionedRefreshPlan::Grouped))
}
/// Stream one unit into uncommitted fragments. Persist the returned receipt only
/// on success. A failed or interrupted call must be retried from the unit's start.
pub async fn write_refresh_partition(
view: &Table,
unit: u32,
plan: &PartitionedRefreshPlan,
) -> Result<WrittenPartition> {
write_refresh_partition_with_inputs(view, unit, plan, &[]).await
}
/// Write a unit after its durable selection dependencies have completed.
/// Supply exactly the first `plan.dependency_units()` receipts for a dependent
/// unit. A missing, duplicate or foreign receipt fails before output is staged.
pub async fn write_refresh_partition_with_inputs(
view: &Table,
unit: u32,
plan: &PartitionedRefreshPlan,
dependencies: &[WrittenPartition],
) -> Result<WrittenPartition> {
match plan {
PartitionedRefreshPlan::VectorDedup(plan) => {
let dependencies = dependencies
.iter()
.map(|unit| match unit {
WrittenPartition::VectorDedup(unit) => Ok(unit.clone()),
_ => Err(Error::InvalidInput {
message: "foreign dedup dependency".into(),
}),
})
.collect::<Result<Vec<_>>>()?;
vector_dedup::write_unit(view, unit, plan, &dependencies)
.await
.map(WrittenPartition::VectorDedup)
}
PartitionedRefreshPlan::DuplicatePairs(plan) => {
duplicate_pairs::write_unit(view, unit, plan)
.await
.map(WrittenPartition::DuplicatePairs)
}
PartitionedRefreshPlan::Grouped(plan) => super::write_grouped_unit(view, unit, plan)
.await
.map(WrittenPartition::Grouped),
}
}
/// Publish a complete set of unit receipts through the common MV commit sink.
/// `expected` fences the destination incarnation. No partial set is publishable.
pub async fn commit_partitioned_refresh(
view: &Table,
plan: &PartitionedRefreshPlan,
units: Vec<WrittenPartition>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let foreign = || Error::InvalidInput {
message: "refresh receipt does not match the plan kind".into(),
};
match plan {
PartitionedRefreshPlan::VectorDedup(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::VectorDedup(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
vector_dedup::commit(view, plan, units, expected).await
}
PartitionedRefreshPlan::DuplicatePairs(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::DuplicatePairs(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
duplicate_pairs::commit(view, plan, units, expected).await
}
PartitionedRefreshPlan::Grouped(plan) => {
let units = units
.into_iter()
.map(|unit| match unit {
WrittenPartition::Grouped(unit) => Ok(unit),
_ => Err(foreign()),
})
.collect::<Result<Vec<_>>>()?;
super::commit_grouped_refresh(view, plan, units, expected).await
}
}
}
rust/lancedb/src/materialized_view/query.rs
+1 -84
View File
@@ -264,36 +264,7 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let mut from = select.from.iter();
let mut vector_source = None;
let (source_namespace, source_table) = match from.next().map(|f| &f.relation) {
Some(TableFactor::Table {
name,
args: Some(args),
alias: None,
with_hints,
version: None,
with_ordinality: false,
partitions,
json_path: None,
sample: None,
index_hints,
..
}) if matches!(
name.to_string().as_str(),
super::duplicate_pairs::FUNCTION_NAME | super::duplicate_pairs::DEDUP_FUNCTION_NAME
) && with_hints.is_empty()
&& partitions.is_empty()
&& index_hints.is_empty() =>
{
let kind = if name.to_string() == super::duplicate_pairs::FUNCTION_NAME {
super::VectorSourceKind::Pairs
} else {
super::VectorSourceKind::Dedup
};
let (namespace, table, config) = super::duplicate_pairs::source(args, kind)?;
vector_source = Some(Box::new(config));
(namespace, table)
}
Some(TableFactor::Table { args: Some(_), .. }) => {
return Err(invalid(
"a view reads a table; a Function in FROM position follows it: \
@@ -445,7 +416,6 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
};
let definition = MaterializedViewDefinition {
vector_source,
source_table,
source_namespace,
lateral,
@@ -454,9 +424,6 @@ fn extract(query: &Query) -> Result<MaterializedViewDefinition> {
group_by,
limit,
};
if definition.vector_source.is_some() {
super::duplicate_pairs::check_shape(&definition)?;
}
check_grouping(&definition)?;
Ok(definition)
}
@@ -516,18 +483,7 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
.chain(std::iter::once(&definition.source_table))
.map(|part| ident_sql(part))
.collect();
if let Some(native) = &definition.vector_source {
sql.push_str(&format!(
"{}('{}', {}, '{}', {})",
native.function_name(),
table.join(".").replace('\'', "''"),
native.dataset_version,
native.column.replace('\'', "''"),
native.distance_threshold
));
} else {
sql.push_str(&table.join("."));
}
sql.push_str(&table.join("."));
if let Some(lateral) = &definition.lateral {
match &lateral.source {
LateralSource::Unnest { column } => {
@@ -555,45 +511,6 @@ pub fn render(definition: &MaterializedViewDefinition) -> String {
mod tests {
use super::*;
#[test]
fn native_pairs_are_a_table_source_with_literal_snapshot_arguments() {
for threshold in ["4", "4.0", "-1", "0.05"] {
let sql =
format!("SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', {threshold})");
let definition = parse(&sql).unwrap();
assert_eq!(definition.source_table, "images");
assert!(definition.lateral.is_none());
assert_eq!(
definition.vector_source.as_ref().unwrap().dataset_version,
3
);
assert_eq!(parse(&render(&definition)).unwrap(), definition);
let metadata = std::collections::HashMap::from([(
super::super::DEFINITION_META_KEY.to_string(),
definition.to_json().unwrap(),
)]);
assert_eq!(
super::super::read_definition(&metadata).unwrap(),
Some(super::super::StoredDefinition::Query(definition))
);
}
for sql in [
"SELECT row_id_a FROM vector_duplicate_pairs('images', 3, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 0, 'phash', 4)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 1e100)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', cutoff)",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) LIMIT 1",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) WHERE distance < 2",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4) AS p",
"SELECT * FROM vector_duplicate_pairs('images', 3, 'phash', 4), f(phash) AS p",
] {
assert!(
parse(sql).is_err(),
"accepted unsupported native view: {sql}"
);
}
}
#[test]
fn a_query_round_trips_through_its_canonical_form() {
for (sql, canonical) in [
rust/lancedb/src/materialized_view/refresh.rs
+294 -4449
View File
File diff suppressed because it is too large Load Diff
rust/lancedb/src/materialized_view/vector_dedup.rs
-636
View File
@@ -1,636 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! Two dependent waves in the ordinary refresh fleet: index partition selection,
//! then source fragment materialization. Only the latter publishes data files.
use std::collections::BTreeMap;
use std::io::Cursor;
use std::sync::Arc;
use arrow_array::{Array, RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use datafusion::prelude::{SessionConfig, SessionContext, col};
use datafusion_common::DataFusionError;
use datafusion_execution::{
SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder,
};
use datafusion_physical_plan::{ExecutionPlan, stream::RecordBatchStreamAdapter};
use futures::TryStreamExt;
use lance::Dataset;
use lance::dataset::rowids::{get_row_id_index, load_row_id_sequence};
use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap};
use lance::dataset::transaction::Operation;
use lance::dataset::{InsertBuilder, WriteDestination, WriteMode, WriteParams};
use lance_core::datatypes::BlobHandling;
use lance_core::utils::address::RowAddress;
use lance_datafusion::exec::{OneShotExec, SessionContextExt};
use lance_table::format::Fragment;
use roaring::{RoaringBitmap, RoaringTreemap};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;
use super::duplicate_pairs::{self, DuplicatePairsRefreshPlan, VectorSourceKind};
use super::{GroupedRefreshPlan, RefreshMaterializedViewResult};
use crate::table::Table;
use crate::table::datafusion::udtf::duplicate_pairs::{DuplicatePairsExec, plan_duplicate_pairs};
use crate::{Error, Result};
const POLICY_VERSION: u32 = 1;
const SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MATERIALIZE_BATCH_ROWS: usize = 1024;
const MASK_FILE: &str = "dedup.bitmap";
const MASK_MAGIC: &[u8] = b"LANCEDB_DEDUP_1";
fn invalid(message: impl Into<String>) -> Error {
Error::InvalidInput {
message: message.into(),
}
}
fn io_error(error: std::io::Error) -> Error {
invalid(format!("invalid dedup bitmap: {error}"))
}
/// Immutable snapshot, partition tasks and source fragments for both waves.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct VectorDedupPlan {
pub snapshot: GroupedRefreshPlan,
native: DuplicatePairsRefreshPlan,
source_fragments: Vec<u32>,
policy_version: u32,
}
impl VectorDedupPlan {
pub(super) fn dependencies(&self) -> u32 {
self.native.snapshot.units
}
fn identity(&self) -> Result<String> {
Ok(format!(
"{:x}",
Sha256::digest(serde_json::to_vec(self).map_err(|e| invalid(e.to_string()))?)
))
}
}
/// Small receipt for a durable selection mask. Row IDs and pairs are not stored
/// in the host's task registry. UUID directories use ordinary index-file GC.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct MaskFile {
source_fragment: u32,
object_id: Uuid,
bytes: usize,
digest: String,
}
/// A selection task produces masks; a materialization task produces data files.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct WrittenDedupUnit {
unit: u32,
plan_id: String,
masks: Vec<MaskFile>,
pub(super) fragments: Vec<Fragment>,
pub(super) rows: u64,
}
pub(super) async fn plan_refresh(view: &Table, pinned: Option<u64>) -> Result<VectorDedupPlan> {
let native = duplicate_pairs::plan_refresh(view, pinned).await?;
let (_, _, source) = duplicate_pairs::open_planned(view, &native).await?;
let source_fragments = source
.get_fragments()
.iter()
.map(|f| f.id() as u32)
.collect::<Vec<_>>();
let mut snapshot = native.snapshot.clone();
snapshot.units = snapshot
.units
.checked_add(
u32::try_from(source_fragments.len())
.map_err(|_| invalid("too many source fragments"))?,
)
.ok_or_else(|| invalid("too many dedup refresh units"))?;
Ok(VectorDedupPlan {
snapshot,
native,
source_fragments,
policy_version: POLICY_VERSION,
})
}
async fn open(
view: &Table,
plan: &VectorDedupPlan,
) -> Result<(Dataset, super::MaterializedViewDefinition, Arc<Dataset>)> {
let (view_ds, definition, source) = duplicate_pairs::open_planned(view, &plan.native).await?;
let mut snapshot = plan.native.snapshot.clone();
snapshot.units = snapshot
.units
.checked_add(
u32::try_from(plan.source_fragments.len())
.map_err(|_| invalid("too many dedup fragments"))?,
)
.ok_or_else(|| invalid("too many dedup units"))?;
if plan.policy_version != POLICY_VERSION
|| snapshot != plan.snapshot
|| definition.vector_source.as_ref().map(|s| s.kind) != Some(VectorSourceKind::Dedup)
|| source
.get_fragments()
.iter()
.map(|f| f.id() as u32)
.collect::<Vec<_>>()
!= plan.source_fragments
{
return Err(invalid("dedup task belongs to another snapshot or policy"));
}
Ok((view_ds, definition, source))
}
fn edge_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt64, false),
Field::new("b", DataType::UInt64, false),
]))
}
/// Sorting canonical edges makes the first unremoved endpoint the smallest
/// unassigned representative. The sort spills; only vertex masks stay in RAM.
pub async fn select_representatives(
stream: SendableRecordBatchStream,
memory: usize,
) -> Result<(RoaringTreemap, RoaringTreemap)> {
let schema = edge_schema();
let batch_schema = schema.clone();
let canonical = stream
.map_ok(move |batch| {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
RecordBatch::try_new(
batch_schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)),
)),
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)),
)),
],
)
.map_err(DataFusionError::from)
})
.and_then(futures::future::ready);
let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical));
let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS);
// Leave enough room to merge spilled runs while the final in-memory run
// is still resident. Small batches also bound each merge input buffer.
config.options_mut().execution.sort_spill_reservation_bytes = memory / 2;
let ctx = SessionContext::new_with_config_rt(
config,
RuntimeEnvBuilder::new()
.with_memory_limit(memory, 1.0)
.build_arc()?,
);
let mut sorted = ctx
.read_one_shot(canonical)?
.sort_by(vec![col("a"), col("b")])?
.execute_stream()
.await?;
let mut removed = RoaringTreemap::new();
let mut representatives = RoaringTreemap::new();
while let Some(batch) = sorted.try_next().await? {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
for (&a, &b) in a.values().iter().zip(b.values()) {
if a != b && !removed.contains(a) && !removed.contains(b) {
representatives.insert(a);
removed.insert(b);
}
}
}
Ok((removed, representatives))
}
#[derive(Default)]
struct Masks {
removed: RoaringBitmap,
representatives: RoaringBitmap,
}
async fn persist_masks(
view: &Dataset,
source: &Dataset,
removed: RoaringTreemap,
representatives: RoaringTreemap,
) -> Result<Vec<MaskFile>> {
let row_index = get_row_id_index(source).await?;
let mut masks: BTreeMap<u32, Masks> = BTreeMap::new();
for (ids, is_removed) in [(removed, true), (representatives, false)] {
for id in ids {
let address = match &row_index {
Some(index) => index
.get(id)?
.ok_or_else(|| invalid("pair row id missing from pinned source"))?,
None => RowAddress::new_from_u64(id),
};
let mask = masks.entry(address.fragment_id()).or_default();
if is_removed {
mask.removed.insert(address.row_offset());
} else {
mask.representatives.insert(address.row_offset());
}
}
}
let store = view.object_store(None).await?;
let mut files = Vec::new();
for (source_fragment, mask) in masks {
if source.get_fragment(source_fragment as usize).is_none() {
return Err(invalid("pair row address is outside the pinned snapshot"));
}
let mut bytes = MASK_MAGIC.to_vec();
mask.removed.serialize_into(&mut bytes).map_err(io_error)?;
mask.representatives
.serialize_into(&mut bytes)
.map_err(io_error)?;
let object_id = Uuid::new_v4();
store
.put(
&view
.indices_dir()
.join(object_id.to_string())
.join(MASK_FILE),
&bytes,
)
.await?;
files.push(MaskFile {
source_fragment,
object_id,
bytes: bytes.len(),
digest: format!("{:x}", Sha256::digest(&bytes)),
});
}
Ok(files)
}
fn validate_receipts<'a>(
plan: &VectorDedupPlan,
receipts: &'a [WrittenDedupUnit],
count: u32,
) -> Result<Vec<&'a WrittenDedupUnit>> {
let identity = plan.identity()?;
let mut ordered = vec![None; count as usize];
for receipt in receipts {
if receipt.plan_id != identity {
return Err(invalid("dedup receipt belongs to another plan"));
}
let slot = ordered
.get_mut(receipt.unit as usize)
.ok_or_else(|| invalid("unexpected dedup receipt"))?;
if slot.is_some() {
return Err(invalid("duplicate dedup receipt"));
}
if receipt.unit < plan.dependencies() {
if !receipt.fragments.is_empty()
|| receipt.rows != 0
|| receipt
.masks
.iter()
.any(|m| !plan.source_fragments.contains(&m.source_fragment))
{
return Err(invalid("invalid selection receipt"));
}
} else if !receipt.masks.is_empty() {
return Err(invalid("invalid materialization receipt"));
}
*slot = Some(receipt);
}
ordered
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or_else(|| invalid("missing dedup dependency receipt"))
}
async fn load_masks(
view: &Dataset,
source_fragment: u32,
dependencies: &[&WrittenDedupUnit],
) -> Result<RoaringBitmap> {
let store = view.object_store(None).await?;
let mut masks = Masks::default();
for file in dependencies
.iter()
.flat_map(|u| &u.masks)
.filter(|m| m.source_fragment == source_fragment)
{
let reader = store
.open(
&view
.indices_dir()
.join(file.object_id.to_string())
.join(MASK_FILE),
)
.await?;
if reader.size().await? != file.bytes {
return Err(invalid("dedup mask length mismatch"));
}
let bytes = reader.get_range(0..file.bytes).await?;
if !bytes.starts_with(MASK_MAGIC) || format!("{:x}", Sha256::digest(&bytes)) != file.digest
{
return Err(invalid("dedup mask checksum mismatch"));
}
let mut cursor = Cursor::new(&bytes[MASK_MAGIC.len()..]);
masks.removed |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?;
masks.representatives |= RoaringBitmap::deserialize_from(&mut cursor).map_err(io_error)?;
if cursor.position() as usize != bytes.len() - MASK_MAGIC.len() {
return Err(invalid("dedup mask has trailing data"));
}
}
// Native index scopes normally have disjoint rows. Fail closed if an
// overlapping scope would delete a representative used by another scope.
if !masks.removed.is_disjoint(&masks.representatives) {
return Err(invalid(
"overlapping index scopes disagree on retained representatives",
));
}
Ok(masks.removed)
}
pub(super) async fn write_unit(
view: &Table,
unit: u32,
plan: &VectorDedupPlan,
dependencies: &[WrittenDedupUnit],
) -> Result<WrittenDedupUnit> {
let (view_ds, definition, source) = open(view, plan).await?;
let mut written = WrittenDedupUnit {
unit,
plan_id: plan.identity()?,
masks: Vec::new(),
fragments: Vec::new(),
rows: 0,
};
if unit < plan.dependencies() {
let config = definition
.vector_source
.as_ref()
.expect("dedup source")
.config()?;
let exec = DuplicatePairsExec::try_new(
source.clone(),
config,
vec![plan.native.tasks[unit as usize].clone()],
)?;
let stream = exec.execute(0, Arc::new(TaskContext::default()))?;
let (removed, representatives) = select_representatives(stream, SORT_MEMORY_BYTES).await?;
written.masks = persist_masks(&view_ds, &source, removed, representatives).await?;
return Ok(written);
}
let source_fragment = *plan
.source_fragments
.get((unit - plan.dependencies()) as usize)
.ok_or_else(|| invalid("dedup unit is outside the plan"))?;
let dependencies = validate_receipts(plan, dependencies, plan.dependencies())?;
let removed = load_masks(&view_ds, source_fragment, &dependencies).await?;
let fragment = source
.get_fragment(source_fragment as usize)
.ok_or_else(|| invalid("source fragment missing"))?;
// The scanner prefilter uses snapshot row IDs. Our persisted masks use
// physical offsets so each writer only loads its own fragment's artifacts.
let blocked = if source.manifest.uses_stable_row_ids() {
let sequence = load_row_id_sequence(&source, fragment.metadata()).await?;
if removed
.max()
.is_some_and(|offset| offset as u64 >= sequence.len())
{
return Err(invalid("dedup mask offset is outside the source fragment"));
}
RowAddrTreeMap::from_iter(sequence.select(removed.iter().map(|offset| offset as usize)))
} else {
RowAddrTreeMap::from_iter(
removed
.iter()
.map(|offset| u64::from(RowAddress::new_from_parts(source_fragment, offset))),
)
};
let mut scanner = source.scan();
scanner
.with_fragments(vec![fragment.metadata().clone()])
.with_row_addr_prefilter(RowAddrMask::from_block(blocked))
.blob_handling(BlobHandling::AllBinary)
.batch_size(MATERIALIZE_BATCH_ROWS)
.strict_batch_size(true);
let stream = scanner.try_into_stream().await?;
let schema = lance_io::stream::RecordBatchStream::schema(&stream);
let rows = Arc::new(std::sync::atomic::AtomicU64::new(0));
let count = rows.clone();
let stream = stream
.map_err(|e| DataFusionError::External(Box::new(e)))
.inspect_ok(move |batch| {
count.fetch_add(
batch.num_rows() as u64,
std::sync::atomic::Ordering::Relaxed,
);
});
let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new(schema, stream));
// Materialized binary payloads must become the destination's logical blob
// structs before the Lance writer encodes them. Reuse the normal add path's
// coercions, including nulls and nested fields.
let coerced = crate::table::datafusion::cast::cast_to_table_schema(
Arc::new(OneShotExec::new(stream)),
&Schema::from(view_ds.schema()),
)?;
let stream = coerced.execute(0, Arc::new(TaskContext::default()))?;
let txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds)))
.with_params(&WriteParams {
mode: WriteMode::Append,
..Default::default()
})
.execute_uncommitted_stream(stream)
.await?;
let Operation::Append { fragments } = txn.operation else {
return Err(invalid(
"dedup materialization did not stage append fragments",
));
};
written.fragments = fragments;
written.rows = rows.load(std::sync::atomic::Ordering::Relaxed);
Ok(written)
}
pub(super) async fn commit(
view: &Table,
plan: &VectorDedupPlan,
units: Vec<WrittenDedupUnit>,
expected: Option<&str>,
) -> Result<RefreshMaterializedViewResult> {
let (_, definition, source) = open(view, plan).await?;
let config = definition
.vector_source
.as_ref()
.expect("dedup source")
.config()?;
if plan_duplicate_pairs(source.clone(), &config).await? != plan.native.tasks {
return Err(invalid("dedup plan does not cover the pinned index"));
}
validate_receipts(plan, &units, plan.snapshot.units)?;
let rows = units.iter().map(|u| u.rows).sum();
let fragments = units.into_iter().flat_map(|u| u.fragments).collect();
super::grouped_units::commit_fragments(
view,
&plan.snapshot,
&definition,
source.manifest.timestamp_nanos,
fragments,
rows,
expected,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
fn edges(values: &[(u64, u64)], batch_size: usize) -> SendableRecordBatchStream {
let schema = edge_schema();
let batches = values
.chunks(batch_size)
.map(|edges| {
Ok(RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.0))),
Arc::new(UInt64Array::from_iter_values(edges.iter().map(|e| e.1))),
],
)
.unwrap())
})
.collect::<Vec<_>>();
Box::pin(RecordBatchStreamAdapter::new(
schema,
futures::stream::iter(batches),
))
}
#[tokio::test]
async fn direct_representatives_handle_chains_orientation_and_uint64() {
let high = u64::MAX - 2;
let pairs = [
(2, 1),
(1, 0),
(high + 2, high + 1),
(high, high + 1),
(1, 0),
];
let (removed, representatives) =
select_representatives(edges(&pairs, 1), SORT_MEMORY_BYTES)
.await
.unwrap();
assert_eq!(removed.iter().collect::<Vec<_>>(), [1, high + 1]);
assert_eq!(representatives.iter().collect::<Vec<_>>(), [0, high]);
}
#[tokio::test]
async fn masks_reject_overlapping_representatives_and_corrupt_artifacts() {
let dir = tempfile::tempdir().unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::UInt64, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(UInt64Array::from(vec![0, 1, 2]))],
)
.unwrap();
let ds = Dataset::write(
arrow_array::RecordBatchIterator::new(vec![Ok(batch)], schema),
dir.path().join("source.lance").to_str().unwrap(),
None,
)
.await
.unwrap();
let first = persist_masks(
&ds,
&ds,
[1].into_iter().collect(),
[0].into_iter().collect(),
)
.await
.unwrap();
let second = persist_masks(
&ds,
&ds,
[2].into_iter().collect(),
[1].into_iter().collect(),
)
.await
.unwrap();
let receipt = |masks| WrittenDedupUnit {
unit: 0,
plan_id: "test".into(),
masks,
fragments: vec![],
rows: 0,
};
let first = receipt(first);
let second = receipt(second);
assert_eq!(
load_masks(&ds, 0, &[&first]).await.unwrap(),
[1].into_iter().collect()
);
assert!(
load_masks(&ds, 0, &[&first, &second])
.await
.unwrap_err()
.to_string()
.contains("overlapping")
);
let file = &first.masks[0];
ds.object_store(None)
.await
.unwrap()
.put(
&ds.indices_dir()
.join(file.object_id.to_string())
.join(MASK_FILE),
&vec![0; file.bytes],
)
.await
.unwrap();
assert!(
load_masks(&ds, 0, &[&first])
.await
.unwrap_err()
.to_string()
.contains("checksum")
);
}
#[tokio::test]
async fn external_sort_handles_more_edges_than_its_budget() {
let pairs = (0..262144)
.rev()
.map(|i| ((i + 1) as u64, i as u64))
.collect::<Vec<_>>();
let (removed, reps) = select_representatives(edges(&pairs, 4096), 2 * 1024 * 1024)
.await
.unwrap();
assert_eq!(removed.len(), 131072);
assert_eq!(reps.len(), 131072);
assert!(removed.iter().all(|id| id % 2 == 1));
assert!(reps.iter().all(|id| id % 2 == 0));
}
}
rust/lancedb/src/remote/db.rs
+30 -60
View File
@@ -47,6 +47,10 @@ use super::table::RemoteTable;
use super::util::parse_server_version;
use super::{ARROW_STREAM_CONTENT_TYPE, extract_job_id};
fn quote_sql_identifier(identifier: &str) -> String {
format!("\"{}\"", identifier.replace('"', "\"\""))
}
// Request structure for the remote clone table API
#[derive(serde::Serialize)]
struct RemoteCloneTableRequest {
@@ -860,42 +864,22 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
&self,
request: CreateMaterializedViewRequest,
) -> Result<Job> {
let identifier = build_table_identifier(&request.name, &request.namespace_path)?;
let req = self
.client
.post(&format!("/v1/materialized_view/{identifier}/create"))
.json(&serde_json::json!({
"query": request.query,
"with_no_data": request.with_no_data,
}));
let (request_id, response) = self.client.send(req).await?;
let response = self.client.check_response(&request_id, response).await?;
let status = response.status();
let body = response.text().await.err_to_http(request_id.clone())?;
let job_id = extract_job_id(&body);
if request.with_no_data {
return Ok(match job_id {
Some(job_id) => Job::new(Box::new(RemoteJob::new(self.client.clone(), job_id))),
None => Job::new_done(),
return Err(Error::NotSupported {
message: "remote materialized views are always populated by their SQL job"
.to_string(),
});
}
if status != StatusCode::ACCEPTED {
return Err(Error::Http {
source: "materialized-view creation with data must return 202 Accepted".into(),
request_id,
status_code: Some(status),
});
}
let job_id = job_id.ok_or_else(|| Error::Http {
source: "materialized-view creation response did not contain a valid job_id".into(),
request_id,
status_code: Some(status),
let client = self.sql_client.clone().ok_or_else(|| Error::NotSupported {
message: "SQL is unavailable for this remote database client".to_string(),
})?;
Ok(Job::new(Box::new(RemoteJob::new(
self.client.clone(),
job_id,
))))
let namespace = request.namespace_path;
let statement = format!(
"CREATE MATERIALIZED VIEW {} AS {}",
quote_sql_identifier(&request.name),
request.query
);
Ok(client.submit_as_job(statement, namespace, || async { Ok(()) }))
}
async fn drop_materialized_view_async(
@@ -1547,12 +1531,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let version = parse_server_version(&request_id, &rsp)?;
let table_identifier = build_table_identifier(&request.name, &request.namespace_path)?;
let cache_key = build_cache_key(&request.name, &request.namespace_path);
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name.clone(),
request.namespace_path.clone(),
table_identifier,
version.clone(),
self.sql_client.clone(),
));
self.table_cache.insert(cache_key, version).await;
@@ -1589,12 +1574,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let version = parse_server_version(&request_id, &rsp)?;
let cache_key = build_cache_key(&request.target_table_name, &request.target_namespace_path);
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.target_table_name.clone(),
request.target_namespace_path.clone(),
table_identifier,
version.clone(),
self.sql_client.clone(),
));
self.table_cache.insert(cache_key, version).await;
@@ -1607,12 +1593,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
// Every open gets its own checkout, schema cache, and freshness state.
if let Some(version) = self.table_cache.get(&cache_key).await {
Ok(Arc::new(RemoteTable::new(
Ok(Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name,
request.namespace_path,
identifier,
version,
self.sql_client.clone(),
)))
} else {
// Describe the table to confirm it exists before moving on.
@@ -1625,12 +1612,13 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let rsp = self.client.check_response(&request_id, rsp).await?;
let version = parse_server_version(&request_id, &rsp)?;
let describe_body = rsp.text().await.ok();
let table = Arc::new(RemoteTable::new(
let table = Arc::new(RemoteTable::new_with_sql_client(
self.client.clone(),
request.name.clone(),
request.namespace_path.clone(),
identifier,
version.clone(),
self.sql_client.clone(),
));
// This describe already carries the schema, so hand it to the table
// instead of making the first schema read fetch it again. A version or
@@ -1957,29 +1945,11 @@ mod tests {
}
#[tokio::test]
async fn test_create_materialized_view_uses_item_route_and_job() {
let db = super::RemoteDatabase::new_mock(|request| {
assert_eq!(request.method(), "POST");
assert_eq!(
request.url().path(),
"/v1/materialized_view/analytics$adults/create"
);
let body = request
.body()
.and_then(reqwest::Body::as_bytes)
.and_then(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).ok())
.unwrap();
assert_eq!(
body["query"],
"SELECT age AS \"age\" FROM \"raw\".\"people\" WHERE age >= 18 LIMIT 10"
);
assert_eq!(body["with_no_data"], false);
http::Response::builder()
.status(202)
.body(serde_json::json!({"job_id": "j1-mv-create"}).to_string())
.unwrap()
async fn test_create_materialized_view_requires_sql_client() {
let db = super::RemoteDatabase::new_mock(|request| -> http::Response<String> {
panic!("unexpected REST request: {}", request.url().path())
});
let job = db
let error = db
.create_materialized_view_async(CreateMaterializedViewRequest {
name: "adults".into(),
namespace_path: vec!["analytics".into()],
@@ -1988,8 +1958,8 @@ mod tests {
with_no_data: false,
})
.await
.unwrap();
assert_eq!(job.id(), Some("j1-mv-create"));
.unwrap_err();
assert!(error.to_string().contains("SQL is unavailable"));
}
#[tokio::test]
rust/lancedb/src/remote/sql.rs
+40 -1
View File
@@ -3,6 +3,7 @@
use std::collections::HashMap;
use std::fs;
use std::future::{Future, pending};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex as StdMutex, OnceLock};
use std::time::{Duration, Instant};
@@ -20,12 +21,14 @@ use arrow_schema::{Schema, SchemaRef};
use futures::TryStreamExt;
use http::header::{HeaderMap, HeaderName, HeaderValue};
use prost::Message;
use tokio::sync::{Mutex, Notify, OnceCell, mpsc};
use serde::{Serialize, de::DeserializeOwned};
use tokio::sync::{Mutex, Notify, OnceCell, mpsc, watch};
use tonic::transport::{Certificate, Channel, ClientTlsConfig, Endpoint, Identity};
use uuid::Uuid;
use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
use crate::error::{Error, Result};
use crate::job::Job;
use crate::remote::client::{ClientConfig, TlsConfig};
use crate::remote::retry::ResolvedRetryConfig;
use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus};
@@ -209,6 +212,42 @@ impl SqlClient {
.await
}
pub(super) fn submit_as_job<T, F, Fut>(
&self,
statement: String,
namespace: Vec<String>,
finish: F,
) -> Job<T>
where
T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static,
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<T>> + Send + 'static,
{
let client = self.clone();
let (query_tx, query_rx) = watch::channel(None::<Arc<Query>>);
let task = tokio::spawn(async move {
let query = Arc::new(client.submit(&statement, &namespace).await?);
query_tx.send_replace(Some(query.clone()));
let mut reader = query.reader().await?;
while reader.try_next().await?.is_some() {}
finish().await
});
Job::spawned_with_cancellation(task, move || {
let mut query_rx = query_rx.clone();
async move {
loop {
let query = query_rx.borrow().clone();
if let Some(query) = query {
return query.cancel().await;
}
if query_rx.changed().await.is_err() {
pending::<()>().await;
}
}
}
})
}
pub(super) async fn describe(&self, query_id: Uuid) -> Result<QueryDescription> {
let query = self
.queries
rust/lancedb/src/remote/sql_test.rs
+151
View File
@@ -21,7 +21,10 @@ use super::*;
use crate::database::Database;
use crate::remote::RemoteCatalogOptions;
use crate::remote::client::HeaderProvider;
use crate::remote::client::test_utils::client_with_handler;
use crate::remote::db::RemoteDatabase;
use crate::remote::table::RemoteTable;
use crate::table::BaseTable;
#[derive(Debug, Default)]
struct DelayedHeaderProvider {
@@ -73,6 +76,7 @@ struct TestSqlService {
first_continuation_count: Arc<AtomicUsize>,
transient_poll_failures: Arc<AtomicUsize>,
headers: Arc<std::sync::Mutex<Vec<CapturedHeaders>>>,
gate_refresh_release: Arc<Notify>,
result: RecordBatch,
large_result: RecordBatch,
dictionary_result: RecordBatch,
@@ -121,6 +125,7 @@ impl Default for TestSqlService {
first_continuation_count: Arc::new(AtomicUsize::new(0)),
transient_poll_failures: Arc::new(AtomicUsize::new(0)),
headers: Arc::new(std::sync::Mutex::new(Vec::new())),
gate_refresh_release: Arc::new(Notify::new()),
result,
large_result,
dictionary_result,
@@ -279,6 +284,7 @@ impl FlightService for TestSqlService {
) -> std::result::Result<Response<<Self as FlightService>::DoGetStream>, Status> {
self.do_get_count.fetch_add(1, Ordering::SeqCst);
let ticket = request.get_ref().ticket.as_ref();
let gate_refresh = ticket == b"REFRESH MATERIALIZED VIEW \"cancel_view\"";
let empty = ticket == b"SELECT empty";
let slow = ticket == b"SELECT slow get";
let large = ticket == b"SELECT large message";
@@ -290,7 +296,11 @@ impl FlightService for TestSqlService {
self.result.clone()
};
let schema = result.schema();
let gate_refresh_release = self.gate_refresh_release.clone();
let input = futures::stream::once(async move {
if gate_refresh {
gate_refresh_release.notified().await;
}
if slow {
tokio::time::sleep(Duration::from_millis(250)).await;
}
@@ -446,6 +456,147 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() {
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn remote_refresh_twice_keeps_latest_selector() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let incoming = futures::stream::try_unfold(listener, |listener| async {
let (socket, _) = listener.accept().await?;
Ok::<_, std::io::Error>(Some((socket, listener)))
});
let service = TestSqlService::default();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(FlightServiceServer::new(service))
.serve_with_incoming_shutdown(incoming, async {
let _ = shutdown_rx.await;
}),
);
let sql_client = SqlClient::new(
"analytics".into(),
None,
"test-key".into(),
None,
Some(format!("grpc://{address}")),
ClientConfig::default(),
);
let selectors = Arc::new(std::sync::Mutex::new(Vec::new()));
let captured = selectors.clone();
let client = client_with_handler(move |request| {
let body = request
.body()
.and_then(reqwest::Body::as_bytes)
.and_then(|bytes| serde_json::from_slice::<serde_json::Value>(bytes).ok())
.unwrap();
captured.lock().unwrap().push(body["version"].clone());
let text = match request.url().path() {
"/v1/table/ns$view/describe/" => r#"{"version": 42, "schema": {"fields": []}}"#,
"/v1/table/ns$view/count_rows/" => "2",
other => panic!("unexpected REST request: {other}"),
};
http::Response::builder()
.status(200)
.body(text.to_string())
.unwrap()
});
let table = RemoteTable::new_with_sql_client(
client,
"view".into(),
vec!["ns".into()],
"ns$view".into(),
Default::default(),
Some(sql_client),
);
let first = table
.refresh_materialized_view_async(None)
.await
.unwrap()
.wait()
.await
.unwrap();
let second = table
.refresh_materialized_view_async(None)
.await
.unwrap()
.wait()
.await
.unwrap();
assert_eq!(first.version, 42);
assert_eq!(second.version, 42);
assert!(
selectors
.lock()
.unwrap()
.iter()
.all(serde_json::Value::is_null)
);
shutdown_tx.send(()).unwrap();
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn cancelling_remote_refresh_cancels_sql_query() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let incoming = futures::stream::try_unfold(listener, |listener| async {
let (socket, _) = listener.accept().await?;
Ok::<_, std::io::Error>(Some((socket, listener)))
});
let service = TestSqlService::default();
let cancel_count = service.cancel_count.clone();
let do_get_count = service.do_get_count.clone();
let gate_refresh_release = service.gate_refresh_release.clone();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(FlightServiceServer::new(service))
.serve_with_incoming_shutdown(incoming, async {
let _ = shutdown_rx.await;
}),
);
let sql_client = SqlClient::new(
"analytics".into(),
None,
"test-key".into(),
None,
Some(format!("grpc://{address}")),
ClientConfig::default(),
);
let client = client_with_handler(|request| -> http::Response<String> {
panic!(
"cancelled refresh must not read result metadata: {}",
request.url().path()
);
});
let table = RemoteTable::new_with_sql_client(
client,
"cancel_view".into(),
vec![],
"cancel_view".into(),
Default::default(),
Some(sql_client),
);
let job = table.refresh_materialized_view_async(None).await.unwrap();
tokio::time::timeout(Duration::from_secs(10), async {
while do_get_count.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
job.cancel().await.unwrap();
assert!(matches!(job.wait().await, Err(Error::JobCancelled { .. })));
assert_eq!(cancel_count.load(Ordering::SeqCst), 1);
gate_refresh_release.notify_one();
shutdown_tx.send(()).unwrap();
server.await.unwrap().unwrap();
}
#[tokio::test]
async fn submits_polls_fetches_cancels_and_reuses_client() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
rust/lancedb/src/remote/table.rs
+69 -147
View File
@@ -8,6 +8,7 @@ use self::insert::{RemoteWriteExec, WriteOp};
use super::client::RequestResultExt;
use super::client::{HttpSend, RestfulLanceDbClient, Sender};
use super::db::ServerVersion;
use super::sql::SqlClient;
use super::{ARROW_FILE_CONTENT_TYPE, ARROW_STREAM_CONTENT_TYPE, extract_job_id};
use crate::blob::BlobFile;
use crate::data::scannable::{PeekedScannable, Scannable, estimate_write_partitions};
@@ -18,7 +19,7 @@ use crate::index::scalar::FtsQuery;
use crate::index::waiter::wait_for_index;
use crate::job::Job;
use crate::materialized_view::{
MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, ViewProjection,
MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult,
};
use crate::query::wal_fusion::PkFusionMemory; // WAL-PK-FUSION: delete.
use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest};
@@ -264,7 +265,6 @@ struct FreshnessJob<S: HttpSend> {
enum TrackedJobResult {
None,
RefreshColumn,
MaterializedView,
}
#[async_trait]
@@ -305,11 +305,6 @@ impl<S: HttpSend> crate::job::JobHandle for FreshnessJob<S> {
})
})
}),
TrackedJobResult::MaterializedView => result.value().and_then(|value| {
serde_json::from_value::<RefreshMaterializedViewResult>(value.clone())
.ok()
.map(|result| result.version)
}),
}
.filter(|version| *version != 0);
if let Some(version) = result_version {
@@ -347,6 +342,10 @@ fn compute_min_timestamp(
}
}
fn quote_sql_identifier(identifier: &str) -> String {
format!("\"{}\"", identifier.replace('"', "\"\""))
}
fn freshness_headers_snapshot(
freshness: &Mutex<FreshnessState>,
interval: Option<Duration>,
@@ -496,6 +495,7 @@ pub struct RemoteTable<S: HttpSend = Sender> {
namespace: Vec<String>,
identifier: String,
server_version: ServerVersion,
sql_client: Option<SqlClient>,
version: Arc<RwLock<Option<u64>>>,
location: RwLock<Option<String>>,
@@ -654,12 +654,13 @@ impl<S: HttpSend> RemoteTable<S> {
Ok(job_id)
}
pub fn new(
pub(super) fn new_with_sql_client(
client: RestfulLanceDbClient<S>,
name: String,
namespace: Vec<String>,
identifier: String,
server_version: ServerVersion,
sql_client: Option<SqlClient>,
) -> Self {
Self {
client,
@@ -667,6 +668,7 @@ impl<S: HttpSend> RemoteTable<S> {
namespace,
identifier,
server_version,
sql_client,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -701,6 +703,7 @@ impl<S: HttpSend> RemoteTable<S> {
namespace: self.namespace.clone(),
identifier: self.identifier.clone(),
server_version: self.server_version.clone(),
sql_client: self.sql_client.clone(),
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1619,6 +1622,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: version.map(ServerVersion).unwrap_or_default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1644,6 +1648,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: ServerVersion::default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -1678,6 +1683,7 @@ mod test_utils {
namespace: vec![],
identifier: name,
server_version: version.map(ServerVersion).unwrap_or_default(),
sql_client: None,
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
@@ -2106,29 +2112,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
&self.identifier
}
async fn materialized_view_info(&self) -> Result<MaterializedViewInfo> {
#[derive(Deserialize)]
struct Projection {
output_column: String,
expression: String,
}
#[derive(Deserialize)]
struct DescribeMaterializedViewResponse {
source_table: String,
#[serde(default)]
source_namespace: Vec<String>,
#[serde(default)]
projections: Vec<Projection>,
#[serde(default)]
filter: Option<String>,
#[serde(default)]
limit: Option<u64>,
/// The defining query, once the server describes a view by it;
/// takes precedence over the structured fields.
#[serde(default)]
query: Option<String>,
#[serde(default)]
incarnation: Option<String>,
query: String,
}
let request = self.client.post(&format!(
@@ -2139,79 +2125,64 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
let response = self.check_table_response(&request_id, response).await?;
let response: DescribeMaterializedViewResponse =
response.json().await.err_to_http(request_id)?;
let definition = match response.query {
Some(query) => MaterializedViewDefinition::from_sql(&query)?,
None => MaterializedViewDefinition {
vector_source: None,
source_table: response.source_table,
source_namespace: response.source_namespace,
lateral: None,
projections: response
.projections
.into_iter()
.map(|projection| ViewProjection {
output: projection.output_column,
expression: projection.expression,
})
.collect(),
filter: response.filter,
group_by: Vec::new(),
limit: response.limit,
},
};
let parsed_definition = MaterializedViewDefinition::from_sql(&response.query).ok();
Ok(MaterializedViewInfo {
definition,
incarnation: response.incarnation,
definition_sql: response.query,
parsed_definition,
})
}
async fn refresh_materialized_view_async(
&self,
full: bool,
source_version: Option<u64>,
expected_incarnation: Option<&str>,
) -> Result<Job<RefreshMaterializedViewResult>> {
self.check_mutable().await?;
let mut body = serde_json::json!({ "full": full });
if let Some(source_version) = source_version {
body["source_version"] = source_version.into();
}
if let Some(expected_incarnation) = expected_incarnation {
body["expected_incarnation"] = expected_incarnation.into();
}
let request = self
.client
.post(&format!(
"/v1/materialized_view/{}/refresh",
self.identifier
))
.json(&body);
let freshness_request = self.snapshot_freshness_headers();
let (request_id, response) = self
.send_with_freshness(request, true, freshness_request)
.await?;
let response = self.check_table_response(&request_id, response).await?;
let status = response.status();
let body = response.text().await.err_to_http(request_id.clone())?;
if status != StatusCode::ACCEPTED {
return Err(Error::Http {
source: "materialized-view refresh must return 202 Accepted".into(),
request_id,
status_code: Some(status),
if source_version.is_some() {
return Err(Error::NotSupported {
message: "source-version pinning is available only for local materialized views"
.to_string(),
});
}
let job_id = extract_job_id(&body).ok_or_else(|| Error::Http {
source: "materialized-view refresh response did not contain a valid job_id".into(),
request_id,
status_code: Some(status),
let sql_client = self.sql_client.clone().ok_or_else(|| Error::NotSupported {
message: "SQL is unavailable for this remote table client".to_string(),
})?;
Ok(Job::new_typed(Box::new(FreshnessJob {
inner: RemoteJob::new(self.client.clone(), job_id),
freshness: self.freshness.clone(),
version: self.version.clone(),
tracked_result: TrackedJobResult::MaterializedView,
freshness_request,
})))
let statement = format!(
"REFRESH MATERIALIZED VIEW {}",
quote_sql_identifier(&self.name)
);
let namespace = self.namespace.clone();
let client = self.client.clone();
let name = self.name.clone();
let identifier = self.identifier.clone();
let server_version = self.server_version.clone();
let freshness = self.freshness.clone();
let freshness_request = self.snapshot_freshness_headers();
let job_sql_client = sql_client.clone();
Ok(
sql_client.submit_as_job(statement, namespace.clone(), move || async move {
let table = Self::new_with_sql_client(
client,
name,
namespace,
identifier,
server_version,
Some(job_sql_client),
);
table.checkout_latest().await?;
let version = table.version().await?;
let rows_written =
u64::try_from(table.count_rows(None).await?).map_err(|_| Error::Runtime {
message: "materialized-view row count exceeds u64".to_string(),
})?;
freshness_request.observe_version(&freshness, version);
Ok(RefreshMaterializedViewResult {
mode: crate::materialized_view::RefreshMode::Rebuild,
rows_written,
source_version: 0,
version,
})
}),
)
}
async fn query_snapshot(&self) -> Result<Arc<dyn BaseTable>> {
let description = self.describe().await?;
@@ -12976,54 +12947,16 @@ mod tests {
}
#[tokio::test]
async fn test_materialized_view_describe_and_refresh() {
async fn test_materialized_view_describe_and_refresh_requires_sql() {
const QUERY: &str =
"SELECT s.x, d.label FROM analytics.source s JOIN analytics.dim d ON s.id = d.id";
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
"/v1/materialized_view/my_table/describe" => http::Response::builder()
.status(200)
.body(
json!({
"name": "my_table",
"source_table": "source",
"source_namespace": ["analytics"],
"projections": [{
"output_column": "double_x",
"expression": "x * 2"
}],
"filter": "x > 0",
"limit": 10,
"inputs": ["x"],
"incarnation": "inc-1"
})
.to_string(),
)
.unwrap(),
"/v1/materialized_view/my_table/refresh" => {
assert_eq!(request.method(), "POST");
assert_eq!(
request_body_json(&request),
json!({
"full": true,
"source_version": 7,
"expected_incarnation": "inc-1"
})
);
http::Response::builder()
.status(202)
.body(json!({"job_id": "j1-mv-refresh"}).to_string())
.unwrap()
}
"/v1/jobs/describe" => http::Response::builder()
.status(200)
.body(
json!({
"job_id": "j1-mv-refresh",
"job_state": "DONE",
"result": {
"mode": "rebuild",
"rows_written": 2,
"source_version": 7,
"version": 9
}
"query": QUERY
})
.to_string(),
)
@@ -13031,21 +12964,10 @@ mod tests {
path => panic!("unexpected request: {path}"),
});
let view = crate::MaterializedView::from_table(table).await.unwrap();
assert_eq!(view.definition().source_table, "source");
assert_eq!(view.definition().source_namespace, ["analytics"]);
assert_eq!(view.incarnation(), Some("inc-1"));
let result = view
.refresh()
.full(true)
.source_version(7)
.expect_incarnation("inc-1")
.execute()
.await
.unwrap();
assert_eq!(result.mode, crate::RefreshMode::Rebuild);
assert_eq!(result.rows_written, 2);
assert_eq!(result.version, 9);
assert_eq!(view.definition_sql(), QUERY);
assert!(view.definition().is_err());
let error = view.refresh().execute().await.unwrap_err();
assert!(error.to_string().contains("SQL is unavailable"));
}
#[tokio::test]
rust/lancedb/src/table.rs
+1 -3
View File
@@ -563,7 +563,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
fn id(&self) -> &str;
/// Get the arrow [Schema] of the table.
async fn schema(&self) -> Result<SchemaRef>;
/// Read this table's materialized-view definition and incarnation.
/// Read this table's materialized-view definition.
#[doc(hidden)]
async fn materialized_view_info(
&self,
@@ -578,9 +578,7 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
#[doc(hidden)]
async fn refresh_materialized_view_async(
&self,
_full: bool,
_source_version: Option<u64>,
_expected_incarnation: Option<&str>,
) -> Result<Job<crate::materialized_view::RefreshMaterializedViewResult>> {
Err(Error::NotSupported {
message: "remote materialized-view refresh is not supported on this table type".into(),
rust/lancedb/src/table/datafusion/udtf/vector_query.rs
+90 -9
View File
@@ -6,11 +6,15 @@
use std::sync::Arc;
use arrow_schema::SchemaRef;
use arrow_array::{RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::prelude::{SessionConfig, SessionContext, col};
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::{Result, plan_err};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_common::{DataFusionError, Result, plan_err};
use datafusion_execution::{
SendableRecordBatchStream, TaskContext, runtime_env::RuntimeEnvBuilder,
};
use datafusion_expr::{Expr, TableType};
use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
use datafusion_physical_plan::{
@@ -25,7 +29,9 @@ use futures::{TryStreamExt, stream};
use lance::dataset::scanner::{RowAddrMask, RowAddrTreeMap};
use lance::{Dataset, index::DatasetIndexInternalExt};
use lance_core::datatypes::BlobHandling;
use lance_datafusion::exec::SessionContextExt;
use lance_index::metrics::NoOpMetricsCollector;
use roaring::RoaringTreemap;
use super::duplicate_pairs::{
DuplicatePairTask, DuplicatePairsConfig, DuplicatePairsExec, DuplicatePairsOutput,
@@ -36,6 +42,84 @@ use crate::table::{NativeTableExt, datafusion::BaseTableAdapter};
const REPRESENTATIVE_SORT_MEMORY_BYTES: usize = 256 * 1024 * 1024;
const MATERIALIZE_BATCH_ROWS: usize = 1024;
fn edge_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt64, false),
Field::new("b", DataType::UInt64, false),
]))
}
/// Sorting canonical edges makes the first unremoved endpoint the smallest
/// unassigned representative. The sort spills; only vertex masks stay in RAM.
async fn select_representatives(
stream: SendableRecordBatchStream,
memory: usize,
) -> crate::Result<(RoaringTreemap, RoaringTreemap)> {
let schema = edge_schema();
let batch_schema = schema.clone();
let canonical = stream
.map_ok(move |batch| {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("native row id");
RecordBatch::try_new(
batch_schema.clone(),
vec![
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.min(b)),
)),
Arc::new(UInt64Array::from_iter_values(
a.values().iter().zip(b.values()).map(|(&a, &b)| a.max(b)),
)),
],
)
.map_err(DataFusionError::from)
})
.and_then(futures::future::ready);
let canonical = Box::pin(RecordBatchStreamAdapter::new(schema, canonical));
let mut config = SessionConfig::default().with_batch_size(MATERIALIZE_BATCH_ROWS);
config.options_mut().execution.sort_spill_reservation_bytes = memory / 2;
let ctx = SessionContext::new_with_config_rt(
config,
RuntimeEnvBuilder::new()
.with_memory_limit(memory, 1.0)
.build_arc()?,
);
let mut sorted = ctx
.read_one_shot(canonical)?
.sort_by(vec![col("a"), col("b")])?
.execute_stream()
.await?;
let mut removed = RoaringTreemap::new();
let mut representatives = RoaringTreemap::new();
while let Some(batch) = sorted.try_next().await? {
let a = batch
.column(0)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
let b = batch
.column(1)
.as_any()
.downcast_ref::<UInt64Array>()
.expect("sorted id");
for (&a, &b) in a.values().iter().zip(b.values()) {
if a != b && !removed.contains(a) && !removed.contains(b) {
representatives.insert(a);
removed.insert(b);
}
}
}
Ok((removed, representatives))
}
/// Select nonempty physical IVF scopes before dispatch, never by filtering the
/// resulting pairs. A count spans all segments; partition IDs are segment-local.
#[derive(Debug, Clone, Default)]
@@ -322,12 +406,9 @@ impl ExecutionPlan for VectorDedupExec {
let dataset = self.dataset.clone();
let pairs = CoalescePartitionsExec::new(self.pairs.clone()).execute(0, context)?;
let stream = stream::once(async move {
let (removed, _) = crate::materialized_view::vector_dedup::select_representatives(
pairs,
REPRESENTATIVE_SORT_MEMORY_BYTES,
)
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
let (removed, _) = select_representatives(pairs, REPRESENTATIVE_SORT_MEMORY_BYTES)
.await
.map_err(|e| datafusion_common::DataFusionError::External(Box::new(e)))?;
let mut scanner = dataset.scan();
scanner
.with_row_addr_prefilter(RowAddrMask::from_block(RowAddrTreeMap::from_iter(
rust/lancedb/tests/duplicate_pairs.rs
+3 -205
View File
@@ -289,127 +289,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
);
let retry = ctx.sql(&sql).await?.collect().await?;
assert_eq!(pairs(&retry), pairs(&actual));
// The same table function can be maintained through the MV unit
// lifecycle. Staging attempts do not make any rows visible.
use lancedb::materialized_view::{
MaterializedViewDefinition, commit_partitioned_refresh, plan_partitioned_refresh,
prepare_definition, write_refresh_partition,
};
let view_name = format!("pairs_{}", threshold.to_bits());
let view = prepare_definition(&table, MaterializedViewDefinition::from_sql(&sql)?)
.await?
.create(&view_name)
.await?;
assert_eq!(view.table().schema().await?.fields().len(), 3);
let plan = plan_partitioned_refresh(view.table(), Some(version))
.await?
.unwrap();
assert_eq!(plan.units() as usize, tasks.len());
let plan = serde_json::from_slice(&serde_json::to_vec(&plan)?)?;
let first = write_refresh_partition(view.table(), 0, &plan).await?;
assert_eq!(view.table().count_rows(None).await?, 0);
assert!(
commit_partitioned_refresh(
view.table(),
&plan,
vec![first.clone(), first.clone()],
view.incarnation()
)
.await
.is_err()
);
// Lost receipt: its files remain unreferenced; the retry writes
// another attempt, of which only one receipt is ever committed.
drop(write_refresh_partition(view.table(), 0, &plan).await?);
let mut units = vec![first];
for unit in 1..plan.units() {
units.push(write_refresh_partition(view.table(), unit, &plan).await?);
}
assert_eq!(view.table().count_rows(None).await?, 0);
let outcome =
commit_partitioned_refresh(view.table(), &plan, units.clone(), view.incarnation())
.await?;
assert_eq!(outcome.rows_written as usize, pairs(&expected).len());
assert!(
commit_partitioned_refresh(view.table(), &plan, units, view.incarnation())
.await
.is_err()
);
use lancedb::query::ExecutableQuery;
let written = view
.table()
.query()
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
assert_eq!(pairs(&written), pairs(&expected));
// Repeating a refresh replaces the view rather than appending.
view.refresh().execute().await?;
assert_eq!(
view.table().count_rows(None).await? as usize,
pairs(&expected).len()
);
// The complete declarative pipeline keeps direct representatives
// and materializes original rows, including isolated source rows.
use lancedb::materialized_view::write_refresh_partition_with_inputs;
use lancedb::query::QueryBase;
let definition = MaterializedViewDefinition::from_sql(&format!(
"SELECT * FROM vector_dedup('source', {version}, 'vector', {threshold})"
))?;
let clean = prepare_definition(&table, definition)
.await?
.create(&format!("clean_{}", threshold.to_bits()))
.await?;
let clean_plan = plan_partitioned_refresh(clean.table(), Some(version))
.await?
.unwrap();
assert_eq!(clean_plan.dependency_units() as usize, tasks.len());
assert_eq!(clean_plan.units() as usize, tasks.len() + 2);
let clean_plan: lancedb::materialized_view::PartitionedRefreshPlan =
serde_json::from_slice(&serde_json::to_vec(&clean_plan)?)?;
assert!(
write_refresh_partition(clean.table(), tasks.len() as u32, &clean_plan)
.await
.is_err()
);
let mut receipts = Vec::new();
for unit in 0..clean_plan.units() {
let dependencies = &receipts[..receipts.len().min(tasks.len())];
if unit == 0 || unit == tasks.len() as u32 {
// Both selection and materialization can lose a receipt.
drop(
write_refresh_partition_with_inputs(
clean.table(),
unit,
&clean_plan,
dependencies,
)
.await?,
);
}
let receipt = write_refresh_partition_with_inputs(
clean.table(),
unit,
&clean_plan,
dependencies,
)
.await?;
receipts.push(serde_json::from_slice(&serde_json::to_vec(&receipt)?)?);
}
assert_eq!(clean.table().count_rows(None).await?, 0);
assert!(
commit_partitioned_refresh(
clean.table(),
&clean_plan,
receipts[..receipts.len() - 1].to_vec(),
clean.incarnation()
)
.await
.is_err()
);
commit_partitioned_refresh(clean.table(), &clean_plan, receipts, clean.incarnation())
.await?;
let mut edges = pairs(&expected)
.into_iter()
.map(|(a, b, _)| (a.min(b), a.max(b)))
@@ -449,27 +328,6 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
}
}
}
let kept = clean
.table()
.query()
.select(lancedb::query::Select::columns(&["id"]))
.execute()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut actual_ids = kept
.iter()
.flat_map(|b| {
b.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.iter()
.copied()
})
.collect::<Vec<_>>();
actual_ids.sort_unstable();
expected_ids.sort_unstable();
let mut query_ids = clean_query
.iter()
@@ -486,14 +344,10 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
query_ids.sort_unstable();
assert_eq!(query_ids, expected_ids);
assert_eq!(actual_ids, expected_ids);
if threshold == 1.0 {
assert!(actual_ids.contains(&0) && actual_ids.contains(&2));
assert!(!actual_ids.contains(&1), "A-B-C retains A and C");
assert!(query_ids.contains(&0) && query_ids.contains(&2));
assert!(!query_ids.contains(&1), "A-B-C retains A and C");
}
clean.refresh().execute().await?;
assert_eq!(clean.table().count_rows(None).await?, expected_ids.len());
assert_eq!(table.count_rows(None).await?, n * 2);
assert!(
DuplicatePairsExec::try_new(
ds.clone(),
@@ -527,10 +381,8 @@ async fn sql_pairs_match_native_snapshot_and_partition_tasks() -> anyhow::Result
}
#[tokio::test]
async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions() -> anyhow::Result<()>
{
async fn dedup_query_materializes_blob_payloads_and_null_vectors() -> anyhow::Result<()> {
use arrow_array::{LargeBinaryArray, StringArray, StructArray, types::Float32Type};
use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition};
for stable in [false, true] {
let dir = tempfile::tempdir()?;
let conn = lancedb::connect(dir.path().to_str().unwrap())
@@ -607,7 +459,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
true,
)
.await?;
let version = ds.version().version;
use lancedb::table::datafusion::udtf::{
duplicate_pairs::DuplicatePairsOutput,
vector_query::{PartitionSelection, VectorQueryOptions, VectorQueryTable},
@@ -692,59 +543,6 @@ async fn dedup_materializes_blob_payloads_null_vectors_and_snapshot_deletions()
assert_eq!(images.value(i), vec![ids.value(i) as u8; 4096]);
}
}
let definition = MaterializedViewDefinition::from_sql(&format!(
"SELECT * FROM vector_dedup('source', {version}, 'vector', 1)"
))?;
let clean = prepare_definition(&table, definition)
.await?
.create("clean")
.await?;
// A later source delete does not change the declared result snapshot.
table.delete("id = 0").await?;
let adapter = BaseTableAdapter::try_new(table.base_table().clone()).await?;
let explicit = adapter.vector_query_snapshot(Some(version)).await?;
let latest = adapter.vector_query_snapshot(None).await?;
assert_eq!(explicit.version().version, version);
assert!(latest.version().version > version);
assert_eq!(explicit.count_rows(None).await?, 5);
assert_eq!(latest.count_rows(None).await?, 4);
clean.refresh().execute().await?;
assert_eq!(clean.table().count_rows(None).await?, 4);
let output = Dataset::open(&clean.table().uri().await?).await?;
let batches = output
.scan()
.blob_handling(lance_core::datatypes::BlobHandling::AllBinary)
.try_into_stream()
.await?
.try_collect::<Vec<_>>()
.await?;
let mut ids = Vec::new();
for batch in batches {
let row_ids = batch
.column_by_name("id")
.unwrap()
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
let images = batch
.column_by_name("image")
.unwrap()
.as_any()
.downcast_ref::<LargeBinaryArray>()
.unwrap();
for i in 0..batch.num_rows() {
let id = row_ids.value(i);
assert_eq!(images.value(i), vec![id as u8; 4096]);
ids.push(id);
}
}
ids.sort_unstable();
assert_eq!(ids, [0, 2, 3, 4]);
assert_eq!(
table.count_rows(None).await?,
4,
"refresh must not mutate the source"
);
}
Ok(())
}