diff --git a/docs/src/js/classes/Connection.md b/docs/src/js/classes/Connection.md index f1aace87b..4bb769a0e 100644 --- a/docs/src/js/classes/Connection.md +++ b/docs/src/js/classes/Connection.md @@ -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 diff --git a/docs/src/js/classes/MaterializedView.md b/docs/src/js/classes/MaterializedView.md index 58eaa3c13..0e2413328 100644 --- a/docs/src/js/classes/MaterializedView.md +++ b/docs/src/js/classes/MaterializedView.md @@ -65,21 +65,14 @@ refresh(options?): Promise 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 diff --git a/docs/src/js/interfaces/RefreshMaterializedViewResult.md b/docs/src/js/interfaces/RefreshMaterializedViewResult.md index cb7100cd8..1c8b8bc21 100644 --- a/docs/src/js/interfaces/RefreshMaterializedViewResult.md +++ b/docs/src/js/interfaces/RefreshMaterializedViewResult.md @@ -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". *** diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f946fbb13..814bcb12a 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -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 diff --git a/nodejs/__test__/materialized_view.test.ts b/nodejs/__test__/materialized_view.test.ts index 6739320aa..f0c97e13a 100644 --- a/nodejs/__test__/materialized_view.test.ts +++ b/nodejs/__test__/materialized_view.test.ts @@ -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); }); }); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index d2a34669e..708f10fb4 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -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, diff --git a/nodejs/lancedb/materialized_view.ts b/nodejs/lancedb/materialized_view.ts index a59a3a5f6..7cf1690e0 100644 --- a/nodejs/lancedb/materialized_view.ts +++ b/nodejs/lancedb/materialized_view.ts @@ -9,10 +9,9 @@ export const DEFINITION_META_KEY = "mv.definition"; /** * The newest stored layout this version reads: `{"format": N, "query": ""}`, - * 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 { validateNonNegativeInteger(options?.sourceVersion, "sourceVersion"); - return await this.inner.refreshMaterializedView( - options?.full, - options?.sourceVersion, - ); + return await this.inner.refreshMaterializedView(options?.sourceVersion); } } diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index d4c5db509..6bdc69418 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -685,7 +685,6 @@ export abstract class Table { * @ignore */ abstract refreshMaterializedView( - full?: boolean, sourceVersion?: number, ): Promise; @@ -1420,10 +1419,9 @@ export class LocalTable extends Table { } async refreshMaterializedView( - full?: boolean, sourceVersion?: number, ): Promise { - return await this.inner.refreshMaterializedView(full, sourceVersion); + return await this.inner.refreshMaterializedView(sourceVersion); } async materializedViewDefinition(): Promise { diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index f5c3b2987..79b335c37 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -469,13 +469,12 @@ impl Table { #[napi(catch_unwind)] pub async fn refresh_materialized_view( &self, - full: Option, source_version: Option, ) -> napi::Result { 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 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 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(), diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 45a69a92e..213794ef3 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -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 diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index bbd6a3c5d..8ff53bcbe 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 1af3c1b4a..3a9815cb4 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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, diff --git a/python/python/lancedb/materialized_view.py b/python/python/lancedb/materialized_view.py index b4d455870..6342641ff 100644 --- a/python/python/lancedb/materialized_view.py +++ b/python/python/lancedb/materialized_view.py @@ -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": ""}``, -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))) diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index 1ee21fbac..bef173323 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -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, - }, - ) - ] diff --git a/python/src/connection.rs b/python/src/connection.rs index 744ab34b3..737ef05e2 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -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, limit: Option, with_no_data: bool, - vector_source_json: Option, ) -> PyResult> { 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, limit: Option, with_no_data: bool, - vector_source_json: Option, ) -> PyResult> { 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); } diff --git a/python/src/table.rs b/python/src/table.rs index d7519bdab..16f62bf8d 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -592,8 +592,6 @@ impl From 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, ) -> PyResult> { 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, ) -> PyResult> { 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(), + )) }) } diff --git a/rust/lancedb/src/job.rs b/rust/lancedb/src/job.rs index 54b46fc11..711a73033 100644 --- a/rust/lancedb/src/job.rs +++ b/rust/lancedb/src/job.rs @@ -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>) -> Self { Self::new_typed(Box::new(SpawnedJob::new(task))) } + + pub(crate) fn spawned_with_cancellation(task: JoinHandle>, cancel: F) -> Self + where + F: Fn() -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + let cancellation: Cancellation = Arc::new(move || Box::pin(cancel())); + Self::new_typed(Box::new(SpawnedJob::new_with_cancellation( + task, + cancellation, + ))) + } } impl Job @@ -501,10 +515,28 @@ impl Outcome { struct SpawnedJob { outcome: watch::Receiver>, abort: AbortHandle, + cancellation: Option, } +type CancellationFuture = Pin> + Send>>; +type Cancellation = Arc CancellationFuture + Send + Sync>; + impl SpawnedJob { fn new(task: JoinHandle>) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, None) + } + + fn new_with_cancellation(task: JoinHandle>, cancellation: Cancellation) -> Self + where + T: Serialize + Send + 'static, + { + Self::new_inner(task, Some(cancellation)) + } + + fn new_inner(task: JoinHandle>, cancellation: Option) -> 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(()) } diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index 7dec300a2..70b39506b 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -9,96 +9,20 @@ //! query this version cannot maintain reads back as unrefreshable, not as a //! plain table. Queries, indexes and search work on the view unchanged. -mod duplicate_pairs; mod grouped; -mod grouped_units; -mod partitioned; -pub(crate) mod vector_dedup; -pub use duplicate_pairs::{VectorSource, VectorSourceKind}; pub use grouped::IVF_PARTITION; -/// Refreshing a view grouped by [`IVF_PARTITION`] in units, one per index -/// partition, so the work can be spread over several processes. The caller -/// plans once, computes every unit of that plan wherever it likes, and -/// commits the whole set; a unit is bound to the view and plan it was -/// computed for, and the commit publishes all of them or none. -/// -/// ```no_run -/// # #![recursion_limit = "256"] -/// # use lancedb::materialized_view::{ -/// # commit_grouped_refresh, plan_grouped_refresh, write_grouped_unit, WrittenUnit, -/// # }; -/// # use lancedb::materialized_view::MaterializedView; -/// # async fn refresh_in_units(view: &MaterializedView) -> Result<(), Box> { -/// // The view this refresh was requested for. Every path below carries it, -/// // including the fallbacks: a view dropped and recreated meanwhile is a -/// // different view, and refreshing it in place of this one is the identity -/// // crossing the units path refuses. -/// let incarnation = view.incarnation().map(str::to_string); -/// -/// async fn in_one_pass(view: &MaterializedView, incarnation: Option<&str>) -> lancedb::Result<()> { -/// let mut refresh = view.refresh(); -/// if let Some(token) = incarnation { -/// refresh = refresh.expect_incarnation(token); -/// } -/// refresh.execute().await?; -/// Ok(()) -/// } -/// -/// // No plan means this view cannot be refreshed in units: it is not -/// // grouped by the partition of an indexed column, its index no longer -/// // says which fragments it covers, or it carries no incarnation to bind -/// // the units to. -/// let Some(plan) = plan_grouped_refresh(view.table(), None).await? else { -/// in_one_pass(view, incarnation.as_deref()).await?; -/// return Ok(()); -/// }; -/// -/// // Each unit is independent; this loop stands in for dispatching them. -/// // A unit whose view was dropped and recreated, or whose plan is stale, -/// // fails here rather than producing rows for the wrong view. -/// let mut units: Vec = Vec::new(); -/// for unit in 0..plan.units { -/// units.push(write_grouped_unit(view.table(), unit, &plan).await?); -/// } -/// -/// match commit_grouped_refresh(view.table(), &plan, units, incarnation.as_deref()).await { -/// Ok(result) => println!("{} rows at version {}", result.rows_written, result.version), -/// Err(err) => { -/// // The refresh is unrecorded, but it is not necessarily inert: a -/// // commit that raced a concurrent write can have landed without a -/// // watermark. Refresh from scratch rather than assuming either. -/// eprintln!("refresh in units unrecorded ({err}); refreshing in one pass"); -/// in_one_pass(view, incarnation.as_deref()).await?; -/// } -/// } -/// # Ok(()) -/// # } -/// ``` -pub use grouped_units::{ - GroupedRefreshPlan, WrittenUnit, commit_grouped_refresh, plan_grouped_refresh, - write_grouped_unit, -}; -pub use partitioned::{ - PartitionedRefreshPlan, WrittenPartition, commit_partitioned_refresh, plan_partitioned_refresh, - write_refresh_partition, write_refresh_partition_with_inputs, -}; mod query; pub mod refresh; -#[cfg(test)] -mod differential; - use std::collections::HashMap; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, FieldRef, Schema as ArrowSchema, SchemaRef}; use datafusion_common::ScalarValue; -use lance_core::ROW_ID; use lance_datafusion::planner::Planner; use serde::{Deserialize, Serialize}; use crate::connection::Connection; -use crate::database::listing::OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS; use crate::database::{CreateTableRequest, Database, OpenTableRequest}; use crate::embeddings::EmbeddingDefinition; use crate::function::FunctionBinding; @@ -116,26 +40,6 @@ pub use refresh::{RefreshMaterializedViewResult, RefreshMode}; /// Schema metadata key holding the view definition; see [`DEFINITION_FORMAT`]. pub const DEFINITION_META_KEY: &str = "mv.definition"; -/// Schema metadata key holding the view's incarnation: a token minted at each -/// physical creation of a view table, so a view dropped and recreated under -/// the same name and definition is still told apart from the one a caller -/// captured. A view whose metadata was replaced wholesale, or one declared -/// before tokens existed, carries none until its next refresh mints one. -pub const INCARNATION_META_KEY: &str = "mv.incarnation"; - -/// Schema metadata key holding the source table version the view was last -/// refreshed to. Absent until the first refresh. -pub const SOURCE_VERSION_META_KEY: &str = "mv.source_version"; - -/// Schema metadata key holding the wall-clock time of the last refresh, -/// in milliseconds since the epoch. -pub const REFRESHED_AT_MS_META_KEY: &str = "mv.refreshed_at_ms"; - -/// Column recording which source row produced each view row: the source's -/// stable `_rowid` at refresh time, which is why sources must keep stable -/// row ids. -pub const SOURCE_ROW_ID_COLUMN: &str = "__source_row_id"; - /// Field metadata namespace for declarations about schema structure, such as /// an unenforced primary key. const SCHEMA_DECLARATION_META_PREFIX: &str = "lance-schema:"; @@ -154,26 +58,16 @@ const COLUMN_DEFINITIONS_META_KEY: &str = "lancedb::column_definitions"; /// The newest layout this version reads under [`DEFINITION_META_KEY`]: /// `{"format": N, "query": ""}`, the query as -/// [`MaterializedViewDefinition::to_sql`] renders it. Format 3 identifies a -/// native duplicate-pair source, format 4 a dedup result source, format 2 a query with `GROUP BY`, and -/// format 1 an ordinary query. A reader refuses a newer format rather than +/// [`MaterializedViewDefinition::to_sql`] renders it. Format 2 identifies a +/// query with `GROUP BY`, and format 1 an ordinary query. A reader refuses a newer format rather than /// guess at it. The layout also carries `"kind": "query"`, which readers /// older than the format number report as an unrefreshable view instead of /// failing to read the metadata. -pub const DEFINITION_FORMAT: u64 = 4; +pub const DEFINITION_FORMAT: u64 = 2; /// The format `definition` is written in; see [`DEFINITION_FORMAT`]. fn format_of(definition: &MaterializedViewDefinition) -> u64 { - if let Some(source) = &definition.vector_source { - match source.kind { - VectorSourceKind::Pairs => 3, - VectorSourceKind::Dedup => 4, - } - } else if definition.is_grouped() { - 2 - } else { - 1 - } + if definition.is_grouped() { 2 } else { 1 } } /// Legacy `kind` tag of the structured layout written before @@ -236,9 +130,8 @@ pub enum LateralSource { /// The list column. column: String, }, - /// `name(args)`: a Function in `FROM` position, returning rows. The - /// server stages its output in a hidden table, recorded under - /// [`STAGING_META_KEY`]; a local database cannot refresh this form. + /// `name(args)`: a Function in `FROM` position, returning rows. Local + /// databases do not execute table functions during refresh. Function { /// The Function's name. name: String, @@ -254,37 +147,17 @@ pub(crate) struct ViewUnnest { pub alias: String, } -/// Schema metadata key holding a [`StagingBinding`], present only on a view -/// whose query calls a Function in `FROM` position. -pub const STAGING_META_KEY: &str = "mv.staging"; - -/// Where a Function in `FROM` position has its output staged: a hidden -/// table carrying every source column plus `column`, the Function's list -/// output. Refresh scans this table in place of the query's source. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct StagingBinding { - /// The staging table's name. - pub table: String, - /// Its namespace path; empty is the root namespace. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub namespace: Vec, - /// The list column holding the Function's output. - pub column: String, -} - -/// The list column refresh unnests for `definition`, given the staging its -/// Function output lives in. `None` when the query has no lateral item. +/// The list column refresh unnests for `definition`. `None` when the query +/// has no lateral item. pub(crate) fn physical_unnest( definition: &MaterializedViewDefinition, - staging: Option<&StagingBinding>, ) -> Result> { let Some(lateral) = &definition.lateral else { return Ok(None); }; - let column = match (&lateral.source, staging) { - (LateralSource::Unnest { column }, _) => column.clone(), - (LateralSource::Function { .. }, Some(staging)) => staging.column.clone(), - (LateralSource::Function { name, .. }, None) => { + let column = match &lateral.source { + LateralSource::Unnest { column } => column.clone(), + LateralSource::Function { name, .. } => { return Err(Error::NotSupported { message: format!( "'{name}' in FROM position is a Function; views over Function rows are \ @@ -299,25 +172,11 @@ pub(crate) fn physical_unnest( })) } -/// Read the staging binding off a view's schema metadata, if it has one. -pub fn read_staging(metadata: &HashMap) -> Result> { - metadata - .get(STAGING_META_KEY) - .map(|raw| { - serde_json::from_str(raw).map_err(|e| Error::Runtime { - message: format!("unreadable materialized view staging binding: {e}"), - }) - }) - .transpose() -} - /// The query that defines a materialized view, in the relational shape /// refresh maintains. Stored as SQL; see [`MaterializedViewDefinition::from_sql`] /// for the shape. #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaterializedViewDefinition { - /// A table-bound native vector source. It is never evaluated per source row. - pub vector_source: Option>, /// Name of the source table, in the same database as the view. pub source_table: String, /// Namespace path holding the source table; empty is the root namespace. @@ -408,10 +267,10 @@ pub enum StoredDefinition { #[doc(hidden)] #[derive(Debug, Clone, PartialEq, Eq)] pub struct MaterializedViewInfo { - /// The parsed view definition. - pub definition: MaterializedViewDefinition, - /// The current physical incarnation, when one has been minted. - pub incarnation: Option, + /// The defining SQL as stored. + pub definition_sql: String, + /// The definition parsed into the local refresh grammar, when supported. + pub parsed_definition: Option, } /// The backend-independent request used to create a remote materialized view. @@ -456,12 +315,78 @@ struct LegacyDefinition { /// Serialize `definition` into the layout stored under /// [`DEFINITION_META_KEY`]. pub(crate) fn definition_to_metadata(definition: &MaterializedViewDefinition) -> Result { - Ok(serde_json::json!({ + Ok(definition_metadata_from_sql_with_format( + &definition.to_sql(), + format_of(definition), + )) +} + +/// Serialize an arbitrary SQL definition into the materialized-view metadata +/// envelope. Sophon uses this for SQL queries that are broader than the local +/// refresh engine's supported shape. +pub fn definition_metadata_from_sql(sql: &str) -> String { + definition_metadata_from_sql_with_format(sql, 1) +} + +/// Serialize an arbitrary SQL definition together with the resolution +/// defaults needed to plan unqualified identifiers on a later refresh. +pub fn definition_metadata_from_sql_with_defaults( + sql: &str, + default_database: &str, + default_namespace_path: &[String], +) -> String { + serde_json::json!({ "kind": "query", - "format": format_of(definition), - "query": definition.to_sql(), + "format": 1, + "query": sql, + "default_database": default_database, + "default_namespace_path": default_namespace_path, }) - .to_string()) + .to_string() +} + +fn definition_metadata_from_sql_with_format(sql: &str, format: u64) -> String { + serde_json::json!({ + "kind": "query", + "format": format, + "query": sql, + }) + .to_string() +} + +/// Return the defining SQL stored in [`DEFINITION_META_KEY`], without +/// requiring the local refresh parser to understand it. +pub fn read_definition_sql(metadata: &HashMap) -> Result> { + Ok(read_definition_sql_with_defaults(metadata)?.map(|definition| definition.query)) +} + +/// A stored SQL definition and the optional name-resolution defaults recorded +/// by a remote SQL engine. +#[derive(Debug, Clone, PartialEq, Eq, Deserialize)] +pub struct MaterializedViewSqlDefinition { + /// The SQL query that defines the materialized view. + pub query: String, + /// The database used to resolve unqualified references when the view was created. + #[serde(default)] + pub default_database: Option, + /// The namespace used to resolve unqualified references when the view was created. + #[serde(default)] + pub default_namespace_path: Option>, +} + +/// Read the SQL definition and its optional resolution defaults from +/// [`DEFINITION_META_KEY`]. +pub fn read_definition_sql_with_defaults( + metadata: &HashMap, +) -> Result> { + let Some(raw) = metadata.get(DEFINITION_META_KEY) else { + return Ok(None); + }; + serde_json::from_str(raw) + .map(Some) + .map_err(|e| Error::Runtime { + message: format!("unreadable materialized view definition: {e}"), + }) } /// Read a view definition off a schema metadata map, if it carries one. @@ -516,7 +441,6 @@ pub fn read_definition(metadata: &HashMap) -> Result, ) -> Result { - let incarnation = metadata.get(INCARNATION_META_KEY).cloned(); match read_definition(metadata)? { Some(StoredDefinition::Query(definition)) => Ok(MaterializedViewInfo { - definition, - incarnation, + definition_sql: definition.to_sql(), + parsed_definition: Some(definition), }), Some(StoredDefinition::Newer { format }) => Err(Error::NotSupported { message: format!( @@ -563,19 +486,14 @@ pub(crate) struct Planned { pub fields: Vec, pub lineage: Lineage, /// Source columns the query reads; a read through the unnest alias is - /// recorded as the list column, which is what the source has and what - /// incremental refresh watches. + /// recorded as the list column present in the source schema. pub inputs: Vec, } pub(crate) fn plan( source_schema: SchemaRef, definition: &MaterializedViewDefinition, - staging: Option<&StagingBinding>, ) -> Result { - if definition.vector_source.is_some() { - return duplicate_pairs::plan(source_schema, definition); - } let filter = definition .filter .as_deref() @@ -593,7 +511,7 @@ pub(crate) fn plan( // Projections are typed against the source, or for an unnested view // against the source with the list column replaced by its element // under the alias, where `c.chunk` is an ordinary nested path. - let unnest = physical_unnest(definition, staging)?; + let unnest = physical_unnest(definition)?; let source_schema = match &unnest { None => source_schema, Some(unnest) => { @@ -607,12 +525,9 @@ pub(crate) fn plan( } }; let projections: Vec<(String, String)> = if definition.selects_star() { - // `SELECT *`. A source that is itself a view carries its own - // provenance column; the new view records its own, not a copy. source_schema .fields() .iter() - .filter(|f| f.name() != SOURCE_ROW_ID_COLUMN) .map(|f| (f.name().clone(), query::ident_sql(f.name()))) .collect() } else { @@ -653,25 +568,12 @@ pub(crate) fn plan( name: output.clone(), }); } - if output == SOURCE_ROW_ID_COLUMN || output == ROW_ID { - return Err(Error::InvalidInput { - message: format!("view column name '{output}' is reserved"), - }); - } - let parsed = planner .parse_expr(expression) .map_err(|e| Error::InvalidExpression { column: output.clone(), message: e.to_string(), })?; - // Before optimization: the simplifier folds a stable-but-not-immutable - // call like now() into a literal, hiding it from the check while the - // stored definition keeps the call. - ensure_immutable(&parsed, |message| Error::InvalidExpression { - column: output.clone(), - message, - })?; let expr = planner .optimize_expr(parsed) .map_err(|e| Error::InvalidExpression { @@ -734,7 +636,6 @@ pub(crate) fn plan( } let definition = MaterializedViewDefinition { - vector_source: None, source_table: definition.source_table.clone(), source_namespace: definition.source_namespace.clone(), lateral: definition.lateral.clone(), @@ -760,17 +661,14 @@ pub(crate) fn plan( }) } -/// Check that `filter` is an immutable boolean predicate over `source_schema`, -/// returning the columns it reads. +/// Check that `filter` is a boolean predicate over `source_schema`, returning +/// the columns it reads. pub(crate) fn plan_filter(source_schema: &SchemaRef, filter: &str) -> Result> { let expr = Planner::new(source_schema.clone()) .parse_filter(filter) .map_err(|e| Error::InvalidInput { message: format!("invalid view filter: {e}"), })?; - ensure_immutable(&expr, |message| Error::InvalidInput { - message: format!("invalid view filter: {message}"), - })?; let filter_inputs = resolve_inputs(source_schema, &expr, |message| Error::InvalidInput { message: format!("invalid view filter: {message}"), })?; @@ -839,41 +737,6 @@ pub(crate) fn flattened_schema( Ok(Arc::new(ArrowSchema::new(fields))) } -/// Reject any function that is not immutable: a view definition has to -/// evaluate identically across refreshes, or incremental maintenance would -/// mix rows from different evaluations of the same definition. -fn ensure_immutable(expr: &datafusion_expr::Expr, error: impl Fn(String) -> Error) -> Result<()> { - use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; - use datafusion_expr::Volatility; - - // Labeled immutable but not determined by row values alone: version() - // depends on the build, the arrow_* introspectors on schema state. - const NOT_VALUE_DETERMINED: &[&str] = - &["version", "arrow_typeof", "arrow_field", "arrow_metadata"]; - - let mut offending: Option = None; - expr.apply(|node| { - if let datafusion_expr::Expr::ScalarFunction(function) = node { - let name = function.func.name(); - if function.func.signature().volatility != Volatility::Immutable - || NOT_VALUE_DETERMINED.contains(&name) - { - offending = Some(name.to_string()); - return Ok(TreeNodeRecursion::Stop); - } - } - Ok(TreeNodeRecursion::Continue) - }) - .map_err(|e| error(e.to_string()))?; - match offending { - Some(name) => Err(error(format!( - "function '{name}' is not immutable and would evaluate differently \ - across refreshes" - ))), - None => Ok(()), - } -} - /// The root of a possibly-dotted column path: `metadata.age` -> `metadata`. fn root(path: &str) -> &str { path.split('.').next().unwrap_or(path) @@ -1125,7 +988,7 @@ fn project_schema(schema: &ArrowSchema, columns: &[String]) -> SchemaRef { } /// A validated view declaration, ready to become a table: the projected -/// fields plus [`SOURCE_ROW_ID_COLUMN`], definition stamped in metadata. +/// fields with the definition stamped in metadata. /// Produced only by [`prepare_declaration`]. #[derive(Clone)] pub struct PreparedDeclaration { @@ -1158,8 +1021,7 @@ impl PreparedDeclaration { } /// The schema the view will have: the declared columns in order, any - /// internal projections added by [`PreparedDeclaration::input_column`], - /// then [`SOURCE_ROW_ID_COLUMN`]. + /// internal projections added by [`PreparedDeclaration::input_column`]. pub fn schema(&self) -> &SchemaRef { &self.schema } @@ -1168,11 +1030,6 @@ impl PreparedDeclaration { /// read: the column the view projects it to, if any, otherwise an /// internal projection added here, named by [`input_column_name`]. pub fn input_column(&mut self, source_column: &str) -> Result { - if self.definition.vector_source.is_some() { - return Err(Error::InvalidInput { - message: "native vector views have no per-source-row inputs".into(), - }); - } // A grouped view's rows are groups; no source row carries a value into one. if self.definition.is_grouped() { return Err(Error::InvalidInput { @@ -1194,14 +1051,13 @@ impl PreparedDeclaration { if self.schema.field_with_name(&name).is_ok() { return Err(Error::ColumnAlreadyExists { name }); } - let row_id = self.row_id_index()?; let mut fields: Vec = self .schema .fields() .iter() .map(|f| f.as_ref().clone()) .collect(); - fields.insert(row_id, without_declarations(&field.with_name(name.clone()))); + fields.push(without_declarations(&field.with_name(name.clone()))); self.definition.projections.push(ViewProjection { output: name.clone(), expression: source_column @@ -1273,11 +1129,6 @@ impl PreparedDeclaration { columns: Vec<(usize, ArrowField)>, bindings: &[FunctionBinding], ) -> Result { - if self.definition.vector_source.is_some() { - return Err(Error::InvalidInput { - message: "computed columns on native vector views are not supported".into(), - }); - } let invalid = |message: String| Error::InvalidInput { message }; if columns.is_empty() { return Err(invalid("at least one computed column is needed".into())); @@ -1303,10 +1154,7 @@ impl PreparedDeclaration { columns.sort_by_key(|(position, _)| *position); for (inserted, (position, field)) in columns.iter().enumerate() { let name = field.name().as_str(); - if name == SOURCE_ROW_ID_COLUMN - || name == ROW_ID - || name.starts_with(INPUT_COLUMN_PREFIX) - { + if name.starts_with(INPUT_COLUMN_PREFIX) { return Err(invalid(format!("view column name '{name}' is reserved"))); } if fields.iter().any(|f| f.name() == name) { @@ -1348,16 +1196,8 @@ impl PreparedDeclaration { Ok(self) } - fn row_id_index(&self) -> Result { - self.schema - .index_of(SOURCE_ROW_ID_COLUMN) - .map_err(|e| Error::Runtime { - message: e.to_string(), - }) - } - /// Columns the declaration lists: everything before the internal - /// projections and the provenance column. + /// projections. fn visible_count(&self) -> usize { self.definition.projections.len() - self.internal_inputs + computed_columns(&self.schema).len() @@ -1366,9 +1206,7 @@ impl PreparedDeclaration { /// Create the view table and verify it, consuming the declaration. /// /// The view goes at the root of the source's own database, where refresh - /// resolves the recorded source coordinate. Stable row ids are requested - /// at both levels and verified rather than trusted; nothing is rolled - /// back on failure. + /// resolves the recorded source coordinate. pub async fn create(self, name: &str) -> Result { self.create_in(&[], name).await } @@ -1382,53 +1220,16 @@ impl PreparedDeclaration { ) -> Result { let empty: Vec> = vec![]; - // Minted here, not at preparation: a declaration can be cloned and - // create more than one physical table, and each needs its own token. - let incarnation = uuid::Uuid::new_v4().to_string(); - let mut metadata = self.schema.metadata().clone(); - metadata.insert(INCARNATION_META_KEY.to_string(), incarnation.clone()); - let schema = Arc::new(ArrowSchema::new_with_metadata( - self.schema.fields().clone(), - metadata, - )); let reader: Box = - Box::new(arrow_array::RecordBatchIterator::new(empty, schema)); + Box::new(arrow_array::RecordBatchIterator::new(empty, self.schema)); let mut request = CreateTableRequest::new(name.to_string(), Box::new(reader)); request.namespace_path = namespace_path.to_vec(); - let write_params = request - .write_options - .lance_write_params - .get_or_insert_with(Default::default); - write_params.enable_stable_row_ids = true; - let store_params = write_params - .store_params - .get_or_insert_with(Default::default); - crate::connection::merge_storage_options( - store_params, - [( - OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(), - "true".to_string(), - )], - ); let table = self.database.clone().create_table(request).await?; let table = Table::new(table, self.database); - let stable = match table.as_native() { - Some(native) => native.dataset.get().await?.manifest.uses_stable_row_ids(), - None => false, - }; - if !stable { - return Err(Error::Runtime { - message: format!( - "view '{name}' was created without stable row ids: the database \ - ignored the creation option; the table remains and is not \ - usable as a materialized view" - ), - }); - } Ok(MaterializedView { table, - definition: self.definition, - incarnation: Some(incarnation), + definition_sql: self.definition.to_sql(), + parsed_definition: Some(self.definition), }) } } @@ -1484,7 +1285,6 @@ pub async fn prepare_declaration( limit: Option, ) -> Result { let definition = MaterializedViewDefinition { - vector_source: None, source_table: source.name().to_string(), source_namespace: source.namespace().to_vec(), lateral: None, @@ -1508,7 +1308,7 @@ pub async fn prepare_declaration( /// Validate `definition` over `source`, the table it names. The declaration /// is planned against the source as refresh will reach it, and the result /// creates the view with [`PreparedDeclaration::create`]. A query calling a -/// Function in `FROM` position needs [`prepare_staged_definition`]. +/// Function in `FROM` position is not supported by local refresh. /// /// ``` /// # #![recursion_limit = "256"] @@ -1546,27 +1346,6 @@ pub async fn prepare_declaration( /// # Ok(()) /// # } /// ``` -/// -/// An indexed, pinned source can materialize retained original rows using the -/// default greedy direct-representative policy. A row is removed only when it -/// directly matches a retained row; an A-B-C chain keeps A and C. The source is -/// unchanged. Hosts distribute this through [`plan_partitioned_refresh`] and -/// [`write_refresh_partition_with_inputs`], with a barrier between selection -/// tasks and source-fragment materialization tasks. -/// -/// ``` -/// # #![recursion_limit = "256"] -/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_definition}; -/// # async fn dedup(images: &lancedb::Table) -> Result<(), Box> { -/// let version = images.version().await?; -/// let definition = MaterializedViewDefinition::from_sql(&format!( -/// "SELECT * FROM vector_dedup('images', {version}, 'phash', 4)" -/// ))?; -/// let clean = prepare_definition(images, definition).await?.create("images_clean").await?; -/// clean.refresh().execute().await?; -/// # Ok(()) -/// # } -/// ``` pub async fn prepare_definition( source: &Table, definition: MaterializedViewDefinition, @@ -1583,59 +1362,12 @@ pub async fn prepare_definition( ), }); } - prepare_with(source, definition, None).await -} - -/// Validate a `definition` whose query calls a Function in `FROM` position, -/// planned over `staging`: a table carrying every column of the query's -/// source plus `column`, the Function's list output for that row. The view -/// records the query as written and the staging under [`STAGING_META_KEY`]; -/// refresh scans the staging table and unnests `column`. -/// -/// ``` -/// # #![recursion_limit = "256"] -/// use lancedb::materialized_view::{MaterializedViewDefinition, prepare_staged_definition}; -/// -/// // `staging` holds every column of `docs` plus `chunks`, the list -/// // `chunk(body)` returned for each row. -/// # async fn declare(staging: &lancedb::Table) -> Result<(), Box> { -/// let definition = MaterializedViewDefinition::from_sql( -/// "SELECT id, c.text, c.ordinal FROM docs, chunk(body) AS c", -/// )?; -/// let view = prepare_staged_definition(staging, definition, "chunks") -/// .await? -/// .create("chunks") -/// .await?; -/// view.refresh().execute().await?; -/// # Ok(()) -/// # } -/// ``` -pub async fn prepare_staged_definition( - staging: &Table, - definition: MaterializedViewDefinition, - column: impl Into, -) -> Result { - if !matches!( - definition.lateral.as_ref().map(|l| &l.source), - Some(LateralSource::Function { .. }) - ) { - return Err(Error::InvalidInput { - message: "only a query calling a Function in FROM position takes a staging table" - .into(), - }); - } - let binding = StagingBinding { - table: staging.name().to_string(), - namespace: staging.namespace().to_vec(), - column: column.into(), - }; - prepare_with(staging, definition, Some(binding)).await + prepare_with(source, definition).await } async fn prepare_with( source: &Table, definition: MaterializedViewDefinition, - staging: Option, ) -> Result { let Some(caller_native) = source.as_native() else { return Err(Error::NotSupported { @@ -1682,17 +1414,6 @@ async fn prepare_with( ), }); } - if definition.vector_source.is_none() - && !native.dataset.get().await?.manifest.uses_stable_row_ids() - { - return Err(Error::InvalidInput { - message: format!( - "materialized views require stable row ids on the source table; \ - create '{}' with storage option new_table_enable_stable_row_ids=true", - source.name() - ), - }); - } refresh::ensure_no_mem_wal( native.dataset.get().await?.as_ref(), "source table", @@ -1710,45 +1431,23 @@ async fn prepare_with( message: format!("view column name '{}' is reserved", reserved.output), }); } - let source_schema = if let Some(pairs) = &definition.vector_source { - let ds = native - .dataset - .get() - .await? - .checkout_version(pairs.dataset_version) - .await?; - crate::table::datafusion::udtf::duplicate_pairs::plan_duplicate_pairs( - Arc::new(ds.clone()), - &pairs.config()?, - ) - .await?; - Arc::new(ArrowSchema::from(ds.schema())) - } else { - resolved.schema().await? - }; + let source_schema = resolved.schema().await?; let source_metadata = source_schema.metadata().clone(); let Planned { definition, - mut fields, + fields, lineage, .. - } = plan(source_schema.clone(), &definition, staging.as_ref())?; + } = plan(source_schema.clone(), &definition)?; if definition.is_grouped() { grouped::check(native.dataset.get().await?.as_ref(), &definition).await?; } // What later projections (`input_column`) are planned against: for an // unnested view the flattened schema, where the alias is a column. - let planning_schema = match physical_unnest(&definition, staging.as_ref())? { + let planning_schema = match physical_unnest(&definition)? { None => source_schema.clone(), Some(unnest) => flattened_schema(&source_schema, &unnest)?, }; - if definition.vector_source.is_none() { - fields.push(ArrowField::new( - SOURCE_ROW_ID_COLUMN, - DataType::UInt64, - false, - )); - } // Only column-describing metadata comes along: structural declarations // describe how a table is written, and a view is written by refresh alone. let mut metadata: HashMap = HashMap::new(); @@ -1766,14 +1465,6 @@ async fn prepare_with( DEFINITION_META_KEY.to_string(), definition_to_metadata(&definition)?, ); - if let Some(staging) = &staging { - metadata.insert( - STAGING_META_KEY.to_string(), - serde_json::to_string(staging).map_err(|e| Error::Runtime { - message: format!("failed to serialize the staging binding: {e}"), - })?, - ); - } Ok(PreparedDeclaration { schema: Arc::new(ArrowSchema::new_with_metadata(fields, metadata)), definition, @@ -1792,13 +1483,31 @@ pub struct CreateMaterializedViewBuilder { namespace: Vec, source: String, source_namespace: Vec, - vector_source: Option, projections: Vec<(String, String)>, filter: Option, limit: Option, with_no_data: bool, } +fn source_sql_identifier( + namespace_path: &[String], + source: &str, + namespace_in_schema_position: bool, +) -> String { + fn quote(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) + } + + let namespace = if namespace_in_schema_position && !namespace_path.is_empty() { + vec![namespace_path.join("$")] + } else { + namespace_path.to_vec() + }; + let mut parts = namespace.iter().map(|part| quote(part)).collect::>(); + parts.push(quote(source)); + parts.join(".") +} + impl CreateMaterializedViewBuilder { pub(crate) fn new(connection: Connection, name: String, source: String) -> Self { Self { @@ -1807,7 +1516,6 @@ impl CreateMaterializedViewBuilder { namespace: Vec::new(), source, source_namespace: Vec::new(), - vector_source: None, projections: Vec::new(), filter: None, limit: None, @@ -1828,38 +1536,6 @@ impl CreateMaterializedViewBuilder { self } - /// Read native pairs or retained original rows from a fixed indexed snapshot. - /// This source selects its full output schema and cannot be combined with - /// projections, filters or limits. Remote creation uses the ordinary MV - /// endpoint and its registry job; it does not execute a SQL query client-side. - pub fn vector_source(mut self, source: VectorSource) -> Self { - self.vector_source = Some(source); - self - } - - fn native_definition(&self) -> Result> { - let Some(source) = &self.vector_source else { - return Ok(None); - }; - source.config()?; - if !self.projections.is_empty() || self.filter.is_some() || self.limit.is_some() { - return Err(Error::InvalidInput { - message: "native vector sources cannot be combined with select, where or limit" - .into(), - }); - } - Ok(Some(MaterializedViewDefinition { - vector_source: Some(Box::new(source.clone())), - source_table: self.source.clone(), - source_namespace: self.source_namespace.clone(), - projections: vec![ViewProjection::star()], - filter: None, - lateral: None, - group_by: Vec::new(), - limit: None, - })) - } - /// The view's columns, as `(name, SQL expression)` pairs. Not calling /// this selects every source column, expanded at creation time. pub fn select( @@ -1893,9 +1569,6 @@ impl CreateMaterializedViewBuilder { } fn query(&self) -> Result { - if let Some(definition) = self.native_definition()? { - return Ok(definition.to_sql()); - } fn quote(name: &str) -> String { format!("\"{}\"", name.replace('"', "\"\"")) } @@ -1909,13 +1582,11 @@ impl CreateMaterializedViewBuilder { .collect::>() .join(", ") }; - let source = self - .source_namespace - .iter() - .chain(std::iter::once(&self.source)) - .map(|part| quote(part)) - .collect::>() - .join("."); + let source = source_sql_identifier( + &self.source_namespace, + &self.source, + self.connection.uri().starts_with("db://"), + ); let mut query = format!("SELECT {projection} FROM {source}"); if let Some(filter) = &self.filter { query.push_str(" WHERE "); @@ -1928,9 +1599,7 @@ impl CreateMaterializedViewBuilder { } /// Submit creation and initial population, returning a [`Job`] that - /// settles when the view is ready. Ordinary sources must keep stable row - /// ids for provenance across compaction. Native vector sources also support - /// physical row IDs because their definitions pin the dataset version. + /// settles when the view is ready. pub async fn execute_async(self) -> Result { if self.connection.uri().starts_with("db://") { return self @@ -1967,24 +1636,19 @@ impl CreateMaterializedViewBuilder { } async fn execute_native(self) -> Result { - let native_definition = self.native_definition()?; let source = self .connection .open_table(&self.source) .namespace(self.source_namespace.clone()) .execute() .await?; - let prepared = if let Some(definition) = native_definition { - prepare_definition(&source, definition).await? - } else { - prepare_declaration( - &source, - (!self.projections.is_empty()).then_some(self.projections.as_slice()), - self.filter.as_deref(), - self.limit, - ) - .await? - }; + let prepared = prepare_declaration( + &source, + (!self.projections.is_empty()).then_some(self.projections.as_slice()), + self.filter.as_deref(), + self.limit, + ) + .await?; let view = prepared.create_in(&self.namespace, &self.name).await?; if !self.with_no_data { view.refresh().execute().await?; @@ -1993,24 +1657,24 @@ impl CreateMaterializedViewBuilder { } } -/// A handle on a materialized view: the view table plus its parsed definition. +/// A handle on a materialized view: the view table plus its defining SQL. #[derive(Debug, Clone)] pub struct MaterializedView { table: Table, - definition: MaterializedViewDefinition, - incarnation: Option, + definition_sql: String, + parsed_definition: Option, } impl MaterializedView { /// Interpret `table` as a materialized view: [`Error::NotAMaterializedView`] - /// for a plain table, [`Error::NotSupported`] for a query this version - /// cannot refresh. + /// for a plain table. A remote definition broader than the local refresh + /// grammar remains available through [`Self::definition_sql`]. pub async fn from_table(table: Table) -> Result { let info = table.base_table().materialized_view_info().await?; Ok(Self { table, - definition: info.definition, - incarnation: info.incarnation, + definition_sql: info.definition_sql, + parsed_definition: info.parsed_definition, }) } @@ -2024,23 +1688,28 @@ impl MaterializedView { self.table.name() } - /// The query that defines the view. - pub fn definition(&self) -> &MaterializedViewDefinition { - &self.definition + /// The query parsed into the local refresh grammar. + /// + /// Sophon may store SQL broader than the local grammar; use + /// [`Self::definition_sql`] when the original SQL is sufficient. + pub fn definition(&self) -> Result<&MaterializedViewDefinition> { + self.parsed_definition + .as_ref() + .ok_or_else(|| Error::NotSupported { + message: + "this materialized-view definition is supported only by the remote SQL engine" + .to_string(), + }) } - /// The view's incarnation token as of when this handle was opened; see - /// [`RefreshMaterializedViewBuilder::expect_incarnation`]. `None` for a - /// view that has none yet (see [`INCARNATION_META_KEY`]). - pub fn incarnation(&self) -> Option<&str> { - self.incarnation.as_deref() + /// The complete SQL query that defines the view. + pub fn definition_sql(&self) -> &str { + &self.definition_sql } /// Recompute the view from its source. /// - /// By default the refresh is incremental when the source's changes can be - /// reconciled into the view, and otherwise rebuilds; see - /// [`RefreshMaterializedViewBuilder`]. + /// Every refresh rebuilds the view from the selected source version. /// /// ```no_run /// # #![recursion_limit = "256"] @@ -2054,9 +1723,7 @@ impl MaterializedView { pub fn refresh(&self) -> RefreshMaterializedViewBuilder { RefreshMaterializedViewBuilder { view: self.clone(), - full: false, source_version: None, - expected_incarnation: None, } } } @@ -2064,38 +1731,17 @@ impl MaterializedView { /// Builds a refresh. Created by [`MaterializedView::refresh`]. pub struct RefreshMaterializedViewBuilder { view: MaterializedView, - full: bool, source_version: Option, - expected_incarnation: Option, } impl RefreshMaterializedViewBuilder { - /// Rebuild the view even where an incremental refresh would do. - pub fn full(mut self, full: bool) -> Self { - self.full = full; - self - } - - /// Refresh to this source table version instead of the latest. + /// Refresh a local view to this source table version instead of the latest. + /// Remote SQL refreshes do not support source-version pinning. pub fn source_version(mut self, version: u64) -> Self { self.source_version = Some(version); self } - /// Refresh only if the view is still the incarnation that minted `token` - /// (see [`MaterializedView::incarnation`]): a refresh requested against - /// one declaration must not land in a view dropped and recreated since, - /// even under the same name and definition. - /// - /// Best effort. The token is read from the latest stored manifest before - /// planning and again immediately before every commit, but it is not part - /// of the commit's own condition, so a recreation that lands between that - /// final read and the commit is not caught. - pub fn expect_incarnation(mut self, token: impl Into) -> Self { - self.expected_incarnation = Some(token.into()); - self - } - /// Submit the refresh and return a job that settles with its result. pub async fn execute_async(self) -> Result> { if self.view.table.as_native().is_none() { @@ -2103,34 +1749,18 @@ impl RefreshMaterializedViewBuilder { .view .table .base_table() - .refresh_materialized_view_async( - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) + .refresh_materialized_view_async(self.source_version) .await; } Ok(Job::spawned(tokio::spawn(async move { - refresh::execute_refresh( - &self.view.table, - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) - .await + refresh::execute_refresh(&self.view.table, self.source_version).await }))) } /// Refresh the view, waiting for the job to finish. pub async fn execute(self) -> Result { if self.view.table.as_native().is_some() { - return refresh::execute_refresh( - &self.view.table, - self.full, - self.source_version, - self.expected_incarnation.as_deref(), - ) - .await; + return refresh::execute_refresh(&self.view.table, self.source_version).await; } self.execute_async().await?.wait().await } @@ -2249,2033 +1879,84 @@ mod tests { use super::*; use crate::connect; - use crate::table::WriteOptions; - async fn people_db() -> Connection { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("name", Utf8, ["ada", "grace", "alan"]), - ("age", Int32, [36, 85, 41]) - ) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - conn - } - - /// Sources must keep stable row ids; see the create-time gate. - pub(super) fn stable_row_ids() -> WriteOptions { - WriteOptions { - lance_write_params: Some(lance::dataset::WriteParams { - enable_stable_row_ids: true, - ..Default::default() - }), - } - } - - /// The error a doomed declaration against `people` produces. - async fn declare_err( - cfg: impl FnOnce(CreateMaterializedViewBuilder) -> CreateMaterializedViewBuilder, - ) -> Error { - let conn = people_db().await; - cfg(conn.create_materialized_view("bad", "people")) - .execute() - .await - .unwrap_err() - } - - #[tokio::test] - async fn test_create_records_the_definition() { - let conn = people_db().await; - let view = conn - .create_materialized_view("adults", "people") - .select([("name", "name"), ("shout", "upper(name)")]) - .only_if("age >= 18") - .limit(10) - .execute() - .await - .unwrap(); - - assert_eq!(view.name(), "adults"); + #[test] + fn remote_source_namespace_uses_one_sql_schema_position() { + let namespace = vec!["analytics".to_string(), "daily".to_string()]; assert_eq!( - view.definition(), - &MaterializedViewDefinition { - vector_source: None, - source_table: "people".into(), - source_namespace: Vec::new(), - projections: vec![ - ViewProjection { - output: "name".into(), - expression: "name".into() - }, - ViewProjection { - output: "shout".into(), - expression: "upper(name)".into() - }, - ], - filter: Some("age >= 18".into()), - group_by: Vec::new(), - limit: Some(10), - lateral: None, - } - ); - - // The definition round-trips off the stored schema, not the handle. - let reopened = conn.open_materialized_view("adults").await.unwrap(); - assert_eq!(reopened.definition(), view.definition()); - } - - #[tokio::test] - async fn test_view_schema_is_derived_from_the_query() { - let conn = people_db().await; - let view = conn - .create_materialized_view("shapes", "people") - .select([("shout", "upper(name)"), ("next_age", "age + 1")]) - .execute() - .await - .unwrap(); - - let schema = view.table().schema().await.unwrap(); - assert_eq!( - schema.field_with_name("shout").unwrap().data_type(), - &DataType::Utf8 + source_sql_identifier(&namespace, "source", true), + "\"analytics$daily\".\"source\"" ); assert_eq!( - schema.field_with_name("next_age").unwrap().data_type(), - &DataType::Int32 - ); - assert_eq!( - schema - .field_with_name(SOURCE_ROW_ID_COLUMN) - .unwrap() - .data_type(), - &DataType::UInt64 - ); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - } - - #[tokio::test] - async fn test_with_no_data_skips_initial_refresh() { - let conn = people_db().await; - let view = conn - .create_materialized_view("empty", "people") - .with_no_data(true) - .execute() - .await - .unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - } - - #[tokio::test] - async fn test_create_and_refresh_async_jobs() { - let conn = people_db().await; - let create_job = conn - .create_materialized_view("async_view", "people") - .with_no_data(true) - .execute_async() - .await - .unwrap(); - assert!(create_job.id().is_none()); - create_job.wait().await.unwrap(); - - let view = conn.open_materialized_view("async_view").await.unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - - let refresh_job = view.refresh().execute_async().await.unwrap(); - assert!(refresh_job.id().is_none()); - let result = refresh_job.wait().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 3); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - - let drop_job = conn - .drop_materialized_view_async("async_view", &[]) - .await - .unwrap(); - assert!(drop_job.id().is_none()); - drop_job.wait().await.unwrap(); - assert!(conn.open_table("async_view").execute().await.is_err()); - } - - #[tokio::test] - async fn test_drop_materialized_view_rejects_plain_tables() { - let conn = people_db().await; - let error = conn - .drop_materialized_view("people", &[]) - .await - .unwrap_err(); - assert!(matches!(error, Error::NotAMaterializedView { .. })); - - conn.create_materialized_view("drop_me", "people") - .with_no_data(true) - .execute() - .await - .unwrap(); - conn.drop_materialized_view("drop_me", &[]).await.unwrap(); - assert!(conn.open_table("drop_me").execute().await.is_err()); - } - - #[tokio::test] - async fn test_remote_query_quotes_resource_identifiers() { - let conn = people_db().await; - let query = conn - .create_materialized_view("unused", "odd\"source") - .source_namespace(vec!["raw data".into()]) - .select([("double\"age", "age * 2")]) - .only_if("age >= 18") - .limit(10) - .query() - .unwrap(); - assert_eq!( - query, - "SELECT age * 2 AS \"double\"\"age\" FROM \"raw data\".\"odd\"\"source\" WHERE age >= 18 LIMIT 10" + source_sql_identifier(&namespace, "source", false), + "\"analytics\".\"daily\".\"source\"" ); } - #[tokio::test] - async fn native_vector_builder_preserves_source_and_rejects_extra_clauses() { - let conn = people_db().await; - let source = VectorSource { - kind: VectorSourceKind::Dedup, - dataset_version: 7, - column: "vector'field".into(), - distance_threshold: "4.0".into(), - }; - let builder = || { - conn.create_materialized_view("clean", "odd.source'name") - .source_namespace(vec!["raw space".into()]) - .vector_source(source.clone()) - }; - let query = builder().query().unwrap(); - let definition = MaterializedViewDefinition::from_sql(&query).unwrap(); - assert_eq!(definition.source_table, "odd.source'name"); - assert_eq!(definition.source_namespace, ["raw space"]); - assert_eq!(definition.vector_source.as_deref(), Some(&source)); - assert!(definition.selects_star()); - assert!(builder().select([("id", "id")]).query().is_err()); - assert!(builder().only_if("true").query().is_err()); - assert!(builder().limit(1).query().is_err()); - let invalid_source = VectorSource { - dataset_version: 0, - ..source - }; - assert!( - conn.create_materialized_view("bad", "people") - .vector_source(invalid_source) - .query() - .is_err() - ); - } - - /// No projection selects every source column, expanded now: the schema - /// captured at creation is the definition. - #[tokio::test] - async fn test_default_projection_captures_the_source_schema() { - let conn = people_db().await; - let view = conn - .create_materialized_view("copy", "people") - .execute() - .await - .unwrap(); - assert_eq!( - view.definition() - .projections - .iter() - .map(|p| p.output.as_str()) - .collect::>(), - vec!["name", "age"] - ); - } - - #[tokio::test] - async fn test_unknown_column_fails_at_create_time() { - let conn = people_db().await; - let err = conn - .create_materialized_view("bad", "people") - .select([("x", "missing + 1")]) - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "x")); - let names = conn.table_names().execute().await.unwrap(); - assert!(!names.contains(&"bad".to_string())); - } - - #[tokio::test] - async fn test_unknown_filter_column_fails_at_create_time() { - let err = declare_err(|b| b.only_if("missing > 1")).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("missing"))); - } - - #[tokio::test] - async fn test_duplicate_output_is_rejected() { - let err = declare_err(|b| b.select([("dup", "age"), ("dup", "age + 1")])).await; - assert!(matches!(err, Error::ColumnAlreadyExists { name } if name == "dup")); - } - - #[tokio::test] - async fn test_reserved_output_name_is_rejected() { - let err = declare_err(|b| b.select([(SOURCE_ROW_ID_COLUMN, "age")])).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("reserved"))); - } - - #[tokio::test] - async fn test_missing_source_fails() { - let conn = connect("memory://").execute().await.unwrap(); - let err = conn - .create_materialized_view("v", "nope") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::TableNotFound { .. })); - } - - /// Provenance has to survive source compactions and updates, and stable - /// row ids cannot be enabled after a table exists -- so the requirement - /// is checked at the last moment the caller can still act on it. - #[tokio::test] - async fn test_source_without_stable_row_ids_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, [1, 2])).unwrap(); - conn.create_table("plain", batch).execute().await.unwrap(); - - let err = conn - .create_materialized_view("v", "plain") - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("stable row ids")) - ); - assert!( - !conn - .table_names() - .execute() - .await - .unwrap() - .contains(&"v".to_string()) - ); - } - - #[tokio::test] - async fn test_name_collision_fails() { - let conn = people_db().await; - let err = conn - .create_materialized_view("people", "people") - .execute() - .await - .unwrap_err(); - assert!(matches!(err, Error::TableAlreadyExists { .. })); - } - - #[tokio::test] - async fn test_a_plain_table_is_not_a_view() { - let conn = people_db().await; - let table = conn.open_table("people").execute().await.unwrap(); - let err = MaterializedView::from_table(table).await.unwrap_err(); - assert!(matches!(err, Error::NotAMaterializedView { name } if name == "people")); - - let err = conn.open_materialized_view("people").await.unwrap_err(); - assert!(matches!(err, Error::NotAMaterializedView { .. })); - } - - /// The reason the kind is tagged: a definition written by a newer version - /// reads back as a view this one cannot refresh, not as a plain table. - #[tokio::test] - async fn test_unrecognized_kind_is_refused_by_name() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - let table = conn.open_table("v").execute().await.unwrap(); - table - .as_native() - .unwrap() - .replace_schema_metadata(HashMap::from([( - DEFINITION_META_KEY.to_string(), - r#"{"kind": "join"}"#.to_string(), - )])) - .await - .unwrap(); - - let err = conn.open_materialized_view("v").await.unwrap_err(); - assert!(matches!(err, Error::NotSupported { message } if message.contains("join"))); - } - - #[tokio::test] - async fn test_list_reports_views_and_only_views() { - let conn = people_db().await; - conn.create_materialized_view("adults", "people") - .only_if("age >= 18") - .execute() - .await - .unwrap(); - - let views = conn.list_materialized_views().await.unwrap(); - assert_eq!(views, vec!["adults"]); - let view = conn.open_materialized_view("adults").await.unwrap(); - assert_eq!(view.definition().filter.as_deref(), Some("age >= 18")); - } - - /// The creation option outranks a connection configured to create - /// unstable tables: the view still gets stable row ids, on the same - /// store (no fork -- the table must be reachable through the - /// connection afterwards). - #[tokio::test] - async fn test_view_is_stable_despite_connection_override() { - let conn = connect("memory://") - .storage_options([("new_table_enable_stable_row_ids", "false")]) - .execute() - .await - .unwrap(); - let batch = record_batch!(("x", Int32, [1])).unwrap(); - conn.create_table("src", batch) - .storage_option("new_table_enable_stable_row_ids", "true") - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("v", "src") - .execute() - .await - .unwrap(); - let stable = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .manifest - .uses_stable_row_ids(); - assert!(stable); - conn.open_materialized_view("v").await.unwrap(); - } - - /// A committed filter has to be usable as a predicate. - #[tokio::test] - async fn test_non_boolean_filter_is_rejected() { - let err = declare_err(|b| b.only_if("age + 1")).await; - assert!(matches!(err, Error::InvalidInput { message } if message.contains("boolean"))); - } - - /// Nested references stay dotted paths; resolution is by root field. - #[tokio::test] - async fn test_struct_columns_can_be_declared() { - use arrow_array::{ArrayRef, Int32Array, StructArray}; - - let conn = connect("memory://").execute().await.unwrap(); - let ages = StructArray::from(vec![( - Arc::new(ArrowField::new("age", DataType::Int32, false)), - Arc::new(Int32Array::from(vec![36, 17])) as ArrayRef, + #[test] + fn sql_definition_defaults_share_one_metadata_value() { + let metadata = HashMap::from([( + DEFINITION_META_KEY.to_string(), + definition_metadata_from_sql_with_defaults( + "SELECT * FROM source", + "db", + &["analytics".to_string()], + ), )]); - let batch = - arrow_array::RecordBatch::try_from_iter(vec![("metadata", Arc::new(ages) as ArrayRef)]) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("ages", "people") - .select([("age", "metadata.age")]) - .only_if("metadata.age >= 18") - .execute() - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - assert_eq!( - schema.field_with_name("age").unwrap().data_type(), - &DataType::Int32 - ); - } - - /// A newer-kind view must not disappear from the listing. - #[tokio::test] - async fn test_unrecognized_kind_is_listed_by_name() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - let table = conn.open_table("v").execute().await.unwrap(); - table - .as_native() + let definition = read_definition_sql_with_defaults(&metadata) .unwrap() - .replace_schema_metadata(HashMap::from([( - DEFINITION_META_KEY.to_string(), - r#"{"kind": "join"}"#.to_string(), - )])) - .await .unwrap(); - - let views = conn.list_materialized_views().await.unwrap(); - assert_eq!(views, vec!["v"]); - } - - /// A definition must evaluate identically across refreshes; anything - /// less makes incremental maintenance a mix of evaluations. - #[tokio::test] - async fn test_volatile_and_unstable_expressions_are_rejected() { - let conn = people_db().await; - for expression in [ - "random()", - "now()", - "version()", - "arrow_typeof(age)", - "arrow_metadata(age, 'k')", - ] { - let err = conn - .create_materialized_view("bad", "people") - .select([("x", expression)]) - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidExpression { message, .. } - if message.contains("not immutable")), - "{expression} was not rejected" - ); - } - for filter in ["age > random() * 100", "age >= 0 and now() is not null"] { - let err = conn - .create_materialized_view("bad", "people") - .only_if(filter) - .execute() - .await - .unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("not immutable")), - "{filter} was not rejected" - ); - } - } - - /// A column projected as itself stays the column it was: blob discovery - /// and the blob APIs key off field metadata, which a bare rebuild of the - /// field would drop. - #[tokio::test] - async fn test_identity_projection_keeps_field_metadata() { - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("id", DataType::Int32, true), - crate::blob("payload", true), - ], - HashMap::new(), - )); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("v", "src") - .execute() - .await - .unwrap(); - let view_schema = view.table().schema().await.unwrap(); - - let payload = view_schema.field_with_name("payload").unwrap(); - assert!( - crate::blob::is_blob(payload), - "default projection dropped the blob marker: {:?}", - payload.metadata() - ); + assert_eq!(definition.query, "SELECT * FROM source"); + assert_eq!(definition.default_database.as_deref(), Some("db")); assert_eq!( - view.table().blob_columns().await.unwrap(), - vec!["payload".to_string()], - "blob discovery no longer finds the projected column" - ); - assert!(view_schema.metadata().contains_key(DEFINITION_META_KEY)); - // Structural declarations describe how a table is written; a view is - // written by refresh, and its fields are always nullable. - assert!(!view_schema.metadata().contains_key("lance:primary_key")); - - // A computed column is a new value and carries no source metadata. - let computed = conn - .create_materialized_view("c", "src") - .select([("payload", "payload"), ("n", "id + 1")]) - .execute() - .await - .unwrap(); - let computed_schema = computed.table().schema().await.unwrap(); - assert!(crate::blob::is_blob( - computed_schema.field_with_name("payload").unwrap() - )); - assert!( - computed_schema - .field_with_name("n") - .unwrap() - .metadata() - .is_empty() + definition.default_namespace_path, + Some(vec!["analytics".to_string()]) ); } - /// A nested column projected straight through is still that column, and a - /// declaration buried in a struct child binds as hard as one on top. #[tokio::test] - async fn test_nested_projection_metadata_and_declarations() { - // The schema below carries the legacy v1 blob marker, which Lance only - // allows writing at file version <= 2.1. - let conn = connect("memory://") - .storage_options([( - crate::database::listing::OPT_NEW_TABLE_STORAGE_VERSION, - "2.1", - )]) - .execute() - .await - .unwrap(); - let payload = crate::blob("payload", true).with_metadata(HashMap::from([ - ("lance-encoding:blob".to_string(), "true".to_string()), - ( - "lance-schema:unenforced-primary-key".to_string(), - "0".to_string(), - ), - ])); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, true), - ArrowField::new("meta", DataType::Struct(vec![payload].into()), true), - ])); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - // A nested path is a direct projection: the leaf's metadata comes with - // it, so the blob stays a blob rather than a plain struct. - let lifted = conn - .create_materialized_view("lifted", "src") - .select([("payload", "meta.payload")]) - .execute() - .await - .unwrap(); - let field = lifted.table().schema().await.unwrap(); - let field = field.field_with_name("payload").unwrap().clone(); - assert_eq!( - field.metadata().get("lance-encoding:blob"), - Some(&"true".to_string()), - "nested projection lost the leaf's metadata" - ); - assert!( - !field - .metadata() - .contains_key("lance-schema:unenforced-primary-key"), - "a structural declaration rode along" - ); - - // Projecting the struct whole must not carry the child's declaration - // out to a view whose fields are nullable. - let whole = conn - .create_materialized_view("whole", "src") - .select([("meta", "meta")]) - .execute() - .await - .unwrap(); - let schema = whole.table().schema().await.unwrap(); - let DataType::Struct(children) = schema.field_with_name("meta").unwrap().data_type() else { - panic!("meta is not a struct"); - }; - let child = children.iter().find(|c| c.name() == "payload").unwrap(); - assert!( - !child - .metadata() - .contains_key("lance-schema:unenforced-primary-key"), - "a nested declaration survived: {:?}", - child.metadata() - ); - assert_eq!( - child.metadata().get("lance-encoding:blob"), - Some(&"true".to_string()) - ); - } - - /// A map's entries are fields like any other, and a declaration on one - /// binds the view's writes just as hard as one on top. - #[tokio::test] - async fn test_map_declarations_are_stripped() { + async fn refresh_always_rebuilds_without_source_row_ids() { let conn = connect("memory://").execute().await.unwrap(); - let value = - ArrowField::new("value", DataType::Utf8, false).with_metadata(HashMap::from([( - "lance-schema:unenforced-clustering-key:position".to_string(), - "1".to_string(), - )])); - let entries = ArrowField::new( - "entries", - DataType::Struct(vec![ArrowField::new("key", DataType::Utf8, false), value].into()), - false, - ); - let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( - "props", - DataType::Map(Arc::new(entries), false), - true, - )])); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("view", "src") - .execute() - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let DataType::Map(entries, _) = schema.field_with_name("props").unwrap().data_type() else { - panic!("props is not a map"); - }; - let DataType::Struct(children) = entries.data_type() else { - panic!("map entries are not a struct"); - }; - let value = children.iter().find(|c| c.name() == "value").unwrap(); - assert!( - !value - .metadata() - .contains_key("lance-schema:unenforced-clustering-key:position"), - "a declaration survived inside a map: {:?}", - value.metadata() - ); - } - - /// A computed column is declared by field metadata. Projecting one -- - /// as itself or under an alias -- must carry its description without its - /// declaration, which the target table would reject as foreign. - #[tokio::test] - async fn test_view_over_a_computed_column() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("id", Int32, [1, 2])).unwrap(); let source = conn - .create_table("src", batch) - .write_options(stable_row_ids()) + .create_table("src", record_batch!(("x", Int32, [1, 2])).unwrap()) .execute() .await .unwrap(); - source - .add_columns() - .computed("doubled", "id * 2") - .execute() - .await - .unwrap(); - - // Default projection reaches the computed column too. - let whole = conn - .create_materialized_view("whole", "src") - .execute() - .await - .unwrap(); - let schema = whole.table().schema().await.unwrap(); - let field = schema.field_with_name("doubled").unwrap(); - assert!( - !field - .metadata() - .keys() - .any(|k| k.starts_with("computed_column")), - "a computed-column declaration rode along: {:?}", - field.metadata() - ); - - // And under an alias. - conn.create_materialized_view("aliased", "src") - .select([("twice", "doubled")]) - .execute() - .await - .unwrap(); - } - - /// Embedding configuration names columns. It comes along only for the - /// columns a view actually projects, under the names the view gives them. - #[tokio::test] - async fn test_embedding_config_follows_the_projection() { - let config = r#"[{"name":"f","model":{},"source_column":"text","vector_column":"vec"}]"#; - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vec", DataType::Float32, true), - ], - HashMap::from([("embedding_functions".to_string(), config.to_string())]), - )); - conn.create_empty_table("src", schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let carried = |view: &MaterializedView| { - let view = view.table().clone(); - async move { - view.schema() - .await - .unwrap() - .metadata() - .get("embedding_functions") - .cloned() - } - }; - - // Both columns projected as themselves: kept as it stands. - let whole = conn - .create_materialized_view("whole", "src") - .execute() - .await - .unwrap(); - let kept = carried(&whole).await.expect("config dropped"); - assert!(kept.contains(r#""source_column":"text""#), "{kept}"); - assert!(kept.contains(r#""vector_column":"vec""#), "{kept}"); - - // Only the source column: the configuration names a vector column the - // view does not have, so it describes nothing and goes. - let partial = conn - .create_materialized_view("partial", "src") - .select([("text", "text")]) - .execute() - .await - .unwrap(); - assert_eq!(carried(&partial).await, None); - - // Renamed: the configuration follows the names the view uses. - let renamed = conn - .create_materialized_view("renamed", "src") - .select([("body", "text"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - let remapped = carried(&renamed).await.expect("config dropped"); - assert!(remapped.contains(r#""source_column":"body""#), "{remapped}"); - assert!( - remapped.contains(r#""vector_column":"embedding""#), - "{remapped}" - ); - - // The Node bindings spell the same configuration in camelCase, and - // the Rust definition names the destination `dest_column`. - for (config, source_key, dest_key) in [ - ( - r#"[{"name":"f","model":{},"sourceColumn":"text","vectorColumn":"vec"}]"#, - "sourceColumn", - "vectorColumn", - ), - ( - r#"[{"name":"f","model":{},"source_column":"text","dest_column":"vec"}]"#, - "source_column", - "dest_column", - ), - ] { - let schema = Arc::new(ArrowSchema::new_with_metadata( - vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vec", DataType::Float32, true), - ], - HashMap::from([("embedding_functions".to_string(), config.to_string())]), - )); - let name = format!("src_{source_key}"); - conn.create_empty_table(&name, schema) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view(format!("v_{source_key}"), &name) - .select([("body", "text"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - let carried = carried(&view).await.expect("config dropped"); - assert!( - carried.contains(&format!(r#""{source_key}":"body""#)), - "{carried}" - ); - assert!( - carried.contains(&format!(r#""{dest_key}":"embedding""#)), - "{carried}" - ); - } - - // A computed column is not the source column under another name. - let computed = conn - .create_materialized_view("computed", "src") - .select([("body", "upper(text)"), ("embedding", "vec")]) - .execute() - .await - .unwrap(); - assert_eq!(carried(&computed).await, None); - - // One column projected twice is two columns in the view, and the - // configuration has to describe both rather than whichever came last. - let twice = conn - .create_materialized_view("twice", "src") - .select([("body", "text"), ("a", "vec"), ("b", "vec")]) - .execute() - .await - .unwrap(); - let carried = carried(&twice).await.expect("config dropped"); - let entries: Vec = serde_json::from_str(&carried).unwrap(); - let mut vectors: Vec<&str> = entries - .iter() - .filter_map(|e| e["vector_column"].as_str()) - .collect(); - vectors.sort_unstable(); - assert_eq!(vectors, ["a", "b"], "{carried}"); - } - - /// The native Rust producer records embeddings as column definitions - /// rather than as `embedding_functions`, and a query embeds its own text - /// through them. They are positional, so the view's list covers every one - /// of its fields. - #[tokio::test] - async fn test_native_column_definitions_follow_the_projection() { - let conn = connect("memory://").execute().await.unwrap(); - let rich = crate::table::TableDefinition::new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("text", DataType::Utf8, true), - ArrowField::new("vector", DataType::Float32, true), - ])), - vec![ - ColumnDefinition { - kind: ColumnKind::Physical, - }, - ColumnDefinition { - kind: ColumnKind::Embedding(EmbeddingDefinition::new( - "text", - "model", - Some("vector"), - )), - }, - ], - ) - .into_rich_schema(); - conn.create_empty_table("src", rich) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("view", "src") - .select([("body", "text"), ("embedding", "vector")]) + .create_materialized_view("copy", "src") .execute() .await .unwrap(); - let schema = view.table().schema().await.unwrap(); - let raw = schema - .metadata() - .get(COLUMN_DEFINITIONS_META_KEY) - .expect("the view dropped the native column definitions"); - let definitions: Vec = serde_json::from_str(raw).unwrap(); - assert_eq!( - definitions.len(), - schema.fields().len(), - "column definitions are positional" - ); - let ColumnKind::Embedding(embedding) = &definitions[1].kind else { - panic!("the embedding column came back physical: {raw}"); - }; - assert_eq!(embedding.source_column, "body"); - assert_eq!(embedding.dest_column.as_deref(), Some("embedding")); - assert_eq!(embedding.embedding_name, "model"); - assert!(matches!(definitions[0].kind, ColumnKind::Physical)); - assert!(matches!(definitions[2].kind, ColumnKind::Physical)); - // Without the column the function reads, the view cannot recompute - // the embedding, so it carries no definition for it. - let partial = conn - .create_materialized_view("partial", "src") - .select([("embedding", "vector")]) - .execute() - .await - .unwrap(); - assert_eq!( - partial - .table() + assert_eq!(view.table().count_rows(None).await.unwrap(), 2); + assert!( + view.table() .schema() .await .unwrap() - .metadata() - .get(COLUMN_DEFINITIONS_META_KEY), - None - ); - } - - /// A scan takes the cap as i64, so a larger one is refused where it is - /// declared rather than at the refresh that cannot run it. What a cap of - /// zero means is a refresh question, tested there. - #[tokio::test] - async fn test_limit_above_i64_max_is_refused_at_creation() { - let conn = people_db().await; - let err = conn - .create_materialized_view("too_big", "people") - .limit(i64::MAX as u64 + 1) - .execute() - .await - .unwrap_err(); - assert!( - matches!(&err, Error::InvalidInput { message } if message.contains("exceeds the maximum")), - "got {err:?}" + .field_with_name("__source_row_id") + .is_err() ); - // The boundary itself is accepted. - conn.create_materialized_view("at_max", "people") - .limit(i64::MAX as u64) + source + .add(record_batch!(("x", Int32, [3])).unwrap()) .execute() .await .unwrap(); - } + let refreshed = view.refresh().execute().await.unwrap(); + assert_eq!(refreshed.mode, RefreshMode::Rebuild); + assert_eq!(refreshed.rows_written, 3); + assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - #[tokio::test] - async fn test_drop_is_drop_table() { - let conn = people_db().await; - conn.create_materialized_view("v", "people") - .execute() - .await - .unwrap(); - conn.drop_table("v", &[]).await.unwrap(); - assert!(conn.list_materialized_views().await.unwrap().is_empty()); - } + let unchanged = view.refresh().execute().await.unwrap(); + assert_eq!(unchanged.mode, RefreshMode::Rebuild); + assert_eq!(unchanged.rows_written, 3); - /// The public declaration contract: prepare validates the source and - /// create consumes the declaration into a verified view table; a - /// source that cannot anchor refresh and a target outside the - /// source's database are both refused. - #[tokio::test] - async fn prepare_and_create_bind_the_declaration_lifecycle() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("id", Int32, [1, 2]), ("value", Int32, [3, 4])).unwrap(); - let source = conn - .create_table("src", batch.clone()) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let projections = [ - ("id".to_string(), "id".to_string()), - ("double".to_string(), "value * 2".to_string()), - ]; - let prepared = prepare_declaration(&source, Some(&projections), Some("value > 0"), None) - .await - .unwrap(); - assert_eq!(prepared.definition().source_table, "src"); - - let view = prepared.create("v").await.unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["id", "double", SOURCE_ROW_ID_COLUMN]); - assert!(schema.metadata().contains_key(DEFINITION_META_KEY)); - - // The same call rejects a source without stable row ids, so an - // external creation path cannot skip the check. - conn.create_table("plain", batch).execute().await.unwrap(); - let plain = conn.open_table("plain").execute().await.unwrap(); - let err = prepare_declaration(&plain, None, None, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("stable row ids"), "{err}"); - - // A handle whose location does not resolve back through its name is - // refused: the definition would record a name reaching other data. - let plain_uri = plain - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - let masquerade = conn - .open_table("src") - .location(plain_uri) - .execute() - .await - .unwrap(); - let err = prepare_declaration(&masquerade, None, None, None) - .await - .unwrap_err(); - assert!( - err.to_string().contains("does not resolve to itself"), - "{err}" - ); - - // A table created at a custom location is refused outright: its - // recorded name reaches nothing at the database root, so the - // canonical reopen fails before any URI comparison. - let custom = conn - .create_table( - "custom_loc", - record_batch!(("id", Int32, [1, 2]), ("value", Int32, [3, 4])).unwrap(), - ) - .location("memory://elsewhere/custom_loc") - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - let err = prepare_declaration(&custom, None, None, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("custom_loc"), "{err}"); - } - - /// A view declared over a namespaced source records that namespace, and - /// refresh resolves the source through it -- the coordinate round-trips. - #[tokio::test] - async fn a_namespaced_source_round_trips_through_refresh() { - use lance_namespace::models::CreateNamespaceRequest; - - let tmp = tempfile::tempdir().unwrap(); - let mut properties = std::collections::HashMap::new(); - properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string()); - let conn = crate::connect_namespace("dir", properties) - .execute() - .await - .unwrap(); - conn.create_namespace(CreateNamespaceRequest { - id: Some(vec!["ns".into()]), - ..Default::default() - }) - .await - .unwrap(); - - let batch = record_batch!( - ("name", Utf8, ["ada", "grace", "alan"]), - ("age", Int32, [36, 85, 41]) - ) - .unwrap(); - conn.create_table("people", batch) - .namespace(vec!["ns".to_string()]) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - // A decoy of the same name at the root: resolving the source at the - // wrong namespace materializes one row here instead of three. - let decoy = record_batch!(("name", Utf8, ["mallory"]), ("age", Int32, [42])).unwrap(); - conn.create_table("people", decoy) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap(); - - let view = conn - .create_materialized_view("adults", "people") - .with_no_data(true) - .namespace(vec!["ns".to_string()]) - .source_namespace(vec!["ns".to_string()]) - .select([("name", "name")]) - .only_if("age >= 18") - .execute() - .await - .unwrap(); - - assert_eq!(view.definition().source_table, "people"); - assert_eq!(view.definition().source_namespace, vec!["ns".to_string()]); - assert_eq!(view.table().namespace(), &["ns"]); - - // Refresh resolves the source at the recorded namespace, not at root. - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - } - - fn definition(source_namespace: Vec) -> MaterializedViewDefinition { - MaterializedViewDefinition { - vector_source: None, - source_table: "people".to_string(), - source_namespace, - lateral: None, - projections: vec![ViewProjection { - output: "name".to_string(), - expression: "name".to_string(), - }], - filter: None, - group_by: Vec::new(), - limit: None, - } - } - - fn read(stored: impl Into) -> Result> { - read_definition(&HashMap::from([( - DEFINITION_META_KEY.to_string(), - stored.into(), - )])) - } - - /// What is stored is the query, under a format number; the same query - /// reads back whatever namespace the source sits in. - #[test] - fn the_stored_layout_is_the_canonical_query() { - for (namespace, query) in [ - (Vec::new(), "SELECT name FROM people"), - (vec!["ns".to_string()], "SELECT name FROM ns.people"), - ] { - let stored = definition_to_metadata(&definition(namespace.clone())).unwrap(); - let value: serde_json::Value = serde_json::from_str(&stored).unwrap(); - assert_eq!( - value, - serde_json::json!({"kind": "query", "format": 1, "query": query}) - ); - assert_eq!( - read(stored).unwrap(), - Some(StoredDefinition::Query(definition(namespace))) - ); - } - } - - /// The structured layout written before the format number still reads, - /// under both of its kind tags, and a tag that disagrees with its - /// namespace is an error: under `select` old readers resolved it at root. - #[test] - fn legacy_layouts_read_back() { - let legacy = |kind: &str, namespace: Vec<&str>| { - serde_json::json!({ - "kind": kind, - "source_table": "people", - "source_namespace": namespace, - "projections": [{"output": "name", "expression": "name"}], - "inputs": ["name"], - }) - .to_string() - }; - assert_eq!( - read(legacy(SELECT_KIND, vec![])).unwrap(), - Some(StoredDefinition::Query(definition(Vec::new()))) - ); - assert_eq!( - read(legacy(NAMESPACED_SELECT_KIND, vec!["ns"])).unwrap(), - Some(StoredDefinition::Query(definition(vec!["ns".into()]))) - ); - assert!( - read(r#"{"kind":"select","source_table":"people","projections":[]}"#) - .unwrap() - .is_some(), - "a pre-namespace definition carries no namespace key" - ); - for (kind, namespace) in [(SELECT_KIND, vec!["ns"]), (NAMESPACED_SELECT_KIND, vec![])] { - let err = read(legacy(kind, namespace)).unwrap_err(); - assert!( - err.to_string() - .contains("does not match its source namespace"), - "kind '{kind}': {err}" - ); - } - } - - /// A newer writer's definition is reported as such, never guessed at, - /// and a definition that is not a definition at all is an error rather - /// than a plain table. - #[test] - fn a_newer_format_is_reported_not_guessed() { - assert_eq!( - read(r#"{"format":5,"query":"SELECT name FROM people"}"#).unwrap(), - Some(StoredDefinition::Newer { format: "5".into() }) - ); - assert_eq!( - read(r#"{"kind":"join"}"#).unwrap(), - Some(StoredDefinition::Newer { - format: "kind 'join'".into() - }) - ); - for stored in [ - "{}", - r#"{"format":"one"}"#, - r#"{"format":1}"#, - r#"{"format":1,"query":"SELECT name FROM people GROUP BY name"}"#, - r#"{"format":2,"query":"SELECT name FROM people ORDER BY name"}"#, - ] { - assert!(read(stored).is_err(), "{stored}"); - } - } - - /// Planning records what the query reads of the source: a nested path - /// as itself, a read through an unnest alias as the list column. - #[test] - fn planning_records_the_source_columns_read() { - let element = DataType::Struct( - vec![ - ArrowField::new("chunk", DataType::Utf8, true), - ArrowField::new("ordinal", DataType::Int32, true), - ] - .into(), - ); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int64, false), - ArrowField::new( - "meta", - DataType::Struct(vec![ArrowField::new("title", DataType::Utf8, true)].into()), - true, - ), - ArrowField::new( - "chunks", - DataType::List(Arc::new(ArrowField::new("item", element, true))), - true, - ), - ])); - let planned = plan( - schema.clone(), - &MaterializedViewDefinition::from_sql( - "SELECT id, meta.title, c.chunk FROM docs, UNNEST(chunks) AS c WHERE c.ordinal < 3", - ) - .unwrap(), - None, - ) - .unwrap(); - assert_eq!(planned.inputs, ["chunks", "id", "meta.title"]); - assert_eq!( - planned - .fields - .iter() - .map(|f| f.name().as_str()) - .collect::>(), - ["id", "title", "chunk"] - ); - assert_eq!( - planned.fields[2].data_type(), - &DataType::Utf8, - "the element's field is read through the alias" - ); - - let star = plan( - schema, - &MaterializedViewDefinition::from_sql("SELECT * FROM docs, UNNEST(chunks) AS c") - .unwrap(), - None, - ) - .unwrap(); - assert_eq!( - star.definition.to_sql(), - "SELECT id, meta, c FROM docs, UNNEST(chunks) AS c" - ); - let err = plan( - Arc::new(ArrowSchema::new(vec![ArrowField::new( - "id", - DataType::Int64, - false, - )])), - &MaterializedViewDefinition::from_sql("SELECT id FROM docs, UNNEST(id) AS c").unwrap(), - None, - ) - .unwrap_err(); - assert!(err.to_string().contains("not a list"), "{err}"); - } - - /// A binding as the server records it: one Utf8 input over `input` - /// bound to a nullable parameter, one Int32 output named `output`, with - /// the exact schemas the durable contract requires. - pub fn test_binding(binding_id: &str, input: &str, output: &str) -> FunctionBinding { - let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("text", DataType::Utf8, true), - ])) - .unwrap(); - let output_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new(output, DataType::Int32, true), - ])) - .unwrap(); - let input_type = input_schema.fields[0].r#type.r#type.clone(); - let output_type = output_schema.fields[0].r#type.r#type.clone(); - FunctionBinding::from_json( - &serde_json::json!({ - "binding_id": binding_id, - "function": {"name": "embed", "version": "1", "object_id": "fixture", "location": "memory:///fixture", - "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs": [{ - "parameter": "text", "field_id": -1, "field_path": input, - "arrow_type": input_type, "nullable": true, - }], - "outputs": [{ - "result_field": "$value", "output_name": output, "output_field_id": -1, - "output_ordinal": 0, "arrow_type": output_type, "nullable": false, - }], - "input_schema": serde_json::to_value(input_schema).unwrap(), - "output_schema": serde_json::to_value(output_schema).unwrap(), - }) - .to_string(), - ) - .unwrap() - } - - /// A computed column as the server declares it on a table: bound to a - /// registered Function. - pub fn computed_field(name: &str, binding_id: &str, input: &str) -> ArrowField { - ArrowField::new(name, DataType::Int32, true).with_metadata( - crate::table::computed_columns::function_computed_column_metadata( - binding_id, - 0, - &[input.to_string()], - ), - ) - } - - pub async fn people(conn: &Connection) -> Table { - let batch = - record_batch!(("id", Int32, [1, 2, 3]), ("name", Utf8, ["a", "b", "c"])).unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap() - } - - /// `people` with both columns non-nullable, for nullability cases. - pub async fn strict_people(conn: &Connection) -> Table { - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - ArrowField::new("name", DataType::Utf8, false), - ])); - let batch = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1, 2, 3])), - Arc::new(arrow_array::StringArray::from(vec!["a", "b", "c"])), - ], - ) - .unwrap(); - conn.create_table("people", batch) - .write_options(stable_row_ids()) - .execute() - .await - .unwrap() - } - - async fn prepared_people(conn: &Connection) -> PreparedDeclaration { - let source = people(conn).await; - prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("name".to_string(), "name".to_string()), - ]), - None, - None, - ) - .await - .unwrap() - } - - #[tokio::test] - async fn a_computed_column_is_declared_null_with_its_binding() { - let conn = connect("memory://").execute().await.unwrap(); - let view = prepared_people(&conn) - .await - .with_computed_columns( - vec![(2, computed_field("emb", "fb_1", "name"))], - &[test_binding("fb_1", "name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["id", "name", "emb", SOURCE_ROW_ID_COLUMN]); - let declared: Vec = computed_columns(&schema) - .into_iter() - .map(|c| c.name) - .collect(); - assert_eq!(declared, ["emb"]); - let bindings = crate::table::computed_columns::function_bindings(&schema).unwrap(); - assert_eq!(bindings.len(), 1); - assert_eq!(bindings[0].binding_id(), "fb_1"); - // The stored definition is the query alone; bindings live beside it. - let stored: serde_json::Value = - serde_json::from_str(&schema.metadata()[DEFINITION_META_KEY]).unwrap(); - assert_eq!(stored["format"], 1); - assert_eq!(view.definition().projections.len(), 2); + source.delete("true").await.unwrap(); + let emptied = view.refresh().execute().await.unwrap(); + assert_eq!(emptied.mode, RefreshMode::Rebuild); + assert_eq!(emptied.rows_written, 0); assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!(conn.open_materialized_view("v").await.unwrap().name(), "v"); - } - - #[tokio::test] - async fn computed_column_declarations_are_validated() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let binding = test_binding("fb_1", "name", "emb"); - let fails = |prepared: PreparedDeclaration, - columns: Vec<(usize, ArrowField)>, - bindings: &[FunctionBinding]| { - prepared - .with_computed_columns(columns, bindings) - .err() - .map(|e| e.to_string()) - .expect("the declaration should be refused") - }; - let emb = |binding_id: &str| computed_field("emb", binding_id, "name"); - - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1").with_nullable(false))], - std::slice::from_ref(&binding), - ); - assert!(err.contains("must be nullable"), "{err}"); - - let plain = ArrowField::new("emb", DataType::Int32, true); - let err = fails( - prepared.clone(), - vec![(2, plain)], - std::slice::from_ref(&binding), - ); - assert!( - err.contains("does not carry a computed-column declaration"), - "{err}" - ); - - // The rest is the computed-column contract: a binding the field does - // not name, an output the binding does not map to this field, an - // input the view does not hold. - let err = fails( - prepared.clone(), - vec![(2, emb("fb_other"))], - std::slice::from_ref(&binding), - ); - assert!(err.contains("does not match binding 'fb_1'"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1"))], - &[test_binding("fb_1", "name", "different_output")], - ); - assert!(err.contains("different_output"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, emb("fb_1"))], - &[test_binding("fb_1", "bio", "emb")], - ); - assert!(err.contains("'bio'"), "{err}"); - - let err = fails( - prepared.clone(), - vec![(2, computed_field("name", "fb_1", "name"))], - &[test_binding("fb_1", "name", "name")], - ); - assert!(err.contains("already exists"), "{err}"); - let err = fails( - prepared.clone(), - vec![(2, computed_field(SOURCE_ROW_ID_COLUMN, "fb_1", "name"))], - &[test_binding("fb_1", "name", SOURCE_ROW_ID_COLUMN)], - ); - assert!(err.contains("reserved"), "{err}"); - let err = fails( - prepared.clone(), - vec![(7, emb("fb_1"))], - std::slice::from_ref(&binding), - ); - assert!( - err.contains("placed at 7, past the view's 2 columns"), - "{err}" - ); - let err = fails(prepared, Vec::new(), std::slice::from_ref(&binding)); - assert!(err.contains("at least one computed column"), "{err}"); - } - - /// A source column a computed column reads without the view projecting - /// it becomes an internal projection before the provenance column, with - /// the source's nullability; a projected column is read from its - /// projection. - /// `Some(&[])` is a declaration that projects nothing yet: a view of - /// computed columns alone, whose inputs `input_column` places. `None` - /// is `SELECT *`. The two must not collapse into each other. - #[tokio::test] - async fn an_empty_projection_list_is_not_a_star() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - - let mut prepared = prepare_declaration(&source, Some(&[]), None, None) - .await - .unwrap(); - assert!(prepared.definition().projections.is_empty()); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - let view = prepared - .with_computed_columns( - vec![(0, computed_field("emb", "fb_1", "__input_name"))], - &[test_binding("fb_1", "__input_name", "emb")], - ) - .unwrap() - .create("only") - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - assert_eq!( - view.definition().to_sql(), - "SELECT name AS __input_name FROM people" - ); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - let reopened = conn.open_materialized_view("only").await.unwrap(); - assert_eq!(reopened.definition(), view.definition()); - - let all = prepare_declaration(&source, None, None, None) - .await - .unwrap(); - assert!( - !all.definition().selects_star(), - "planning expands the star" - ); - let outputs: Vec<&str> = all - .definition() - .projections - .iter() - .map(|p| p.output.as_str()) - .collect(); - assert_eq!(outputs, ["id", "name"]); - } - - #[tokio::test] - async fn an_unprojected_input_becomes_an_internal_projection() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let mut prepared = prepare_declaration( - &source, - Some(&[("key".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - assert_eq!(prepared.input_column("id").unwrap(), "key"); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - assert_eq!(prepared.input_column("name").unwrap(), "__input_name"); - let err = prepared.input_column("missing").unwrap_err().to_string(); - assert!(err.contains("no column 'missing'"), "{err}"); - - let view = prepared - .with_computed_columns( - vec![(1, computed_field("emb", "fb_1", "__input_name"))], - &[test_binding("fb_1", "__input_name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let schema = view.table().schema().await.unwrap(); - let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - assert_eq!(names, ["key", "emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - let input = schema.field_with_name("__input_name").unwrap(); - assert_eq!(input.data_type(), &DataType::Utf8); - assert!( - !input.is_nullable(), - "the copy keeps the source's nullability" - ); - assert!(!schema.field_with_name("key").unwrap().is_nullable()); - let projections: Vec<(&str, &str)> = view - .definition() - .projections - .iter() - .map(|p| (p.output.as_str(), p.expression.as_str())) - .collect(); - assert_eq!(projections, [("key", "id"), ("__input_name", "name")]); - } - - /// Two outputs of one binding land at consecutive positions: each - /// insertion widens the range the next may take. - #[tokio::test] - async fn sibling_computed_columns_take_consecutive_positions() { - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let prepared = prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - let metadata = |ordinal: u32| { - crate::table::computed_columns::function_computed_column_metadata( - "fb_pair", - ordinal, - &["id".to_string()], - ) - }; - let left = ArrowField::new("left", DataType::Int32, true).with_metadata(metadata(0)); - let right = ArrowField::new("right", DataType::Int32, true).with_metadata(metadata(1)); - let input_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("value", DataType::Int32, true), - ])) - .unwrap(); - let output_schema = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - ArrowField::new("left", DataType::Int32, true), - ArrowField::new("right", DataType::Int32, true), - ])) - .unwrap(); - let int = input_schema.fields[0].r#type.r#type.clone(); - let binding = FunctionBinding::from_json( - &serde_json::json!({ - "binding_id": "fb_pair", - "function": {"name": "pair", "version": "1", "object_id": "fixture", "location": "memory:///fixture", - "manifest_digest": "sha256:7e22f815b6648e14f093a3979a8e5a2082fa773ebe1ec84b135cae7e84d6f8e6"}, - "inputs": [{"parameter": "value", "field_id": -1, "field_path": "id", - "arrow_type": int, "nullable": true}], - "outputs": [ - {"result_field": "left", "output_name": "left", "output_field_id": -1, - "output_ordinal": 0, "arrow_type": int, "nullable": false}, - {"result_field": "right", "output_name": "right", "output_field_id": -1, - "output_ordinal": 1, "arrow_type": int, "nullable": false}, - ], - "input_schema": serde_json::to_value(input_schema).unwrap(), - "output_schema": serde_json::to_value(output_schema).unwrap(), - }) - .to_string(), - ) - .unwrap(); - let view = prepared - .with_computed_columns(vec![(1, left), (2, right)], &[binding]) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view - .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "left", "right", SOURCE_ROW_ID_COLUMN]); - } - - /// A SQL declaration as `add_columns().computed()` records it. - pub fn sql_field( - name: &str, - data_type: DataType, - expression: &str, - inputs: &str, - ) -> ArrowField { - use crate::table::computed_columns::{ - COMPUTED_COLUMN_META_KEY, EXPRESSION_META_KEY, INPUTS_META_KEY, KIND_META_KEY, SQL_KIND, - }; - ArrowField::new(name, data_type, true).with_metadata(HashMap::from([ - (COMPUTED_COLUMN_META_KEY.to_string(), "true".to_string()), - (KIND_META_KEY.to_string(), SQL_KIND.to_string()), - (EXPRESSION_META_KEY.to_string(), expression.to_string()), - (INPUTS_META_KEY.to_string(), inputs.to_string()), - ])) - } - - /// A SQL declaration is re-planned at admission: it must parse against - /// the view, yield the declared type, and read the inputs it declares. - #[tokio::test] - async fn a_sql_declaration_is_planned_at_admission() { - let conn = connect("memory://").execute().await.unwrap(); - let fails = |prepared: PreparedDeclaration, field: ArrowField| { - prepared - .with_computed_columns(vec![(1, field)], &[]) - .err() - .map(|e| e.to_string()) - .expect("the declaration should be refused") - }; - let prepared = prepared_people(&conn).await; - let err = fails( - prepared.clone(), - sql_field("bad", DataType::Int32, "missing + 1", r#"["missing"]"#), - ); - assert!(err.contains("missing"), "{err}"); - let err = fails( - prepared.clone(), - sql_field("wide", DataType::Int64, "id + 1", r#"["id"]"#), - ); - assert!( - err.contains("declared as Int64 but its expression yields Int32"), - "{err}" - ); - let err = fails( - prepared.clone(), - sql_field("lying", DataType::Int32, "id + 1", r#"["name"]"#), - ); - assert!(err.contains("declares inputs"), "{err}"); - - let view = prepared - .with_computed_columns( - vec![(1, sql_field("next", DataType::Int32, "id + 1", r#"["id"]"#))], - &[], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view - .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "next", "name", SOURCE_ROW_ID_COLUMN]); - } - - /// Creation persists a declaration only when it re-plans and the data - /// carries no values for it, whichever door created the table. - #[tokio::test] - async fn a_created_table_cannot_carry_computed_values() { - let conn = connect("memory://").execute().await.unwrap(); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - sql_field("forged", DataType::Int32, "x + 1", r#"["x"]"#), - ])); - let filled = arrow_array::RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("forged", filled) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("cannot be written directly"), "{err}"); - assert!( - !conn - .table_names() - .execute() - .await - .unwrap() - .contains(&"forged".to_string()) - ); - - let unfilled = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::new_null(1)), - ], - ) - .unwrap(); - let table = conn - .create_table("declared", unfilled) - .execute() - .await - .unwrap(); - assert_eq!(table.refresh_column("forged").await.unwrap().rows_filled, 1); - - // A declaration with only its marker is broken, not absent. - let half = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( - crate::table::computed_columns::COMPUTED_COLUMN_META_KEY.to_string(), - "true".to_string(), - )])); - let batch = arrow_array::RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - half, - ])), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("half", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!( - err.contains("incomplete computed-column declaration"), - "{err}" - ); - - let bogus = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("x", DataType::Int32, false), - sql_field("bad", DataType::Int32, "missing + 1", r#"["missing"]"#), - ])); - let batch = arrow_array::RecordBatch::new_empty(bogus); - let err = conn - .create_table("bogus", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("missing"), "{err}"); - } - - /// The internal-input prefix is reserved for the declaration, like the - /// provenance column, so an alias cannot masquerade as an internal input. - #[tokio::test] - async fn the_internal_input_prefix_is_reserved() { - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let err = prepare_declaration( - &source, - Some(&[("__input_x".to_string(), "name".to_string())]), - None, - None, - ) - .await - .unwrap_err() - .to_string(); - assert!(err.contains("'__input_x' is reserved"), "{err}"); - } - - /// A computed column may not read another, through any path: a - /// Function bound to a child of a computed struct is refused like a SQL - /// declaration over it. - #[tokio::test] - async fn a_computed_column_cannot_read_a_computed_root() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let payload = sql_field( - "payload", - DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), - "named_struct('value', name)", - r#"["name"]"#, - ); - let err = prepared - .with_computed_columns( - vec![ - (2, payload), - (3, computed_field("emb", "fb_dependent", "payload.value")), - ], - &[test_binding("fb_dependent", "payload.value", "emb")], - ) - .err() - .map(|e| e.to_string()) - .expect("a computed root as a Function input should be refused"); - assert!(err.contains("reads computed column 'payload'"), "{err}"); - } - - /// The root check uses the canonical path parser: a quoted top-level - /// name containing a dot is one root, not two segments. - #[tokio::test] - async fn a_quoted_computed_root_is_still_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let prepared = prepared_people(&conn).await; - let payload = sql_field( - "payload.dot", - DataType::Struct(vec![ArrowField::new("value", DataType::Utf8, true)].into()), - "named_struct('value', name)", - r#"["name"]"#, - ); - let input = "`payload.dot`.value"; - let err = prepared - .with_computed_columns( - vec![ - (2, payload), - (3, computed_field("emb", "fb_dependent", input)), - ], - &[test_binding("fb_dependent", input, "emb")], - ) - .err() - .map(|e| e.to_string()) - .expect("a quoted computed root should be refused"); - assert!(err.contains("reads computed column 'payload.dot'"), "{err}"); - } - - /// Namespace-backed creation admits declarations by the same rule, and - /// refuses before the namespace records the table. - #[tokio::test] - async fn a_namespace_created_table_cannot_carry_computed_values() { - let tmp = tempfile::tempdir().unwrap(); - let mut properties = std::collections::HashMap::new(); - properties.insert("root".to_string(), tmp.path().to_str().unwrap().to_string()); - let conn = crate::connect_namespace("dir", properties) - .execute() - .await - .unwrap(); - let half = ArrowField::new("half", DataType::Int32, true).with_metadata(HashMap::from([( - crate::table::computed_columns::COMPUTED_COLUMN_META_KEY.to_string(), - "true".to_string(), - )])); - let batch = arrow_array::RecordBatch::try_new( - Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - half, - ])), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("malformed", batch) - .execute() - .await - .unwrap_err() - .to_string(); - assert!( - err.contains("incomplete computed-column declaration"), - "{err}" - ); - assert!(conn.table_names().execute().await.unwrap().is_empty()); - - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int32, false), - sql_field("next", DataType::Int32, "id + 1", r#"["id"]"#), - ])); - let filled = arrow_array::RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::from(vec![999])), - ], - ) - .unwrap(); - let err = conn - .create_table("forged", filled) - .execute() - .await - .unwrap_err() - .to_string(); - assert!(err.contains("cannot be written directly"), "{err}"); - - let unfilled = arrow_array::RecordBatch::try_new( - schema, - vec![ - Arc::new(arrow_array::Int32Array::from(vec![1])), - Arc::new(arrow_array::Int32Array::new_null(1)), - ], - ) - .unwrap(); - let table = conn - .create_table("declared", unfilled) - .execute() - .await - .unwrap(); - assert_eq!(table.refresh_column("next").await.unwrap().rows_filled, 1); - } - - /// A projected column keeps its nullability; a computed value is - /// nullable. - #[tokio::test] - async fn an_identity_projection_keeps_source_nullability() { - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let prepared = prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("n".to_string(), "name".to_string()), - ("next".to_string(), "id + 1".to_string()), - ]), - None, - None, - ) - .await - .unwrap(); - let nullable: Vec = prepared - .schema() - .fields() - .iter() - .map(|f| f.is_nullable()) - .collect(); - assert_eq!(nullable, [false, false, true, false]); } } diff --git a/rust/lancedb/src/materialized_view/differential.rs b/rust/lancedb/src/materialized_view/differential.rs deleted file mode 100644 index 396b4e65b..000000000 --- a/rust/lancedb/src/materialized_view/differential.rs +++ /dev/null @@ -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 { - 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 = 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 = 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 { - 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::>(); - 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::>() - .await - .unwrap(); - batches - .iter() - .flat_map(|batch| { - let ids = batch["id"].as_any().downcast_ref::().unwrap(); - let scores = batch["score"] - .as_any() - .downcast_ref::() - .unwrap(); - // Scores are integer-valued by construction; compare exactly. - (0..batch.num_rows()) - .map(|i| (ids.value(i), scores.value(i) as i32)) - .collect::>() - }) - .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> { - 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 = 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 { - let batches: Vec = 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::() - .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 { - let batches: Vec = 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::() - .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 = (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 = 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 = 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"); -} diff --git a/rust/lancedb/src/materialized_view/duplicate_pairs.rs b/rust/lancedb/src/materialized_view/duplicate_pairs.rs deleted file mode 100644 index d0dc219ad..000000000 --- a/rust/lancedb/src/materialized_view/duplicate_pairs.rs +++ /dev/null @@ -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 { - 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) -> Error { - Error::InvalidInput { - message: message.into(), - } -} - -pub(super) fn source( - args: &TableFunctionArgs, - kind: VectorSourceKind, -) -> Result<(Vec, 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::>>()?; - 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::>(); - 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 { - 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, -} - -#[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, - pub(super) rows: u64, -} - -impl DuplicatePairsRefreshPlan { - pub(super) fn identity(&self) -> Result { - 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)> { - 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, -) -> Result { - 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)> { - 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 { - 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, - expected: Option<&str>, -) -> Result { - 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::>>() - .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 -} diff --git a/rust/lancedb/src/materialized_view/grouped.rs b/rust/lancedb/src/materialized_view/grouped.rs index 095f20f11..3120e2b43 100644 --- a/rust/lancedb/src/materialized_view/grouped.rs +++ b/rust/lancedb/src/materialized_view/grouped.rs @@ -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, /// 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, - pub(super) covered: Option, -} - -/// 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, - pub(super) fragments: Option, -} - -impl IvfSegments { - pub(super) fn is_unindexed(&self, fragment_id: u64) -> Option { - 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, - bitmap: Option<&RoaringBitmap>, -) -> Option { - 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 { - 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 { - 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 = 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 = 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, ) -> Result { 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); - } -} diff --git a/rust/lancedb/src/materialized_view/grouped_units.rs b/rust/lancedb/src/materialized_view/grouped_units.rs deleted file mode 100644 index 7f408be4d..000000000 --- a/rust/lancedb/src/materialized_view/grouped_units.rs +++ /dev/null @@ -1,1368 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright The LanceDB Authors - -//! A grouped view keyed by one `ivf_partition(col)` refreshed in units, one -//! per index partition: the index already holds each partition's row ids, so -//! a unit reads only its rows, computes only its groups and writes its own -//! fragments, and a single commit replaces the view. Nothing is shuffled and -//! no worker holds more than one partition's groups. - -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; - -use arrow_array::Array; -use arrow_array::RecordBatch; -use arrow_array::cast::AsArray; -use arrow_array::types::UInt64Type; -use arrow_schema::{Schema as ArrowSchema, SchemaRef}; -use datafusion::catalog::default_table_source::provider_as_source; -use datafusion::physical_plan::SendableRecordBatchStream; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::prelude::SessionContext; -use futures::{StreamExt, TryStreamExt}; -use lance::Dataset; -use lance::dataset::transaction::{Operation, Transaction}; -use lance::dataset::{ - CommitBuilder, InsertBuilder, ProjectionRequest, WriteDestination, WriteMode, WriteParams, -}; -use lance_core::ROW_ID; -use lance_core::utils::address::RowAddress; -use lance_datafusion::exec::SessionContextExt; -use lance_index::metrics::NoOpMetricsCollector; -use lance_table::format::Fragment; -use serde::{Deserialize, Serialize}; - -use super::grouped::{self, IVF_PARTITION}; -use super::refresh::{ - RefreshMaterializedViewResult, RefreshMode, ensure_incarnation, ensure_no_mem_wal, open_source, - refresh_lock, stamp_watermark, to_view_batch, -}; -use super::{MaterializedViewDefinition, StoredDefinition}; -use crate::table::Table; -use crate::{Error, Result}; - -/// Row ids per `take`: bounds one unit's memory while it reads a partition. -const TAKE_CHUNK: usize = 8192; - -/// Rows per partition, and in total, the plan reads to see whether the index -/// assigns by its own centroids. A sample: [`write_grouped_unit`] proves the -/// rows it takes, and this only keeps a doomed refresh off the fleet. -const PROBE_PER_UNIT: usize = 8; -const PROBE_ROWS: usize = 256; - -/// How a grouped view can be refreshed in units. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct GroupedRefreshPlan { - /// Units `0..units`: unit `p` before the last holds index partition - /// `p`'s groups; the last holds the rows the index cannot place (null or - /// unplaceable vectors), grouped under a NULL key. One per index - /// partition, plus that one. - pub units: u32, - /// The source version every unit reads and the commit records. - pub source_version: u64, - /// The view version the plan was made against; the commit lands on it - /// or not at all. - pub view_version: u64, - /// The view this plan belongs to: its incarnation, which a drop and - /// recreate replaces. Two views of the same shape reach the same - /// counters, so without it a result written for one could be published - /// into the other. - pub incarnation: String, -} - -/// The column a grouped view's single `ivf_partition(column)` key names, if -/// that is its whole `GROUP BY`. -fn ivf_key_column(definition: &MaterializedViewDefinition) -> Option { - let [key] = definition.group_by.as_slice() else { - return None; - }; - let inner = key - .strip_prefix(&format!("{IVF_PARTITION}("))? - .strip_suffix(')')?; - let column = inner.trim().trim_matches('`'); - (!column.is_empty() && !column.contains(',')).then(|| column.to_string()) -} - -/// What a grouped refresh of `view` reads: its definition, source and binding. -struct Grouped { - view_ds: Dataset, - definition: MaterializedViewDefinition, - source: Dataset, - column: String, - inputs: Vec, - binding: grouped::IvfBinding, - segments: grouped::IvfSegments, -} - -async fn open_grouped(view: &Table, pinned: Option) -> Result> { - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { - message: "materialized views are supported only on local tables".into(), - })?; - view_native.dataset.reload().await?; - let view_ds = view_native.dataset.get().await?.as_ref().clone(); - let Some(StoredDefinition::Query(definition)) = - super::read_definition(&view_ds.schema().metadata)? - else { - return Ok(None); - }; - let Some(column) = ivf_key_column(&definition) else { - return Ok(None); - }; - let staging = super::read_staging(&view_ds.schema().metadata)?; - let source = open_source(view, &definition, staging.as_ref()).await?; - let source = match pinned { - Some(version) => source.checkout_version(version).await?, - None => source, - }; - ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; - let source_schema = Arc::new(ArrowSchema::from(source.schema())); - let planned = super::plan(source_schema, &definition, staging.as_ref())?; - let (binding, segments) = grouped::bind_column_with_segments(&source, &column).await?; - Ok(Some(Grouped { - view_ds, - definition: planned.definition, - source, - column, - inputs: planned.inputs, - binding, - segments, - })) -} - -/// The units `view` can be refreshed in against `source_version` (the -/// latest when `None`), or `None` when it is not a grouped view keyed by one -/// `ivf_partition(column)`. -pub async fn plan_grouped_refresh( - view: &Table, - source_version: Option, -) -> Result> { - let Some(grouped) = open_grouped(view, source_version).await? else { - return Ok(None); - }; - // The split reads indexed rows through the index and the rest by - // scanning what it does not cover, so an index that does not say what it - // covers cannot be split at all. The caller refreshes single-pass. - if grouped.segments.covered.is_none() { - return Ok(None); - } - // Without one there is nothing to bind a unit's result to, so the view - // is refreshed in one pass, where the rows never leave the process. - let Some(incarnation) = grouped - .view_ds - .schema() - .metadata - .get(super::INCARNATION_META_KEY) - .cloned() - else { - return Ok(None); - }; - // The split is the index's posting lists and the key is recomputed, so a - // view whose index assigns by something else has no units at all. - if !assigns_by_its_centroids(&grouped).await? { - return Ok(None); - } - Ok(Some(GroupedRefreshPlan { - units: grouped.binding.partitions + 1, - source_version: grouped.source.version().version, - view_version: grouped.view_ds.version().version, - incarnation, - })) -} - -/// Fragments one unit wrote, uncommitted, for [`commit_grouped_refresh`]. -/// Only [`write_grouped_unit`] makes one: it names the unit and the plan, -/// and the commit publishes a set of them as the view's whole contents, so -/// a result belonging to another plan, another view, or to a unit already -/// accounted for has to be refusable. Travels to whoever commits, so it -/// serializes, but its fields are not a caller's to write. -#[derive(Debug, Clone, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct WrittenUnit { - unit: u32, - plan: GroupedRefreshPlan, - fragments: Vec, - rows: u64, -} - -impl WrittenUnit { - /// The unit of the plan this result holds. - pub fn unit(&self) -> u32 { - self.unit - } - - /// The plan it was computed for. - pub fn plan(&self) -> &GroupedRefreshPlan { - &self.plan - } - - /// View rows this unit wrote, its groups. - pub fn rows(&self) -> u64 { - self.rows - } - - /// The fragments it staged, uncommitted until the whole set is published. - pub fn fragments(&self) -> &[Fragment] { - &self.fragments - } -} - -/// The grouped view `plan` was made for, opened at the source version it -/// pins. Every holder of a plan asks the same two questions: is this a -/// grouped view, and is it the view the plan names. -/// -/// The pinned version is also what holds every unit to one index: indices -/// are part of the source manifest, so a rebuild lands in a later version -/// that a unit reading this one never sees, and units written on either -/// side of it still group by the same centroids. -async fn open_planned(view: &Table, plan: &GroupedRefreshPlan) -> Result { - let grouped = open_grouped(view, Some(plan.source_version)) - .await? - .ok_or_else(|| Error::InvalidInput { - message: format!( - "'{}' is not a grouped view keyed by {IVF_PARTITION}", - view.name() - ), - })?; - ensure_incarnation(&grouped.view_ds, Some(&plan.incarnation), view.name()).await?; - Ok(grouped) -} - -/// Compute unit `unit` of `plan` against source version `source_version` -/// and write its groups as uncommitted fragments of the view. -pub async fn write_grouped_unit( - view: &Table, - unit: u32, - plan: &GroupedRefreshPlan, -) -> Result { - let grouped = open_planned(view, plan).await?; - if unit >= plan.units { - return Err(Error::InvalidInput { - message: format!( - "unit {unit} is not in the plan of '{}' ({} units)", - view.name(), - plan.units - ), - }); - } - let schema: SchemaRef = Arc::new(ArrowSchema::from(grouped.view_ds.schema())); - let rows_written = Arc::new(AtomicU64::new(0)); - let rows = unit_rows(&grouped, unit).await?; - let stream = group_rows(&grouped, rows, schema.clone(), rows_written.clone()).await?; - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(grouped.view_ds))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { fragments } = write_txn.operation else { - return Err(Error::Runtime { - message: "expected an append when writing a grouped unit".into(), - }); - }; - Ok(WrittenUnit { - unit, - plan: plan.clone(), - fragments, - rows: rows_written.load(Ordering::Relaxed), - }) -} - -/// The row ids the index places in partition `unit`, at most `limit` of them. -/// A segment's postings outlive its ownership -- a column rewrite attaches a -/// new file and takes the fragment out of its bitmap while leaving the rows -/// listed -- so each segment's rows are kept only while it still holds their -/// fragment. The fragment [`unit_rows`] scans them from instead is the one -/// the rewrite left them in. -async fn indexed_row_ids(grouped: &Grouped, unit: u32, limit: usize) -> Result> { - let addresses = lance::dataset::rowids::get_row_id_index(&grouped.source).await?; - let fragment_of = |row_id: u64| -> Result> { - Ok(match &addresses { - Some(index) => index.get(row_id)?.map(|address| address.fragment_id()), - None => Some(RowAddress::new_from_u64(row_id).fragment_id()), - }) - }; - let mut row_ids = Vec::new(); - for segment in &grouped.segments.segments { - let owned = segment.fragments.as_ref().ok_or_else(|| Error::Runtime { - message: "the index stopped naming the fragments it covers mid-refresh".into(), - })?; - let mut stream = segment - .index - .partition_reader(unit as usize, false, &NoOpMetricsCollector) - .await?; - while let Some(batch) = stream.try_next().await? { - for row_id in batch[ROW_ID].as_primitive::().values() { - if row_ids.len() >= limit { - return Ok(row_ids); - } - if fragment_of(*row_id)?.is_some_and(|fragment| owned.contains(fragment)) { - row_ids.push(*row_id); - } - } - } - } - Ok(row_ids) -} - -/// The first row of `batch` that `unit` holds but `ivf_partition` does not -/// put there, if there is one. `batch` carries the key column. -fn assigned_elsewhere(grouped: &Grouped, batch: &RecordBatch, unit: u32) -> Result> { - let vectors = batch[grouped.column.as_str()].as_fixed_size_list(); - let assigned = grouped.binding.assign(vectors)?; - Ok((0..batch.num_rows()).find(|&row| !(assigned.is_valid(row) && assigned.value(row) == unit))) -} - -/// Whether the index's posting lists are `ivf_partition`'s buckets, over a -/// sample of each partition. Lance also builds an index from precomputed -/// assignments, and records nowhere that it did, so the only way to ask is -/// to recompute some of them. -async fn assigns_by_its_centroids(grouped: &Grouped) -> Result { - let mut expected = HashMap::new(); - for unit in 0..grouped.binding.partitions { - if expected.len() >= PROBE_ROWS { - break; - } - for row_id in indexed_row_ids(grouped, unit, PROBE_PER_UNIT).await? { - expected.insert(row_id, unit); - } - } - if expected.is_empty() { - return Ok(true); - } - let sample: Vec = expected.keys().copied().collect(); - let columns = [ROW_ID, grouped.column.as_str()]; - let projection = ProjectionRequest::from_columns(columns, grouped.source.schema()); - let batch = grouped.source.take_rows(&sample, projection).await?; - let row_ids = batch[ROW_ID].as_primitive::().clone(); - let assigned = grouped - .binding - .assign(batch[grouped.column.as_str()].as_fixed_size_list())?; - Ok((0..batch.num_rows()).all(|row| { - assigned.is_valid(row) && Some(&assigned.value(row)) == expected.get(&row_ids.value(row)) - })) -} - -/// The source rows of one unit: the index's rows for partition `unit` plus -/// the unindexed rows that assign to it, or for the last unit every row the -/// assigner cannot place. Each batch carries the inputs and `_rowid`. -async fn unit_rows(grouped: &Grouped, unit: u32) -> Result> { - let mut columns: Vec<&str> = vec![ROW_ID]; - columns.extend(grouped.inputs.iter().map(String::as_str)); - if !grouped.inputs.contains(&grouped.column) { - columns.push(&grouped.column); - } - let source = Arc::new(grouped.source.clone()); - let mut batches = Vec::new(); - let unassigned = unit == grouped.binding.partitions; - - if !unassigned { - // Rows the index placed in this partition; a deleted row is dropped - // by the take. - let row_ids = indexed_row_ids(grouped, unit, usize::MAX).await?; - let projection = ProjectionRequest::from_columns(&columns, source.schema()); - for chunk in row_ids.chunks(TAKE_CHUNK) { - let batch = source.take_rows(chunk, projection.clone()).await?; - if batch.num_rows() > 0 { - // What the plan sampled, this unit proves: a row grouped - // here that belongs to another key would be aggregated in - // both units and published as two rows of one group. - if let Some(row) = assigned_elsewhere(grouped, &batch, unit)? { - let row_id = batch[ROW_ID].as_primitive::().value(row); - return Err(Error::NotSupported { - message: format!( - "{IVF_PARTITION}({}): the index lists row {row_id} under \ - partition {unit}, which it does not assign to; this view \ - is refreshed in one pass", - grouped.column - ), - }); - } - batches.push(in_column_order(&batch, &columns)?); - } - } - } - - // Rows the index has not seen assign here; the last unit also takes the - // rows no assignment places, from every fragment. - let mut fragments = Vec::new(); - for fragment in source.get_fragments() { - let unindexed = grouped - .segments - .is_unindexed(fragment.id() as u64) - .ok_or_else(|| Error::Runtime { - message: "the index stopped naming the fragments it covers mid-refresh".into(), - })?; - if unassigned || unindexed { - fragments.push(fragment.metadata().clone()); - } - } - if !fragments.is_empty() { - let mut scanner = source.scan(); - scanner.with_fragments(fragments); - scanner.with_row_id(); - scanner.project(&columns)?; - let mut stream = scanner.try_into_stream().await?; - while let Some(batch) = stream.try_next().await? { - let vectors = batch[grouped.column.as_str()].as_fixed_size_list(); - let assigned = grouped.binding.assign(vectors)?; - let keep: arrow_array::BooleanArray = (0..batch.num_rows()) - .map(|row| { - Some(if unassigned { - assigned.is_null(row) - } else { - assigned.is_valid(row) && assigned.value(row) == unit - }) - }) - .collect(); - let kept = arrow_select::filter::filter_record_batch(&batch, &keep)?; - if kept.num_rows() > 0 { - batches.push(in_column_order(&kept, &columns)?); - } - } - } - Ok(batches) -} - -/// `batch` with its columns in `columns` order: the take and the scan place -/// `_rowid` differently, and the query stream coalesces by position. -fn in_column_order(batch: &RecordBatch, columns: &[&str]) -> Result { - let indices = columns - .iter() - .map(|column| batch.schema().index_of(column)) - .collect::, _>>()?; - Ok(batch.project(&indices)?) -} - -/// The grouped query over `rows`, in the view's schema. -async fn group_rows( - grouped: &Grouped, - rows: Vec, - schema: SchemaRef, - rows_written: Arc, -) -> Result { - let input_schema = match rows.first() { - Some(batch) => batch.schema(), - None => { - return Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::empty(), - ))); - } - }; - let rows: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - input_schema, - futures::stream::iter(rows.into_iter().map(Ok)), - )); - let state = grouped::session(grouped::bind(&grouped.source, &grouped.definition).await?); - let ctx = SessionContext::new_with_state(state.clone()); - let source = ctx.read_one_shot(rows)?.into_view(); - let planned = grouped::logical_plan( - &state, - provider_as_source(source), - &grouped.definition, - // Neither the indexed take nor the fragment scan filters: the rows - // reach this query raw, so the predicate belongs in it. - true, - )?; - let groups = ctx - .execute_logical_plan(planned) - .await? - .execute_stream() - .await?; - let out_schema = schema.clone(); - let mapped = groups.map(move |batch| { - let batch = to_view_batch(&batch?, &out_schema)?; - rows_written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); - Ok(batch) - }); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped))) -} - -/// Replace the view's rows with the fragments every unit wrote, in one -/// commit, and record `plan.source_version` as what the view reflects. -/// -/// The commit is refused if the view moved since the plan: a refresh that -/// landed meanwhile replaced rows the units never saw. Concurrent refreshes -/// of one view are serialized by the caller, as for a single-pass rebuild; -/// this check catches the ones that slipped through. -pub async fn commit_grouped_refresh( - view: &Table, - plan: &GroupedRefreshPlan, - units: Vec, - expected_incarnation: Option<&str>, -) -> Result { - let grouped = open_planned(view, plan).await?; - let units = complete_set(view.name(), plan, units)?; - let rows_written = units.iter().map(|unit| unit.rows).sum(); - let new_fragments = units.into_iter().flat_map(|unit| unit.fragments).collect(); - commit_fragments( - view, - plan, - &grouped.definition, - grouped.source.manifest.timestamp_nanos, - new_fragments, - rows_written, - expected_incarnation, - ) - .await -} - -/// The common all-or-nothing unit sink, shared by grouped and native views. -#[allow(clippy::too_many_arguments)] -pub(super) async fn commit_fragments( - view: &Table, - plan: &GroupedRefreshPlan, - definition: &MaterializedViewDefinition, - source_ts: u128, - new_fragments: Vec, - rows_written: u64, - expected_incarnation: Option<&str>, -) -> Result { - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { - message: "materialized views are supported only on local tables".into(), - })?; - view_native.dataset.ensure_mutable()?; - let lock = refresh_lock(view_native.dataset.get().await?.uri()); - let _guard = lock.lock().await; - if let Some(expected) = expected_incarnation - && expected != plan.incarnation - { - return Err(Error::Runtime { - message: format!( - "the refresh of '{}' was requested for another incarnation than the one \ - its units were planned against", - view.name() - ), - }); - } - view_native.dataset.reload().await?; - let ds = view_native.dataset.get().await?.as_ref().clone(); - ensure_incarnation(&ds, Some(&plan.incarnation), view.name()).await?; - if super::read_definition(&ds.schema().metadata)? - != Some(super::StoredDefinition::Query(definition.clone())) - { - return Err(Error::InvalidInput { - message: "view definition changed while its units were written".into(), - }); - } - let ds = Arc::new(ds); - if ds.version().version != plan.view_version { - return Err(Error::Runtime { - message: format!( - "'{}' moved from version {} to {} while its units were written; the \ - refresh is unrecorded and must be replanned", - view.name(), - plan.view_version, - ds.version().version - ), - }); - } - let removed_fragment_ids: Vec = ds.get_fragments().iter().map(|f| f.id() as u64).collect(); - // Fragment ids are assigned at commit; the units' fragments all carry 0. - let transaction = Transaction::new( - plan.view_version, - Operation::Update { - removed_fragment_ids, - updated_fragments: Vec::new(), - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: None, - updated_fragment_offsets: None, - }, - None, - ); - #[cfg(test)] - super::refresh::tests::hold_before_publish(ds.uri()).await; - let committed = CommitBuilder::new(WriteDestination::Dataset(ds)) - .execute(transaction) - .await?; - // The version check above is not the fence: lance rebases this `Update` - // over a concurrent append rather than refusing it, which would leave - // the appended rows in a view certified as the units' whole output. - if committed.version().version != plan.view_version + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced the refresh of '{}' (view version {}); the \ - refresh is unrecorded and must be replanned", - view.name(), - committed.version().version - ), - }); - } - let version = stamp_watermark( - view_native, - committed, - plan.source_version, - source_ts, - Some(definition), - expected_incarnation, - ) - .await?; - Ok(RefreshMaterializedViewResult { - mode: RefreshMode::Rebuild, - rows_written, - source_version: plan.source_version, - version, - }) -} - -/// The units in plan order, exactly one per unit of `plan`. The commit -/// publishes them as the view's whole contents, so a missing, repeated or -/// foreign result would certify a view nobody computed. -fn complete_set( - view: &str, - plan: &GroupedRefreshPlan, - units: Vec, -) -> Result> { - let mut ordered: Vec> = vec![None; plan.units as usize]; - for written in units { - if &written.plan != plan { - return Err(Error::InvalidInput { - message: format!( - "a unit of another plan of '{view}' was offered to this refresh; the \ - refresh is unrecorded and must be replanned" - ), - }); - } - let slot = ordered - .get_mut(written.unit as usize) - .ok_or_else(|| Error::InvalidInput { - message: format!( - "unit {} is not in the plan of '{view}' ({} units)", - written.unit, plan.units - ), - })?; - if slot.is_some() { - return Err(Error::InvalidInput { - message: format!("unit {} of '{view}' was written twice", written.unit), - }); - } - // A unit holds one key, so it writes that key's one group. - if written.rows > 1 { - return Err(Error::InvalidInput { - message: format!( - "unit {} of '{view}' wrote {} groups; the split did not match the \ - grouping and the refresh must be replanned", - written.unit, written.rows - ), - }); - } - *slot = Some(written); - } - ordered - .into_iter() - .enumerate() - .map(|(unit, written)| { - written.ok_or_else(|| Error::InvalidInput { - message: format!( - "unit {unit} of '{view}' is missing; the refresh holds {} units and \ - publishes all of them or none", - plan.units - ), - }) - }) - .collect() -} - -#[cfg(test)] -mod tests { - use super::super::refresh::tests::{clustered_source, declare, rows}; - use super::*; - use crate::connect; - use crate::index::Index; - use crate::index::vector::IvfFlatIndexBuilder; - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array, RecordBatchIterator}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - use lance_index::vector::ivf::IvfBuildParams; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, min(id) AS lo, \ - max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - - /// The clustered source under a two-partition cosine index, then edited - /// behind the index: a row deleted from each cluster, a row appended to - /// each, and one zero vector cosine cannot place. - async fn edited_source(conn: &crate::Connection) -> Table { - let source = clustered_source(conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - source.delete("id IN (3, 13)").await.unwrap(); - let appended = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(1.5), Some(0.0)]), - Some(vec![Some(101.5), Some(100.0)]), - Some(vec![Some(0.0), Some(0.0)]), - ], - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from(vec![20, 21, 22])) as ArrayRef, - ), - ("vec", Arc::new(appended) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - source - } - - async fn refresh_in_units( - view: &Table, - expected_incarnation: Option<&str>, - ) -> Result<( - GroupedRefreshPlan, - Vec, - RefreshMaterializedViewResult, - )> { - let plan = plan_grouped_refresh(view, None) - .await? - .expect("an ivf-keyed view"); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view, unit, &plan).await?); - } - let result = - commit_grouped_refresh(view, &plan, units.clone(), expected_incarnation).await?; - Ok((plan, units, result)) - } - - /// Every partition's unit reads its indexed rows plus the appended rows - /// that assign to it, the deleted rows are gone, and the last unit holds - /// the unplaceable row: together they are the single-pass result. - #[tokio::test] - async fn units_compose_to_the_single_pass_result() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let single = declare(&source, "single", BUCKETS).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - assert_eq!(expected.len(), 3, "{expected:?}"); - assert!(expected.iter().any(|row| row == " 1 22 22"), "{expected:?}"); - - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let (plan, units, result) = refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!(plan.units, 3, "two partitions plus the unplaceable rows"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - assert_eq!(units.iter().map(|u| u.rows).sum::(), 3); - assert_eq!(units.iter().map(|u| u.rows).collect::>(), [1, 1, 1]); - assert_eq!(result.rows_written, 3); - assert_eq!(result.source_version, plan.source_version); - assert_eq!(view.table().version().await.unwrap(), result.version); - - // The commit is a refresh: the view reflects the source version it - // read and a second pass finds nothing to do. - let again = view.refresh().execute().await.unwrap(); - assert_eq!(again.rows_written, 0, "{again:?}"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - } - - /// A second refresh in units replaces the first one's rows rather than - /// appending to them. - #[tokio::test] - async fn a_refresh_in_units_replaces_the_previous_rows() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - source.delete("id = 22").await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - let mut groups = rows(view.table(), &["n", "lo", "hi"]).await; - groups.sort(); - assert_eq!(groups, ["10 0 20", "10 10 21"]); - } - - #[tokio::test] - async fn a_view_not_keyed_by_ivf_partition_has_no_units() { - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - let view = declare( - &source, - "by_id", - "SELECT id % 2 AS k, count(*) AS n FROM images GROUP BY id % 2", - ) - .await - .unwrap(); - assert_eq!( - plan_grouped_refresh(view.table(), None).await.unwrap(), - None - ); - } - - #[tokio::test] - async fn a_unit_outside_the_plan_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let err = write_grouped_unit(view.table(), plan.units, &plan) - .await - .unwrap_err(); - assert!(err.to_string().contains("not in the plan"), "{err}"); - } - - /// A commit that lands on a view someone else refreshed meanwhile is - /// refused rather than recorded over their rows. - #[tokio::test] - async fn a_raced_commit_is_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let units = vec![write_grouped_unit(view.table(), 0, &plan).await.unwrap()]; - view.refresh().execute().await.unwrap(); - let before = rows(view.table(), &["bucket", "n", "lo", "hi"]).await; - let raced = commit_grouped_refresh(view.table(), &plan, units, None).await; - assert!(raced.is_err(), "{raced:?}"); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - before - ); - } - - /// A view with a predicate holds the rows it selects however it is - /// refreshed: the units read their rows through the index and by - /// scanning fragments, neither of which filters, so the predicate has - /// to be in the grouped query itself. - #[tokio::test] - async fn a_filtered_view_refreshes_in_units_like_a_single_pass() { - const FILTERED: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images WHERE id < 5 \ - GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let single = declare(&source, "single", FILTERED).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - - let view = declare(&source, "buckets", FILTERED).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected - ); - } - - /// The commit publishes its units as the view's whole contents, so it - /// takes the plan's units exactly once each and nothing else. - #[tokio::test] - async fn only_a_complete_set_of_this_plan_s_units_is_published() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - - let missing = commit_grouped_refresh(view.table(), &plan, units[..2].to_vec(), None).await; - assert!( - missing.unwrap_err().to_string().contains("is missing"), - "a partial set was published" - ); - let mut twice = units.clone(); - twice[1] = units[0].clone(); - let repeated = commit_grouped_refresh(view.table(), &plan, twice, None).await; - assert!( - repeated.unwrap_err().to_string().contains("written twice"), - "a repeated unit was published" - ); - let mut foreign = units.clone(); - foreign[0].plan.source_version += 1; - let stale = commit_grouped_refresh(view.table(), &plan, foreign, None).await; - assert!( - stale.unwrap_err().to_string().contains("another plan"), - "a unit of another plan was published" - ); - assert_eq!(rows(view.table(), &["bucket"]).await, Vec::::new()); - - // The complete set still lands. - commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap(); - assert_eq!(rows(view.table(), &["bucket"]).await.len(), 3); - } - - /// A column rewrite attaches a new file to the row's fragment and takes - /// that fragment out of the index segment's bitmap, but the segment's - /// postings still name its rows. Taking them anyway, while the fragment - /// is also scanned as unindexed, counts those rows twice and buckets the - /// stale copy by the value the rewrite replaced. - #[tokio::test] - async fn a_rewritten_fragment_is_read_once() { - use lance::dataset::{ - MergeInsertBuilder, MergeInsertWriteMode, WhenMatched, WhenNotMatched, - }; - - let conn = connect("memory://").execute().await.unwrap(); - // A column the rewrite leaves alone: RewriteColumns is only for a - // source that covers some of the row, not all of it. - let vectors = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| { - let base = if i < 10 { 0.0 } else { 100.0 }; - Some(vec![Some(base + 1.0 + i as f32 * 0.1), Some(base)]) - }), - 2, - ); - let source = conn - .create_table( - "images", - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as ArrayRef, - ), - ("vec", Arc::new(vectors) as ArrayRef), - ( - "tag", - Arc::new(Int32Array::from_iter_values((0..20).map(|i| i * 10))) as ArrayRef, - ), - ]) - .unwrap(), - ) - .write_options(super::super::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - source - .create_index( - &["vec"], - Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(2)), - ) - .execute() - .await - .unwrap(); - - // Row 0 keeps its id and its fragment; its vector is replaced in a - // file the index segment does not cover. - let batch = RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ( - "vec", - Arc::new( - FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(101.0), Some(100.0)])], - 2, - ), - ) as ArrayRef, - ), - ]) - .unwrap(); - let schema = batch.schema(); - let dataset = source.as_native().unwrap().dataset.get().await.unwrap(); - let mut builder = - MergeInsertBuilder::try_new(Arc::new(dataset.as_ref().clone()), vec!["id".to_string()]) - .unwrap(); - builder - .when_matched(WhenMatched::UpdateAll) - .when_not_matched(WhenNotMatched::DoNothing) - .write_mode(MergeInsertWriteMode::RewriteColumns); - builder - .try_build() - .unwrap() - .execute_reader(arrow_array::RecordBatchIterator::new([Ok(batch)], schema)) - .await - .unwrap(); - - let single = declare(&source, "single", BUCKETS).await.unwrap(); - single.refresh().execute().await.unwrap(); - let expected = rows(single.table(), &["bucket", "n", "lo", "hi"]).await; - - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - refresh_in_units(view.table(), None).await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - expected, - "the rewritten row was counted twice or bucketed by its old vector" - ); - } - - /// Units written on either side of an index rebuild still group by one - /// model: the plan pins a source version, indices live in the source - /// manifest, and a rebuild lands in a later version the units never - /// read. Without that, two units of one plan could assign their rows by - /// different centroids. - #[tokio::test] - async fn a_rebuilt_index_does_not_reach_the_units_of_a_plan() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let first = write_grouped_unit(view.table(), 0, &plan).await.unwrap(); - - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .replace(true) - .execute() - .await - .unwrap(); - let rebuilt = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - assert_eq!(rebuilt.units, plan.units, "same shape"); - assert!( - rebuilt.source_version > plan.source_version, - "the rebuild is a later source version" - ); - - // The remaining units read the version the plan pinned, so they are - // grouped by the index that version holds, not the new one. - let mut units = vec![first]; - for unit in 1..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await.len(), - 3 - ); - } - - /// Two views of the same shape over the same source reach the same plan - /// counters, so the results of one must not satisfy the other: they are - /// fragments written into a different dataset. - #[tokio::test] - async fn units_written_for_another_view_are_refused() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let mine = declare(&source, "buckets", BUCKETS).await.unwrap(); - let theirs = declare(&source, "twin", BUCKETS).await.unwrap(); - let my_plan = plan_grouped_refresh(mine.table(), None) - .await - .unwrap() - .unwrap(); - let their_plan = plan_grouped_refresh(theirs.table(), None) - .await - .unwrap() - .unwrap(); - assert_eq!( - (my_plan.units, my_plan.view_version, my_plan.source_version), - ( - their_plan.units, - their_plan.view_version, - their_plan.source_version - ), - "the counters alone do not tell these views apart" - ); - - let mut mine_units = Vec::new(); - for unit in 0..my_plan.units { - mine_units.push( - write_grouped_unit(mine.table(), unit, &my_plan) - .await - .unwrap(), - ); - } - let misrouted = - commit_grouped_refresh(theirs.table(), &their_plan, mine_units.clone(), None).await; - assert!( - misrouted.unwrap_err().to_string().contains("another plan"), - "another view's units were published" - ); - assert_eq!( - rows(theirs.table(), &["bucket"]).await, - Vec::::new() - ); - - // And a unit cannot be written for a view the plan does not name. - let foreign = write_grouped_unit(theirs.table(), 0, &my_plan).await; - assert!( - foreign - .unwrap_err() - .to_string() - .contains("not the incarnation"), - "a unit was written under another view's plan" - ); - } - - /// A write that lands on the view between the plan check and the commit - /// is not a conflict lance refuses -- the update rebases over it -- so - /// the publication checks the generation it actually landed on. - #[tokio::test(flavor = "multi_thread")] - async fn a_write_racing_the_publication_is_refused() { - use super::super::refresh::tests::{ - DRIFT_LOCK, DRIFT_PLANNED, DRIFT_RELEASED, DRIFT_TARGET, - }; - let _serial = DRIFT_LOCK.lock().await; - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - // Its own name: memory:// uris repeat across tests, so a view called - // like its siblings would arm the hook for one of their commits. - let view = declare(&source, "racing_buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - - let publishing = { - let table = view.table().clone(); - let plan = plan.clone(); - tokio::spawn(async move { commit_grouped_refresh(&table, &plan, units, None).await }) - }; - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the commit never reached the publication boundary"); - view.table() - .add( - RecordBatch::try_from_iter(vec![ - ("bucket", Arc::new(Int32Array::from(vec![7])) as ArrayRef), - ("n", Arc::new(Int32Array::from(vec![1])) as ArrayRef), - ("lo", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ("hi", Arc::new(Int32Array::from(vec![0])) as ArrayRef), - ( - super::super::SOURCE_ROW_ID_COLUMN, - Arc::new(arrow_array::UInt64Array::from(vec![0u64])) as ArrayRef, - ), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = publishing.await.unwrap().unwrap_err(); - assert!(err.to_string().contains("raced the refresh"), "{err}"); - } - - #[test] - fn the_key_column_is_read_off_a_single_ivf_key() { - let mut definition = MaterializedViewDefinition::from_sql( - "SELECT ivf_partition(vec) AS b, count(*) AS n FROM t GROUP BY ivf_partition(vec)", - ) - .unwrap(); - assert_eq!(ivf_key_column(&definition).as_deref(), Some("vec")); - definition.group_by = vec!["k".to_string()]; - assert_eq!(ivf_key_column(&definition), None); - definition.group_by = vec!["ivf_partition(vec)".to_string(), "k".to_string()]; - assert_eq!(ivf_key_column(&definition), None); - } - - /// A source whose index does not assign by its own centroids: lance - /// takes precomputed partitions, so row 1 is stored under partition 1 - /// and `ivf_partition` puts it in 0. - async fn precomputed_source(conn: &crate::Connection, dir: &std::path::Path) -> Table { - let vectors = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(1.0), Some(0.0)]), - Some(vec![Some(2.0), Some(0.0)]), - Some(vec![Some(100.0), Some(100.0)]), - ], - 2, - ); - let source = conn - .create_table( - "images", - RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![0, 1, 2])) as ArrayRef), - ("vec", Arc::new(vectors) as ArrayRef), - ]) - .unwrap(), - ) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - - let assignments = arrow_array::record_batch!( - ("row_id", UInt64, [0, 1, 2]), - ("partition", UInt32, [0, 1, 1]) - ) - .unwrap(); - let file = dir.join("partitions.lance"); - let file = file.to_str().unwrap(); - lance::Dataset::write( - RecordBatchIterator::new(vec![Ok(assignments.clone())], assignments.schema()), - file, - None, - ) - .await - .unwrap(); - - let centroids = FixedSizeListArray::from_iter_primitive::( - [ - Some(vec![Some(0.0), Some(0.0)]), - Some(vec![Some(100.0), Some(100.0)]), - ], - 2, - ); - let mut ivf = IvfBuildParams::try_with_centroids(2, Arc::new(centroids)).unwrap(); - ivf.precomputed_partitions_file = Some(file.to_string()); - let mut dataset = source - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .as_ref() - .clone(); - dataset - .create_index( - &["vec"], - IndexType::Vector, - None, - &VectorIndexParams::with_ivf_flat_params( - lance_linalg::distance::DistanceType::L2, - ivf, - ), - true, - ) - .await - .unwrap(); - conn.open_table("images").execute().await.unwrap() - } - - async fn incarnation(view: &Table) -> String { - view.as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .schema() - .metadata - .get(crate::materialized_view::INCARNATION_META_KEY) - .cloned() - .unwrap() - } - - /// The split is the index's posting lists and the key is - /// `ivf_partition`. Where they disagree one group is aggregated in two - /// units and published twice, so such a view has no units at all. - #[tokio::test] - async fn an_index_that_does_not_assign_by_its_centroids_has_no_units() { - let dir = tempfile::tempdir().unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = precomputed_source(&conn, dir.path()).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - assert_eq!( - plan_grouped_refresh(view.table(), None).await.unwrap(), - None - ); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - ["0 2 0 1", "1 1 2 2"] - ); - } - - /// The publication is the last place the split can be caught: a unit - /// that holds one key writes one group, and a set that says otherwise is - /// not published however it was produced. - #[tokio::test] - async fn a_unit_that_wrote_more_than_one_group_is_not_published() { - let conn = connect("memory://").execute().await.unwrap(); - let source = edited_source(&conn).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = plan_grouped_refresh(view.table(), None) - .await - .unwrap() - .unwrap(); - let mut units = Vec::new(); - for unit in 0..plan.units { - units.push(write_grouped_unit(view.table(), unit, &plan).await.unwrap()); - } - units[1].rows = 2; - let err = commit_grouped_refresh(view.table(), &plan, units, None) - .await - .unwrap_err(); - assert!(err.to_string().contains("wrote 2 groups"), "{err}"); - assert_eq!(rows(view.table(), &["bucket"]).await, Vec::::new()); - } - - /// The plan reads a sample, so the unit proves the rest: it holds only - /// rows that assign to it, or it refuses rather than write a group - /// another unit holds too. - #[tokio::test] - async fn a_unit_holding_a_row_that_assigns_elsewhere_is_refused() { - let dir = tempfile::tempdir().unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = precomputed_source(&conn, dir.path()).await; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - let plan = GroupedRefreshPlan { - units: 3, - source_version: source.version().await.unwrap(), - view_version: view.table().version().await.unwrap(), - incarnation: incarnation(view.table()).await, - }; - assert_eq!( - write_grouped_unit(view.table(), 0, &plan) - .await - .unwrap() - .rows, - 1 - ); - let err = write_grouped_unit(view.table(), 1, &plan) - .await - .unwrap_err(); - assert!(err.to_string().contains("does not assign to"), "{err}"); - } -} diff --git a/rust/lancedb/src/materialized_view/partitioned.rs b/rust/lancedb/src/materialized_view/partitioned.rs deleted file mode 100644 index 58192089b..000000000 --- a/rust/lancedb/src/materialized_view/partitioned.rs +++ /dev/null @@ -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, -) -> Result> { - 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 { - 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 { - 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::>>()?; - 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, - expected: Option<&str>, -) -> Result { - 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::>>()?; - 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::>>()?; - 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::>>()?; - super::commit_grouped_refresh(view, plan, units, expected).await - } - } -} diff --git a/rust/lancedb/src/materialized_view/query.rs b/rust/lancedb/src/materialized_view/query.rs index 4bfd5542e..c1a987b2c 100644 --- a/rust/lancedb/src/materialized_view/query.rs +++ b/rust/lancedb/src/materialized_view/query.rs @@ -264,36 +264,7 @@ fn extract(query: &Query) -> Result { }; 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 { }; let definition = MaterializedViewDefinition { - vector_source, source_table, source_namespace, lateral, @@ -454,9 +424,6 @@ fn extract(query: &Query) -> Result { 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 [ diff --git a/rust/lancedb/src/materialized_view/refresh.rs b/rust/lancedb/src/materialized_view/refresh.rs index 9f0ea2646..f9634837a 100644 --- a/rust/lancedb/src/materialized_view/refresh.rs +++ b/rust/lancedb/src/materialized_view/refresh.rs @@ -1,63 +1,37 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Refreshing materialized views. +//! Full refresh for materialized views. //! -//! A refresh pins one source version and brings the view to exactly the -//! definition's result at that version: added rows are computed and appended, -//! removed or changed rows are evicted by provenance id and recomputed in the -//! same pass. Compaction outputs cost nothing, which is only sound while -//! [`SOURCE_ROW_ID_COLUMN`] stays valid across the rewrite. Anything the -//! classifier cannot prove intact rebuilds; an indexed rebuild swaps all -//! fragments in one commit that retains index definitions. -//! -//! The watermark ([`SOURCE_VERSION_META_KEY`]) lands in a follow-up commit; a -//! crash or race between the two leaves the view visibly unstamped and the -//! next refresh rebuilds. In-process refreshes serialize on a per-view lock; -//! across processes the commit's inserted-rows filter carries a shared token, -//! so two refreshes of one view conflict and only one lands. +//! A refresh evaluates the definition at one source version, stages the +//! complete result, and atomically replaces every fragment of the view. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex as StdMutex, OnceLock}; -use std::time::{SystemTime, UNIX_EPOCH}; use arrow_array::cast::AsArray; -use arrow_array::types::UInt64Type; -use arrow_array::{RecordBatch, UInt64Array, new_null_array}; -use arrow_schema::{DataType, Field as ArrowField, FieldRef, Schema as ArrowSchema, SchemaRef}; -use datafusion::common::ScalarValue; +use arrow_array::{RecordBatch, new_null_array}; +use arrow_schema::{Schema as ArrowSchema, SchemaRef}; use datafusion::error::DataFusionError; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::SendableRecordBatchStream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::prelude::{col, lit}; -use futures::{StreamExt, TryStreamExt}; +use futures::StreamExt; use lance::Dataset; use lance::dataset::mem_wal::DatasetMemWalExt; -use lance::dataset::transaction::{Operation, Transaction, UpdateMode}; -use lance::dataset::write::delete::DeleteBuilder; +use lance::dataset::transaction::{Operation, Transaction}; use lance::dataset::write::merge_insert::inserted_rows::{ KeyExistenceFilter, KeyExistenceFilterBuilder, KeyValue, }; use lance::dataset::{CommitBuilder, InsertBuilder, WriteDestination, WriteMode, WriteParams}; -use lance_core::{ROW_CREATED_AT_VERSION, ROW_ID, ROW_LAST_UPDATED_AT_VERSION}; - use lance_datafusion::planner::Planner; -use lance_file::version::ConcreteFileVersion; -use lance_table::format::Fragment; use serde::{Deserialize, Serialize}; -use super::{ - DEFINITION_META_KEY, INCARNATION_META_KEY, MaterializedViewDefinition, - REFRESHED_AT_MS_META_KEY, SOURCE_ROW_ID_COLUMN, SOURCE_VERSION_META_KEY, - definition_to_metadata, -}; +use super::{MaterializedViewDefinition, StoredDefinition}; use crate::database::OpenTableRequest; -use crate::table::computed_columns::{ - computed_column_from_field, computed_columns, ensure_declarations_are_planned, -}; -use crate::table::{NativeTable, NativeTableExt, Table}; +use crate::table::computed_columns::{computed_column_from_field, ensure_declarations_are_planned}; +use crate::table::{NativeTableExt, Table}; use crate::{Error, Result}; /// How a refresh brought the view up to date. @@ -66,10 +40,6 @@ use crate::{Error, Result}; pub enum RefreshMode { /// The view was recomputed from scratch. Rebuild, - /// Rows from source fragments added since the last refresh were appended. - Incremental, - /// The view was already at the requested source version. - NoOp, } /// The result of refreshing a materialized view. @@ -77,28 +47,15 @@ pub enum RefreshMode { pub struct RefreshMaterializedViewResult { /// How the view was brought up to date. pub mode: RefreshMode, - /// Rows written to the view: everything on a rebuild, and on an - /// incremental refresh both the rows added and the rows recomputed in - /// place of ones the source changed. + /// Rows written to the replacement view. pub rows_written: u64, - /// The source table version the view now reflects. + /// The source table version the view reflects. pub source_version: u64, /// The view table version after the refresh. pub version: u64, } -/// Schema metadata key holding the view table version a successful refresh -/// left behind. Any other commit on the view is drift, and refresh rebuilds. -pub const VIEW_VERSION_META_KEY: &str = "mv.view_version"; - -/// Schema metadata key holding the commit timestamp of the watermark's source -/// manifest. A dropped and recreated source reuses version numbers but never -/// their timestamps, so a mismatch means the watermark describes a different -/// incarnation and refresh rebuilds. -pub const SOURCE_VERSION_TS_META_KEY: &str = "mv.source_version_ts"; - -/// One refresh per view at a time within this process. -pub(super) fn refresh_lock(uri: &str) -> Arc> { +fn refresh_lock(uri: &str) -> Arc> { static LOCKS: OnceLock>>>> = OnceLock::new(); LOCKS @@ -110,57 +67,27 @@ pub(super) fn refresh_lock(uri: &str) -> Arc> { .clone() } -/// Internal implementation of the refresh logic. +/// Recompute a local materialized view from one source version. pub(crate) async fn execute_refresh( view: &Table, - full: bool, pinned: Option, - expected_incarnation: Option<&str>, ) -> Result { - // Native table functions use the same partition writer and commit sink in - // local execution; Sophon drives these units through its job registry. - if let Some(super::StoredDefinition::Query(definition)) = - super::read_definition(view.schema().await?.metadata())? - && definition.vector_source.is_some() - { - let plan = super::plan_partitioned_refresh(view, pinned) - .await? - .ok_or_else(|| Error::Runtime { - message: "native pair view has no partition plan".into(), - })?; - let mut units = Vec::with_capacity(plan.units() as usize); - for unit in 0..plan.units() { - let dependencies = &units[..units.len().min(plan.dependency_units() as usize)]; - units.push( - super::write_refresh_partition_with_inputs(view, unit, &plan, dependencies).await?, - ); - } - return super::commit_partitioned_refresh(view, &plan, units, expected_incarnation).await; - } - let view_native = view.as_native().ok_or_else(|| Error::NotSupported { + let native = view.as_native().ok_or_else(|| Error::NotSupported { message: "materialized views are supported only on local tables".into(), })?; - view_native.dataset.ensure_mutable()?; - let lock = refresh_lock(view_native.dataset.get().await?.uri()); + native.dataset.ensure_mutable()?; + let lock = refresh_lock(native.dataset.get().await?.uri()); let _guard = lock.lock().await; - // Force-load the latest view state under the lock: each handle caches - // lazily, and a second handle would otherwise plan from a snapshot taken - // before another handle's commit -- appending the same rows again or - // reporting NoOp over a mutated view. - view_native.dataset.reload().await?; - let view_ds = view_native.dataset.get().await?.as_ref().clone(); + native.dataset.reload().await?; + let view_ds = native.dataset.get().await?.as_ref().clone(); + ensure_no_mem_wal(&view_ds, "materialized view", view.name()).await?; - ensure_incarnation(&view_ds, expected_incarnation, view.name()).await?; - - // The definition a handle cached at open may since have been replaced; - // what refresh executes and what it stamps must be one generation. let definition = match super::read_definition(&view_ds.schema().metadata)? { - Some(super::StoredDefinition::Query(definition)) => definition, - Some(super::StoredDefinition::Newer { format }) => { + Some(StoredDefinition::Query(definition)) => definition, + Some(StoredDefinition::Newer { format }) => { return Err(Error::NotSupported { message: format!( - "materialized view '{}' is stored in format {format}, which this \ - version of lancedb cannot refresh", + "materialized view '{}' is stored in format {format}, which this version of lancedb cannot refresh", view.name() ), }); @@ -171,572 +98,216 @@ pub(crate) async fn execute_refresh( }); } }; - let definition = &definition; - let staging = super::read_staging(&view_ds.schema().metadata)?; - ensure_no_mem_wal(&view_ds, "materialized view", view.name()).await?; - - let source_ds = open_source(view, definition, staging.as_ref()).await?; - let source_ds = match pinned { - Some(version) => source_ds.checkout_version(version).await?, - None => source_ds, + let source = open_source(view, &definition).await?; + let source = match pinned { + Some(version) => source.checkout_version(version).await?, + None => source, }; - ensure_no_mem_wal(&source_ds, "source table", &definition.source_table).await?; - let source_version = source_ds.version().version; - let source_ts = source_ds.manifest.timestamp_nanos; + ensure_no_mem_wal(&source, "source table", &definition.source_table).await?; + let source_version = source.version().version; - // Re-plan the persisted definition against the current source schema and - // require its planned output to be exactly the view's physical schema: a - // definition the stored table cannot represent must not be certified. - let source_schema = Arc::new(ArrowSchema::from(source_ds.schema())); + let source_schema = Arc::new(ArrowSchema::from(source.schema())); let super::Planned { - definition: replanned, + definition, fields: planned_fields, inputs, .. - } = super::plan(source_schema.clone(), definition, staging.as_ref()).map_err(|e| match e { - // The stored query planned when the view was declared; what changed - // since is the source. + } = super::plan(source_schema, &definition).map_err(|error| match error { Error::InvalidExpression { column, message } => Error::Schema { message: format!( - "view column '{column}' no longer plans against '{}' (a source column \ - was dropped or renamed): {message}", + "view column '{column}' no longer plans against '{}' (a source column was dropped or renamed): {message}", definition.source_table ), }, Error::InvalidInput { message } => Error::Schema { message: format!( - "the stored query no longer plans against '{}' (a source column was \ - dropped or renamed): {message}", + "the stored query no longer plans against '{}' (a source column was dropped or renamed): {message}", definition.source_table ), }, - e => e, + error => error, })?; - let mut planned_fields = planned_fields; - planned_fields.push(arrow_schema::Field::new( - SOURCE_ROW_ID_COLUMN, - arrow_schema::DataType::UInt64, - false, - )); - // A computed column is not planned from the source: refresh writes it - // NULL and its declaration's owner fills it. Its declaration must still - // be complete, and it must be able to hold NULL. - let physical = ArrowSchema::from(view_ds.schema()); - let mut computed = computed_columns(&physical).into_iter().map(|c| c.name); - if let Some(name) = computed.by_ref().find(|name| { - physical - .field_with_name(name) - .is_ok_and(|f| !f.is_nullable()) - }) { + + let physical = Arc::new(ArrowSchema::from(view_ds.schema())); + ensure_declarations_are_planned(&physical)?; + if let Some(field) = physical + .fields() + .iter() + .find(|field| computed_column_from_field(field).is_some() && !field.is_nullable()) + { return Err(Error::Schema { message: format!( - "computed column '{name}' of view '{}' cannot hold NULL; recreate the view", + "computed column '{}' of view '{}' cannot hold NULL; recreate the view", + field.name(), view.name() ), }); } - ensure_declarations_are_planned(&physical)?; - let physical_planned: Vec<&FieldRef> = physical + let stored_fields: Vec<_> = physical .fields() .iter() - .filter(|f| computed_column_from_field(f).is_none()) + .filter(|field| computed_column_from_field(field).is_none()) .collect(); - // A projected column that became nullable at the source still fits the - // view's nullable field; the reverse would not. - let matches = planned_fields.len() == physical_planned.len() - && planned_fields.iter().zip(&physical_planned).all(|(e, p)| { - e.name() == p.name() - && e.data_type() == p.data_type() - && (p.is_nullable() || !e.is_nullable()) - }); + let matches = planned_fields.len() == stored_fields.len() + && planned_fields + .iter() + .zip(stored_fields) + .all(|(planned, stored)| { + planned.name() == stored.name() + && planned.data_type() == stored.data_type() + && (stored.is_nullable() || !planned.is_nullable()) + }); if !matches { return Err(Error::Schema { message: format!( - "the stored definition of view '{}' does not produce this \ - view's schema; recreate the view", + "the stored definition of view '{}' does not produce this view's schema; recreate the view", view.name() ), }); } - // The stored query is rewritten whenever its stored form differs from - // the current one: a legacy layout, or a spelling the canonicalizer no - // longer produces. Whether the rows change is a separate question: a - // legacy raw filter like `"Party" = 'D'` read the double quotes as a - // string literal, so its watermark certifies different rows than the - // canonical predicate, and only a rebuild can replace them. - let current = definition_to_metadata(&replanned)?; - let persist = view_ds.schema().metadata.get(DEFINITION_META_KEY) != Some(¤t); - let unnest = super::physical_unnest(&replanned, staging.as_ref())?; - let definition_changed = !same_meaning(&source_schema, definition, &replanned, unnest.as_ref()); - let definition = &replanned; - let persist = persist.then_some(definition); - if definition_changed { - return rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - true, - expected_incarnation, - ) - .await; - } - - let metadata = &view_ds.schema().metadata; - let watermark: Option = metadata - .get(SOURCE_VERSION_META_KEY) - .and_then(|raw| raw.parse().ok()); - let recorded_ts: Option = metadata - .get(SOURCE_VERSION_TS_META_KEY) - .and_then(|raw| raw.parse().ok()); - // The watermark speaks only for the view state its refresh left behind; - // any other commit on the view since then is drift, except a fill of its - // computed columns, which rewrites nothing refresh certifies. - let recorded_view_version = metadata - .get(VIEW_VERSION_META_KEY) - .and_then(|raw| raw.parse::().ok()); - let view_intact = match recorded_view_version { - Some(recorded) if recorded == view_ds.version().version => true, - Some(recorded) if recorded < view_ds.version().version => { - only_computed_rewrites_since(&view_ds, recorded).await? - } - _ => false, - }; - - if !full && watermark == Some(source_version) && view_intact && recorded_ts == Some(source_ts) { - return Ok(RefreshMaterializedViewResult { - mode: RefreshMode::NoOp, - rows_written: 0, - source_version, - version: view_ds.version().version, - }); - } - - // A group spans fragments: there is no per-fragment increment. - if definition.is_grouped() { - return rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - None, - persist.is_some(), - expected_incarnation, - ) - .await; - } - let watermark = watermark.filter(|_| view_intact); - match plan_increment( - &source_ds, - source_version, - watermark, - recorded_ts, - full, - definition, + let rows_written = Arc::new(AtomicU64::new(0)); + let unnest = super::physical_unnest(&definition)?; + let stream = compute_stream( + &source, + &definition, &inputs, + unnest.as_ref(), + physical, + rows_written.clone(), ) - .await - { - Some(increment) => { - let reconciled = incremental( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - increment, - definition, - &inputs, - unnest.as_ref(), - persist, - watermark, - expected_incarnation, - ) - .await?; - match reconciled { - Some(result) => Ok(result), - // The delta was too large to reconcile in bounded memory. - None => { - rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - persist.is_some(), - expected_incarnation, - ) - .await - } - } - } - None => { - rebuild( - view_native, - &view_ds, - &source_ds, - source_version, - source_ts, - definition, - &inputs, - unnest.as_ref(), - persist.is_some(), - expected_incarnation, - ) - .await - } - } -} - -/// The source fragments whose rows are new since the watermark, or `None` -/// where the view has to rebuild. Two tiers: the transaction walk is exact -/// where it applies; the fragment-signature check is the fallback for deltas -/// the walk cannot read, and under it any fragment churn rebuilds. -async fn plan_increment( - source_ds: &Dataset, - source_version: u64, - watermark: Option, - recorded_ts: Option, - full: bool, - definition: &MaterializedViewDefinition, - inputs: &[String], -) -> Option { - if full { - return None; - } - let watermark = watermark?; - if watermark > source_version { - return None; - } - let old = source_ds.checkout_version(watermark).await.ok()?; - // A recreated source reuses version numbers, never their timestamps: a - // mismatch means the watermark describes a different incarnation. - if recorded_ts != Some(old.manifest.timestamp_nanos) { - return None; - } - let old_ids: HashSet = old.get_fragments().iter().map(|f| f.id() as u64).collect(); - let live: Vec = source_ds - .get_fragments() - .iter() - .map(|f| f.metadata().clone()) - .collect(); - - if let Some(delta) = appends_and_rewrites(source_ds, watermark, source_version).await { - // A rewrite that consumed a fragment neither present at the watermark - // nor produced by an earlier rewrite swallowed a mid-delta append; - // its rows cannot be told apart from already-materialized ones. - let folded = delta - .rewritten - .iter() - .any(|id| !old_ids.contains(id) && !delta.produced.contains(id)); - if folded { - return None; - } - // An update rewrites a whole fragment. If it touched one the watermark - // never saw, that fragment holds rows appended since -- new rows the - // update did not change, which the recompute does not cover and which - // this fragment's exclusion from the append set would drop. - if delta - .updated_in_place - .iter() - .any(|id| !old_ids.contains(id)) - { - return None; - } - // Rows past the cap left every delta when the watermark advanced, so - // a capped view cannot reconcile a removal incrementally. The rebuild - // is cheap for the same reason it is capped: the scan stops there. - if definition.limit.is_some() && (delta.deleted_rows || delta.updated_rows) { - return None; - } - // Legacy storage cannot serve the row-version columns update - // discovery scans; deletes and appends need none of them. - if delta.updated_rows - && source_ds.manifest.data_storage_format.lance_file_format() == ConcreteFileVersion::V1 - { - return None; - } - // Every other fragment new at head is an append: appends and rewrites - // are the only operations in the delta that add fragments, and the - // rewrite outputs are already-materialized rows rearranged. - return Some(Increment { - appended: live - .into_iter() - .filter(|f| !old_ids.contains(&f.id) && !delta.produced.contains(&f.id)) - .collect(), - evict_deleted: delta.deleted_rows, - replace_updated: delta.updated_rows, - }); - } - - is_pure_append(&old, source_ds, &relevant_field_ids(source_ds, inputs)).then(|| Increment { - appended: live - .into_iter() - .filter(|f| !old_ids.contains(&f.id)) - .collect(), - evict_deleted: false, - replace_updated: false, + .await?; + let committed = replace_fragments(view_ds, stream).await?; + let version = committed.version().version; + native.dataset.update(committed); + Ok(RefreshMaterializedViewResult { + mode: RefreshMode::Rebuild, + rows_written: rows_written.load(Ordering::Relaxed), + source_version, + version, }) } -/// Version gap beyond which per-version transaction reads stop being cheaper -/// than one fragment scan. -const MAX_TRANSACTION_WALK: u64 = 512; - -/// Fragment ids moved by the `Rewrite` operations of the delta. Only rewrite -/// ids are real in transaction files (an Append's are placeholders assigned -/// at commit), so appends are derived as new-at-head minus rewrite outputs. -/// What an incremental refresh must do to bring the view up to date. -struct Increment { - /// Source fragments whose rows are not in the view yet. - appended: Vec, - /// Source rows left the range, so the view holds rows to evict. - evict_deleted: bool, - /// Source rows changed in the range, so the view holds rows to replace. - replace_updated: bool, -} - -struct TxnDelta { - /// Fragments consumed by `Rewrite` operations. - rewritten: HashSet, - /// Fragments produced by `Rewrite` operations. - produced: HashSet, - /// The delta removed source rows, so the view holds rows to evict. - deleted_rows: bool, - /// The delta changed source rows in place, so the view holds rows to - /// recompute. - updated_rows: bool, - /// Fragments an update modified in place, as opposed to produced. - updated_in_place: HashSet, -} - -/// Read the delta from the transaction log, `None` where it holds anything -/// but appends and rewrites or cannot be read; `None` only sends the caller -/// to a slower check. `ReserveFragments` moves no rows and rides along. -async fn appends_and_rewrites(cur: &Dataset, from: u64, to: u64) -> Option { - if to <= from || to - from > MAX_TRANSACTION_WALK { - return None; +/// Atomically replace all fragments of a materialized view with a complete +/// query result. This is shared with Sophon's SQL refresh job. +pub async fn replace_materialized_view_fragments( + view: &Table, + stream: SendableRecordBatchStream, +) -> Result { + let native = view.as_native().ok_or_else(|| Error::NotSupported { + message: "materialized views are supported only on local tables".into(), + })?; + native.dataset.ensure_mutable()?; + let lock = refresh_lock(native.dataset.get().await?.uri()); + let _guard = lock.lock().await; + native.dataset.reload().await?; + let dataset = native.dataset.get().await?.as_ref().clone(); + if super::read_definition_sql(&dataset.schema().metadata)?.is_none() { + return Err(Error::NotAMaterializedView { + name: view.name().to_string(), + }); } - let mut delta = TxnDelta { - rewritten: HashSet::new(), - produced: HashSet::new(), - deleted_rows: false, - updated_rows: false, - updated_in_place: HashSet::new(), + let committed = replace_fragments(dataset, stream).await?; + let version = committed.version().version; + native.dataset.update(committed); + Ok(version) +} + +async fn replace_fragments(dataset: Dataset, stream: SendableRecordBatchStream) -> Result { + let dataset = Arc::new(dataset); + let read_version = dataset.version().version; + let removed_fragment_ids = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.id() as u64) + .collect(); + let write = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted_stream(stream) + .await?; + let Operation::Append { + fragments: new_fragments, + } = write.operation + else { + return Err(Error::Runtime { + message: "expected an append while staging replacement rows".into(), + }); }; - for version in (from + 1)..=to { - let Ok(Some(txn)) = cur.read_transaction_by_version(version).await else { - return None; - }; - match txn.operation { - Operation::Append { .. } | Operation::ReserveFragments { .. } => {} - // A delete removes source rows without changing the ones that - // remain, so the view's other rows stay valid: the refresh - // evicts exactly the ids that left. - Operation::Delete { .. } => delta.deleted_rows = true, - // Update outputs carry no new rows, so they are excluded from - // the append set like rewrite outputs; the changed rows are - // replaced individually below. + let inserted_rows_filter = refresh_conflict_filter(&dataset)?; + let committed = CommitBuilder::new(WriteDestination::Dataset(dataset)) + .execute(Transaction::new( + read_version, Operation::Update { removed_fragment_ids, + updated_fragments: Vec::new(), new_fragments, - updated_fragments, - .. - } => { - delta.updated_rows = true; - // merge_insert reaches here too, and its by-source arm deletes - // rows rather than changing them. - delta.deleted_rows = true; - delta.rewritten.extend(removed_fragment_ids.iter().copied()); - // Only pre-existing fragment ids are real here; created ones - // are placeholders. Rewritten rows are excluded by creation - // version below, not by fragment identity. - delta - .produced - .extend(updated_fragments.iter().map(|f| f.id)); - delta - .updated_in_place - .extend(updated_fragments.iter().map(|f| f.id)); - let _ = new_fragments; - } - Operation::Rewrite { groups, .. } => { - for group in groups { - delta - .rewritten - .extend(group.old_fragments.iter().map(|f| f.id)); - delta - .produced - .extend(group.new_fragments.iter().map(|f| f.id)); - } - } - _ => return None, - } + fields_modified: Vec::new(), + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: None, + inserted_rows_filter: Some(inserted_rows_filter), + updated_fragment_offsets: None, + }, + None, + )) + .await?; + if committed.version().version != read_version + 1 { + return Err(Error::Runtime { + message: format!( + "a concurrent commit raced this refresh (view version {}, expected {})", + committed.version().version, + read_version + 1 + ), + }); } - Some(delta) + Ok(committed) } -/// Fallback pure-append check: every old fragment still present with an -/// identical signature over the columns the view reads. Compaction, deletes -/// and updates each break it and force a rebuild; a change to a column the -/// view does not read leaves it alone, which is what lets this tier pass -/// deltas the transaction walk cannot. -fn is_pure_append(old: &Dataset, cur: &Dataset, relevant: &HashSet) -> bool { - let signature = |fragment: &lance::dataset::fragment::FileFragment| { - fragment_signature(fragment.metadata(), relevant) - }; - let current: HashSet<(u64, String)> = cur.get_fragments().iter().map(signature).collect(); - old.get_fragments() - .iter() - .all(|fragment| current.contains(&signature(fragment))) +const REFRESH_TOKEN_ID: u64 = u64::MAX; + +fn refresh_conflict_filter(dataset: &Dataset) -> Result { + let field_id = dataset + .schema() + .fields + .first() + .map(|field| field.id) + .ok_or_else(|| Error::Runtime { + message: "a materialized view must have at least one field".into(), + })?; + let mut filter = KeyExistenceFilterBuilder::new(vec![field_id]); + filter + .insert(KeyValue::UInt64(REFRESH_TOKEN_ID)) + .map_err(|error| Error::Runtime { + message: format!( + "failed to build the materialized-view refresh conflict filter: {error}" + ), + })?; + Ok(filter.build()) } -/// A fragment's identity as the view observes it: data files and overlays -/// touching the columns it reads, plus the deletion file. Overlays change no -/// file path, so they must be part of the signature. -fn fragment_signature(metadata: &Fragment, relevant: &HashSet) -> (u64, String) { - let touches_relevant = - |fields: &[i32]| relevant.is_empty() || fields.iter().any(|id| relevant.contains(id)); - let mut files: Vec<&str> = metadata - .files - .iter() - .filter(|file| touches_relevant(&file.fields)) - .map(|file| file.path.as_str()) - .collect(); - files.sort_unstable(); - let mut overlays: Vec = metadata - .overlays - .iter() - .filter(|overlay| touches_relevant(&overlay.data_file.fields)) - .map(|overlay| format!("{}@{}", overlay.data_file.path, overlay.committed_version)) - .collect(); - overlays.sort_unstable(); - ( - metadata.id, - format!( - "{}|{}|{:?}", - files.join(","), - overlays.join(","), - metadata.deletion_file - ), - ) -} - -/// Field ids (with struct descendants) of the source columns the view reads. -fn relevant_field_ids(source: &Dataset, inputs: &[String]) -> HashSet { - fn collect(field: &lance_core::datatypes::Field, ids: &mut HashSet) { - ids.insert(field.id); - for child in &field.children { - collect(child, ids); - } - } - let mut ids = HashSet::new(); - for input in inputs { - if let Some(field) = source.schema().field(input) { - collect(field, &mut ids); - } - } - ids -} - -/// Whether two plannings of a view compute the same rows and columns: the -/// same source, unnest and limit, and expressions the planner reads as the -/// same logical expression, whatever their spelling. A definition that does -/// not plan compares as different. -fn same_meaning( - source_schema: &SchemaRef, - stored: &MaterializedViewDefinition, - replanned: &MaterializedViewDefinition, - unnest: Option<&super::ViewUnnest>, -) -> bool { - if stored.source_table != replanned.source_table - || stored.source_namespace != replanned.source_namespace - || stored.lateral != replanned.lateral - || stored.limit != replanned.limit - || stored.projections.len() != replanned.projections.len() - || stored.filter.is_some() != replanned.filter.is_some() - || stored.group_by.len() != replanned.group_by.len() - { - return false; - } - // Lance's planner does not read aggregates; a grouped definition is - // compared in its canonical spelling, which planning produced. - if stored.is_grouped() { - return stored.projections == replanned.projections - && stored.group_by == replanned.group_by - && stored.filter == replanned.filter; - } - let schema = match unnest { - None => source_schema.clone(), - Some(unnest) => match super::flattened_schema(source_schema, unnest) { - Ok(schema) => schema, - Err(_) => return false, - }, - }; - let planner = Planner::new(schema); - let same_expr = |a: &str, b: &str| match (planner.parse_expr(a), planner.parse_expr(b)) { - (Ok(a), Ok(b)) => a == b, - _ => false, - }; - stored - .projections - .iter() - .zip(&replanned.projections) - .all(|(a, b)| a.output == b.output && same_expr(&a.expression, &b.expression)) - && match (&stored.filter, &replanned.filter) { - (Some(a), Some(b)) => same_expr(a, b), - _ => true, - } -} - -/// Reject MemWAL/LSM state on a refresh participant: un-compacted tiers are -/// invisible to the fragment-planned refresh scan. An active write spec and -/// retained rows both disqualify; shard directories on storage are the -/// durable evidence of the latter. +/// Reject MemWAL/LSM state because an ordinary dataset scan cannot see its +/// un-compacted tiers. pub(crate) async fn ensure_no_mem_wal(dataset: &Dataset, role: &str, name: &str) -> Result<()> { let retained = !dataset.list_mem_wal_latest_shard_ids().await?.is_empty(); if retained || dataset.mem_wal_index_details().await?.is_some() { return Err(Error::NotSupported { - message: format!( - "{role} '{name}' has an LSM write spec or retained un-compacted \ - rows: rows in un-compacted tiers are invisible to refresh" - ), + message: format!("{role} '{name}' has an LSM write spec or retained un-compacted rows"), }); } Ok(()) } -/// The table refresh scans: the staging table when the query calls a -/// Function in FROM position, otherwise the query's source. -pub(super) async fn open_source( - view: &Table, - definition: &MaterializedViewDefinition, - staging: Option<&super::StagingBinding>, -) -> Result { +async fn open_source(view: &Table, definition: &MaterializedViewDefinition) -> Result { let database = view.database_opt().ok_or_else(|| Error::InvalidInput { message: "the view was not opened through a database connection".into(), })?; - let (name, namespace_path) = match staging { - Some(staging) => (staging.table.clone(), staging.namespace.clone()), - None => ( - definition.source_table.clone(), - definition.source_namespace.clone(), - ), - }; let source = database .open_table(OpenTableRequest { - name, - namespace_path, + name: definition.source_table.clone(), + namespace_path: definition.source_namespace.clone(), index_cache_size: None, lance_read_params: None, location: None, @@ -747,584 +318,7 @@ pub(super) async fn open_source( let native = source.as_native().ok_or_else(|| Error::NotSupported { message: "materialized views are supported only on local tables".into(), })?; - let dataset = native.dataset.get().await?.as_ref().clone(); - if definition.vector_source.is_none() && !dataset.manifest.uses_stable_row_ids() { - return Err(Error::InvalidInput { - message: format!( - "source table '{}' does not have stable row ids; it is not the \ - table this view was declared over", - definition.source_table - ), - }); - } - Ok(dataset) -} - -#[allow(clippy::too_many_arguments)] -async fn incremental( - view_native: &NativeTable, - view_ds: &Dataset, - source_ds: &Dataset, - source_version: u64, - source_ts: u128, - increment: Increment, - definition: &MaterializedViewDefinition, - inputs: &[String], - unnest: Option<&super::ViewUnnest>, - persist: Option<&MaterializedViewDefinition>, - watermark: Option, - expected_incarnation: Option<&str>, -) -> Result> { - let new_fragments = increment.appended; - let watermark_version = watermark.unwrap_or(0); - // Provenance ids this refresh removes: dropped rows, plus changed rows - // recomputed in the same commit. One view row per source row makes the - // eviction exact. Staged, not committed: removals ride with the rows - // that replace them, so a reader never sees the view without either. - let mut eviction = Eviction::new(view_ds, EVICTION_CHUNK); - let mut updated_rows = false; - if (increment.evict_deleted || increment.replace_updated) - && let Some(watermark) = watermark - { - let delta = source_ds - .delta() - .with_begin_version(watermark) - .with_end_version(source_version) - .build()?; - // Reconciling holds the delta's provenance ids in staged deletion - // vectors; past the cap, the streamed rebuild is the bounded path. - let cap = eviction_rebuild_cap(); - let mut evicted = 0usize; - if increment.evict_deleted { - let mut stream = delta.get_deleted_row_ids().await?; - while let Some(batch) = stream.try_next().await? { - let ids = row_ids_of(&batch)?; - evicted += ids.len(); - if evicted > cap { - return Ok(None); - } - eviction.push(ids).await?; - } - } - if increment.replace_updated { - // Ids only: `get_updated_rows` carries every column of every - // updated row, and discovery needs none of them. - let mut scanner = source_ds.scan(); - scanner.with_row_id().project(&[ROW_CREATED_AT_VERSION])?; - // A fixed bound: the configured default could make one discovery - // batch arbitrarily large before the fallback cap is consulted. - scanner.batch_size(8192); - scanner.filter(&format!( - "{ROW_CREATED_AT_VERSION} <= {watermark} - AND {ROW_LAST_UPDATED_AT_VERSION} > {watermark} - AND {ROW_LAST_UPDATED_AT_VERSION} <= {source_version}" - ))?; - let mut stream = scanner.try_into_stream().await?; - while let Some(batch) = stream.try_next().await? { - let ids = row_ids_of(&batch)?; - evicted += ids.len(); - if evicted > cap { - return Ok(None); - } - updated_rows |= !ids.is_empty(); - eviction.push(ids).await?; - } - } - } - let eviction = eviction.finish().await?; - - // The cap counts rows already materialized, in first-materialized order. - let remaining = match definition.limit { - Some(limit) => { - let held = view_ds.count_rows(None).await? as u64; - Some(limit.saturating_sub(held)) - } - None => None, - }; - - let mut result = RefreshMaterializedViewResult { - mode: RefreshMode::Incremental, - rows_written: 0, - source_version, - version: view_ds.version().version, - }; - let nothing_to_add = (new_fragments.is_empty() && !updated_rows) || remaining == Some(0); - if nothing_to_add && eviction.is_none() { - result.version = stamp_watermark( - view_native, - view_ds.clone(), - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - } - // Rows left but none arrive: the removals still have to be published. - if nothing_to_add { - let filter = refresh_filter(&empty_keys(view_ds)?)?; - let published = publish( - view_ds, - eviction, - Vec::new(), - Some(filter), - expected_incarnation, - ) - .await?; - result.version = stamp_watermark( - view_native, - published, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - } - - // Appends carry the view's schema as it stands; the watermark moves in a - // follow-up commit (see the module docs for the crash window). - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let rows_written = Arc::new(AtomicU64::new(0)); - // compute_stream counts what it produces; the truncation below can drop - // some of that, so the written count comes from the tee instead. - let computed = Arc::new(AtomicU64::new(0)); - let mut stream = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - fragments: Some(new_fragments), - // An update rewrites whole fragments, so a fragment new at head - // can hold rows the view already has. Their creation version - // does not change, so it -- not fragment identity -- says which - // rows are new. - created_after: increment.replace_updated.then_some(watermark_version), - limit: remaining, - ..Default::default() - }, - schema.clone(), - computed.clone(), - ) - .await?; - - // The updated rows' current values, computed the same way and appended - // in the same commit as the new fragments' rows. - if updated_rows { - let recomputed = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - updated_between: Some((watermark_version, source_version)), - ..Default::default() - }, - schema.clone(), - computed.clone(), - ) - .await?; - stream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - recomputed.chain(stream), - )); - } - - // Nothing survived the filter: the watermark still has to advance or the - // same fragments would be rescanned forever, but any removals still do. - let Some(first) = stream.try_next().await? else { - let published = if eviction.is_some() { - publish( - view_ds, - eviction, - Vec::new(), - Some(refresh_filter(&empty_keys(view_ds)?)?), - expected_incarnation, - ) - .await? - } else { - view_ds.clone() - }; - result.version = stamp_watermark( - view_native, - published, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - return Ok(Some(result)); - }; - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter([Ok(first)]).chain(stream), - )); - - // Any two refreshes of one view must not both commit: the filter's - // shared token makes lance reject the loser on key overlap however - // their planned rows relate. - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(view_ds)?, - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), rows_written.clone()); - - let ds = Arc::new(view_ds.clone()); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone())) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { - fragments: new_fragments, - } = write_txn.operation - else { - return Err(Error::Runtime { - message: "expected an append when staging the view's new rows".into(), - }); - }; - let filter = refresh_filter(&keys)?; - let appended = publish( - view_ds, - eviction, - new_fragments, - Some(filter), - expected_incarnation, - ) - .await?; - result.rows_written = rows_written.load(Ordering::Relaxed); - result.version = stamp_watermark( - view_native, - appended, - source_version, - source_ts, - persist, - expected_incarnation, - ) - .await?; - Ok(Some(result)) -} - -#[allow(clippy::too_many_arguments)] -async fn rebuild( - view_native: &NativeTable, - view_ds: &Dataset, - source_ds: &Dataset, - source_version: u64, - source_ts: u128, - definition: &MaterializedViewDefinition, - inputs: &[String], - unnest: Option<&super::ViewUnnest>, - persist_definition: bool, - expected_incarnation: Option<&str>, -) -> Result { - let rows_written = Arc::new(AtomicU64::new(0)); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let stream = compute_stream( - source_ds, - definition, - inputs, - unnest, - RowScope { - limit: definition.limit, - ..Default::default() - }, - schema, - rows_written.clone(), - ) - .await?; - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(view_ds)?, - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - // Every rebuild is one fragment swap, indexed or not: an Update commit - // carries no schema metadata, so it cannot erase a definition update - // that raced in the way an overwrite (which adopts its stream's schema) - // durably would -- and it must land on the planned generation or abort. - let replaced = - replace_retaining_indices(view_ds.clone(), stream, keys, expected_incarnation).await?; - let version = stamp_watermark( - view_native, - replaced, - source_version, - source_ts, - persist_definition.then_some(definition), - expected_incarnation, - ) - .await?; - Ok(RefreshMaterializedViewResult { - mode: RefreshMode::Rebuild, - rows_written: rows_written.load(Ordering::Relaxed), - source_version, - version, - }) -} - -/// Replace all of the view's data in one commit that retains its index -/// definitions: new fragments staged uncommitted, one `Update` removing every -/// old fragment. `Update` prunes index bitmaps only for modified fields and -/// none are modified here, so readers never see the view unindexed or empty. -async fn replace_retaining_indices( - view_ds: Dataset, - stream: SendableRecordBatchStream, - keys: Arc>, - expected_incarnation: Option<&str>, -) -> Result { - let ds = Arc::new(view_ds); - let read_version = ds.version().version; - #[cfg(test)] - tests::hold_before_publish(ds.uri()).await; - ensure_incarnation(&ds, expected_incarnation, ds.uri()).await?; - let removed_fragment_ids: Vec = ds.get_fragments().iter().map(|f| f.id() as u64).collect(); - - let write_txn = InsertBuilder::new(WriteDestination::Dataset(ds.clone())) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await?; - let Operation::Append { - fragments: new_fragments, - } = write_txn.operation - else { - return Err(Error::Runtime { - message: "expected an append when staging the view's replacement rows".into(), - }); - }; - - // Built only now: the tee fills as the staging drains the stream. - let filter = refresh_filter(&keys)?; - let transaction = Transaction::new( - read_version, - Operation::Update { - removed_fragment_ids, - updated_fragments: Vec::new(), - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - // Two refreshes that materialized the same source rows must not - // both land -- a raced first rebuild would double the view. - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - ); - let committed = CommitBuilder::new(WriteDestination::Dataset(ds)) - .execute(transaction) - .await?; - if committed.version().version != read_version + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {}); the - refresh is unrecorded and the next one will rebuild", - committed.version().version - ), - }); - } - Ok(committed) -} - -/// Record that the view now reflects `source_version`, including the view -/// Refuse to act on a view that is not `expected`'s incarnation, judged from -/// the latest stored manifest. Not a commit condition; see -/// `RefreshMaterializedViewBuilder::expect_incarnation`. -pub(super) async fn ensure_incarnation( - view_ds: &Dataset, - expected: Option<&str>, - what: &str, -) -> Result<()> { - let Some(expected) = expected else { - return Ok(()); - }; - let mut latest = view_ds.clone(); - latest.checkout_latest().await?; - match latest.schema().metadata.get(INCARNATION_META_KEY) { - Some(actual) if actual == expected => Ok(()), - Some(_) => Err(Error::Runtime { - message: format!( - "materialized view '{what}' is not the incarnation this refresh was \ - requested for: it was dropped and recreated" - ), - }), - None => Err(Error::Runtime { - message: format!( - "materialized view '{what}' carries no incarnation token: its schema \ - metadata was replaced since the token was captured" - ), - }), - } -} - -/// version this very commit produces. The version is predicted and then -/// verified; on a mismatch another commit raced in between, and the stamp -/// ABORTS rather than certify that commit as the refresh's own generation. -/// The view is left visibly unstamped, so the next refresh rebuilds. -pub(super) async fn stamp_watermark( - view_native: &NativeTable, - mut dataset: Dataset, - source_version: u64, - source_ts: u128, - definition: Option<&MaterializedViewDefinition>, - expected_incarnation: Option<&str>, -) -> Result { - ensure_incarnation(&dataset, expected_incarnation, dataset.uri()).await?; - let predicted = dataset.version().version + 1; - // A view with no token (declared before tokens existed, or its metadata - // replaced wholesale) starts a new incarnation here. - let incarnation = dataset - .schema() - .metadata - .get(INCARNATION_META_KEY) - .cloned() - .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let mut metadata = vec![(INCARNATION_META_KEY.to_string(), Some(incarnation))]; - if let Some(definition) = definition { - metadata.push(( - DEFINITION_META_KEY.to_string(), - Some(definition_to_metadata(definition)?), - )); - } - metadata.extend([ - ( - SOURCE_VERSION_META_KEY.to_string(), - Some(source_version.to_string()), - ), - ( - SOURCE_VERSION_TS_META_KEY.to_string(), - Some(source_ts.to_string()), - ), - ( - REFRESHED_AT_MS_META_KEY.to_string(), - Some(now_ms().to_string()), - ), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]); - dataset.update_schema_metadata(metadata).await?; - let actual = dataset.version().version; - if actual != predicted { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {actual}, \ - expected {predicted}); the refresh is unrecorded and the next \ - one will rebuild" - ), - }); - } - view_native.dataset.update(dataset); - Ok(predicted) -} - -fn now_ms() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() -} - -/// Evaluate the definition over `source`, restricted to `fragments` when -/// given, as batches in the view's schema. The filter is pushed into the -/// scan; projections are `(name, expression)` pairs, never spliced into SQL; -/// [`SOURCE_ROW_ID_COLUMN`] is filled by the scan's row id. -/// Which source rows a compute pass reads. -#[derive(Default)] -struct RowScope { - /// Read only these fragments. - fragments: Option>, - /// Read only rows created after this version. - created_after: Option, - /// Read only rows changed in place after the first version and no later - /// than the second. - updated_between: Option<(u64, u64)>, - /// Stop after this many rows. - limit: Option, -} - -/// Whether every commit on the view after `recorded` is a fill of its -/// computed columns: a column rewrite or data replacement touching only -/// those fields and neither adding nor removing rows, or the freshness -/// stamp a fill leaves on them. A version whose transaction cannot be read -/// is not proven, so it counts as drift. -async fn only_computed_rewrites_since(view_ds: &Dataset, recorded: u64) -> Result { - // A fill may write any field under a computed column, so the whole - // subtree counts, not only the root. - let physical = ArrowSchema::from(view_ds.schema()); - fn subtree(field: &lance_core::datatypes::Field, ids: &mut Vec) { - ids.push(field.id as u32); - for child in &field.children { - subtree(child, ids); - } - } - let mut computed_fields = Vec::new(); - for column in computed_columns(&physical) { - if let Some(field) = view_ds.schema().field(&column.name) { - subtree(field, &mut computed_fields); - } - } - if computed_fields.is_empty() { - return Ok(false); - } - for version in recorded + 1..=view_ds.version().version { - let Some(transaction) = view_ds.read_transaction_by_version(version).await? else { - return Ok(false); - }; - let fill = match &transaction.operation { - Operation::Update { - removed_fragment_ids, - new_fragments, - fields_modified, - update_mode: Some(UpdateMode::RewriteColumns), - .. - } => { - removed_fragment_ids.is_empty() - && new_fragments.is_empty() - && !fields_modified.is_empty() - && fields_modified - .iter() - .all(|field| computed_fields.contains(field)) - } - // What `refresh_column` commits for a SQL declaration. - Operation::DataReplacement { replacements } => { - !replacements.is_empty() - && replacements.iter().all(|group| { - !group.1.fields.is_empty() - && group - .1 - .fields - .iter() - .all(|field| computed_fields.contains(&(*field as u32))) - }) - } - // The stamp `refresh_column` writes after its fill (see - // `table::freshness`): field metadata on computed columns, no data. - Operation::UpdateConfig { - config_updates: None, - table_metadata_updates: None, - schema_metadata_updates: None, - field_metadata_updates, - } => { - !field_metadata_updates.is_empty() - && field_metadata_updates - .keys() - .all(|field| computed_fields.contains(&(*field as u32))) - } - _ => false, - }; - if !fill { - return Ok(false); - } - } - Ok(true) + Ok(native.dataset.get().await?.as_ref().clone()) } async fn compute_stream( @@ -1332,58 +326,23 @@ async fn compute_stream( definition: &MaterializedViewDefinition, inputs: &[String], unnest: Option<&super::ViewUnnest>, - scope: RowScope, schema: SchemaRef, rows_written: Arc, ) -> Result { if definition.is_grouped() { return super::grouped::stream(source, definition, inputs, schema, rows_written).await; } - let RowScope { - fragments, - created_after, - updated_between, - limit, - } = scope; let mut scanner = source.scan(); - if let Some(fragments) = fragments { - scanner.with_fragments(fragments); + if let Some(filter) = definition.filter.as_deref().filter(|_| unnest.is_none()) { + scanner.filter(filter)?; } - scanner.with_row_id(); - // Narrowing keeps the definition's filter, so a row updated out of the - // view simply does not come back. Changed rows are named by the predicate - // `DatasetDelta::get_updated_rows` uses, not its streamed ids: an id list - // grows with the delta, this does not. - let updated_filter = updated_between.map(|(from, to)| { - format!( - "{ROW_CREATED_AT_VERSION} <= {from} \ - AND {ROW_LAST_UPDATED_AT_VERSION} > {from} \ - AND {ROW_LAST_UPDATED_AT_VERSION} <= {to}" - ) - }); - let created_filter = - created_after.map(|version| format!("{ROW_CREATED_AT_VERSION} > {version}")); - let clauses: Vec = definition - .filter - .clone() - .filter(|_| unnest.is_none()) - .map(|f| format!("({f})")) - .into_iter() - .chain(updated_filter) - .chain(created_filter) - .collect(); - if !clauses.is_empty() { - scanner.filter(&clauses.join(" AND "))?; - } - // An expanded view cannot project or filter in the scan: both read the - // unnested element, which exists only after the per-batch expansion. let expanded = match unnest { Some(unnest) => Some(UnnestPlan::new(source, definition, inputs, unnest)?), None => { let transforms: Vec<(&str, &str)> = definition .projections .iter() - .map(|p| (p.output.as_str(), p.expression.as_str())) + .map(|projection| (projection.output.as_str(), projection.expression.as_str())) .collect(); scanner.project_with_transform(&transforms)?; None @@ -1392,54 +351,43 @@ async fn compute_stream( if let Some(expanded) = &expanded { scanner.project(&expanded.raw_inputs)?; } - // A scan reads a limit of zero as no limit at all, so a view capped at - // nothing is answered without one. - if limit == Some(0) { + if definition.limit == Some(0) { return Ok(Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), + schema, futures::stream::empty(), ))); } - if let Some(limit) = limit { + if let Some(limit) = definition.limit { let limit = i64::try_from(limit).map_err(|_| Error::InvalidInput { message: format!("view limit {limit} exceeds the maximum of {}", i64::MAX), })?; scanner.limit(Some(limit), None)?; } - - let out_schema = schema.clone(); + let output_schema = schema.clone(); let mapped = scanner.try_into_stream().await?.map(move |batch| { - let batch = batch.map_err(|e| DataFusionError::External(Box::new(e)))?; + let batch = batch.map_err(|error| DataFusionError::External(Box::new(error)))?; let batch = match &expanded { - None => batch, Some(expanded) => expanded.apply(&batch)?, + None => batch, }; - let batch = to_view_batch(&batch, &out_schema)?; + let batch = to_view_batch(&batch, &output_schema)?; rows_written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); Ok(batch) }); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, mapped))) } -/// `batch` in the view's physical schema: the row id becomes the -/// provenance column and computed columns are written NULL for their owner -/// to fill. pub(super) fn to_view_batch( batch: &RecordBatch, - out_schema: &SchemaRef, + output_schema: &SchemaRef, ) -> datafusion::common::Result { - let mut columns = Vec::with_capacity(out_schema.fields().len()); - for field in out_schema.fields() { + let mut columns = Vec::with_capacity(output_schema.fields().len()); + for field in output_schema.fields() { if computed_column_from_field(field).is_some() { columns.push(new_null_array(field.data_type(), batch.num_rows())); continue; } - let name = if field.name() == SOURCE_ROW_ID_COLUMN { - ROW_ID - } else { - field.name() - }; - let column = batch.column_by_name(name).ok_or_else(|| { + let column = batch.column_by_name(field.name()).ok_or_else(|| { DataFusionError::Internal(format!( "view column '{}' is not produced by the view's definition", field.name() @@ -1447,13 +395,9 @@ pub(super) fn to_view_batch( })?; columns.push(column.clone()); } - Ok(RecordBatch::try_new(out_schema.clone(), columns)?) + Ok(RecordBatch::try_new(output_schema.clone(), columns)?) } -/// The post-scan half of an unnested view's refresh: the scan reads -/// `raw_inputs` plus the row id, and each batch is unnested on the list -/// column, filtered, then projected by expressions typed against -/// `read_schema`, where the element sits under the alias. struct UnnestPlan { column: String, raw_inputs: Vec, @@ -1469,7 +413,6 @@ impl UnnestPlan { inputs: &[String], unnest: &super::ViewUnnest, ) -> Result { - // Whole root columns: a nested input is projected by the expression. let mut raw_inputs: Vec = inputs .iter() .map(|input| super::root(input).to_string()) @@ -1478,9 +421,7 @@ impl UnnestPlan { raw_inputs.sort(); raw_inputs.dedup(); let flattened = super::flattened_schema(&ArrowSchema::from(source.schema()), unnest)?; - // Physical expressions index columns by position, so the schema is - // exactly the scan's output: `raw_inputs` in order, then the row id. - let mut read_fields = Vec::with_capacity(raw_inputs.len() + 1); + let mut read_fields = Vec::with_capacity(raw_inputs.len()); for name in &raw_inputs { let name = if *name == unnest.column { &unnest.alias @@ -1494,25 +435,27 @@ impl UnnestPlan { })?; read_fields.push(field.clone()); } - read_fields.push(ArrowField::new(ROW_ID, DataType::UInt64, false)); let read_schema = Arc::new(ArrowSchema::new(read_fields)); let planner = Planner::new(read_schema.clone()); let physical = |what: &str, sql: &str| -> Result> { - let err = |e: lance::Error| Error::Runtime { - message: format!("{what}: {e}"), + let error = |error: lance::Error| Error::Runtime { + message: format!("{what}: {error}"), }; - let parsed = planner.parse_expr(sql).map_err(err)?; - let optimized = planner.optimize_expr(parsed).map_err(err)?; - planner.create_physical_expr(&optimized).map_err(err) + let parsed = planner.parse_expr(sql).map_err(error)?; + let optimized = planner.optimize_expr(parsed).map_err(error)?; + planner.create_physical_expr(&optimized).map_err(error) }; - let mut projections = Vec::with_capacity(definition.projections.len()); - for projection in &definition.projections { - let expr = physical( - &format!("view column '{}'", projection.output), - &projection.expression, - )?; - projections.push((projection.output.clone(), expr)); - } + let projections = definition + .projections + .iter() + .map(|projection| { + physical( + &format!("view column '{}'", projection.output), + &projection.expression, + ) + .map(|expression| (projection.output.clone(), expression)) + }) + .collect::>()?; let filter = definition .filter .as_deref() @@ -1529,8 +472,7 @@ impl UnnestPlan { fn apply(&self, batch: &RecordBatch) -> datafusion::common::Result { let unnested = unnest_batch(batch, &self.column) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - // Same columns, renamed: the list column is now the element under the alias. + .map_err(|error| DataFusionError::External(Box::new(error)))?; let unnested = RecordBatch::try_new(self.read_schema.clone(), unnested.columns().to_vec())?; let unnested = match &self.filter { None => unnested, @@ -1544,25 +486,20 @@ impl UnnestPlan { arrow_select::filter::filter_record_batch(&unnested, keep)? } }; - let mut columns = Vec::with_capacity(self.projections.len() + 1); - for (output, expr) in &self.projections { - let value = expr.evaluate(&unnested)?.into_array(unnested.num_rows())?; - columns.push((output.clone(), value)); - } - let row_id = unnested - .column_by_name(ROW_ID) - .expect("scan carries the row id") - .clone(); - columns.push((ROW_ID.to_string(), row_id)); + let columns = self + .projections + .iter() + .map(|(output, expression)| { + expression + .evaluate(&unnested)? + .into_array(unnested.num_rows()) + .map(|value| (output.clone(), value)) + }) + .collect::>>()?; Ok(RecordBatch::try_from_iter(columns)?) } } -/// Expand `list_column` one row per element, repeating every other column -/// for each element; an empty or null list contributes no rows. The list -/// column is replaced by its element type, so a projection reads the -/// element's fields as `alias.field` after this. This is the row-cardinality -/// step of an `expanded_select` view, applied per batch on the scan stream. fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { use arrow_array::{Array, ListArray, UInt32Array}; use arrow_select::take::take; @@ -1583,30 +520,24 @@ fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { batch.column(list_index).data_type() ), })?; - - // One take index per element, naming the source row it came from. A null - // list has no elements; its offsets are equal, so it repeats nothing. let offsets = list.value_offsets(); let mut repeat = Vec::with_capacity(list.values().len()); for row in 0..list.len() { if list.is_valid(row) { - let count = (offsets[row + 1] - offsets[row]) as usize; - repeat.extend(std::iter::repeat_n(row as u32, count)); + repeat.extend(std::iter::repeat_n( + row as u32, + (offsets[row + 1] - offsets[row]) as usize, + )); } } let repeat = UInt32Array::from(repeat); - // The flattened elements, in the same order as `repeat`: only the ranges - // valid rows cover, so a null row's stale range (if any) is skipped. - let elements = { - let mut ranges = Vec::new(); - for row in 0..list.len() { - if list.is_valid(row) { - ranges.extend((offsets[row] as u32)..(offsets[row + 1] as u32)); - } + let mut ranges = Vec::new(); + for row in 0..list.len() { + if list.is_valid(row) { + ranges.extend((offsets[row] as u32)..(offsets[row + 1] as u32)); } - take(list.values().as_ref(), &UInt32Array::from(ranges), None)? - }; - + } + let elements = take(list.values().as_ref(), &UInt32Array::from(ranges), None)?; let mut fields = Vec::with_capacity(batch.num_columns()); let mut columns = Vec::with_capacity(batch.num_columns()); for (index, field) in batch.schema().fields().iter().enumerate() { @@ -1636,3168 +567,38 @@ fn unnest_batch(batch: &RecordBatch, list_column: &str) -> Result { Ok(RecordBatch::try_new(schema, columns)?) } -/// Commit the view's removals and additions as one change, on the exact -/// generation the refresh planned from. Lance rejects an overlapping -/// provenance key, but an unrelated write to the view is not a key conflict, -/// so the generation is checked here too. -async fn publish( - view_ds: &Dataset, - eviction: Option<(Vec, Vec)>, - new_fragments: Vec, - keys: Option, - expected_incarnation: Option<&str>, -) -> Result { - let planned = view_ds.version().version; - #[cfg(test)] - tests::hold_before_publish(view_ds.uri()).await; - #[cfg(test)] - tests::hold_until_peers_planned(); - ensure_incarnation(view_ds, expected_incarnation, view_ds.uri()).await?; - let (updated_fragments, removed_fragment_ids) = eviction.unwrap_or_default(); - let committed = CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - planned, - Operation::Update { - removed_fragment_ids, - updated_fragments, - new_fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: keys, - updated_fragment_offsets: None, - }, - None, - )) - .await?; - if committed.version().version != planned + 1 { - return Err(Error::Runtime { - message: format!( - "a concurrent commit raced this refresh (view version {}); - the refresh is unrecorded and the next one will rebuild", - committed.version().version - ), - }); - } - Ok(committed) -} - -/// A delta batch's row id column, borrowed: a delta batch is as large as one -/// fragment's deletions, so copying it out would be the unbounded step the -/// chunking exists to avoid. -fn row_ids_of(batch: &RecordBatch) -> Result<&UInt64Array> { - let column = batch.column_by_name(ROW_ID).ok_or_else(|| Error::Runtime { - message: format!("'{ROW_ID}' is missing from a delta batch"), - })?; - column - .as_primitive_opt::() - .ok_or_else(|| Error::Runtime { - message: "row ids are not UInt64".into(), - }) -} - -/// A provenance id no source row can hold, carried in every refresh's -/// inserted-rows filter: any two refreshes of one view overlap on it, so -/// the loser of a race conflicts at commit whatever rows each planned. -const REFRESH_TOKEN_ID: u64 = u64::MAX; - -/// A builder with no collected keys, for publishes that only remove rows. -fn empty_keys(view_ds: &Dataset) -> Result>> { - Ok(Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new( - vec![source_row_id_field_id(view_ds)?], - )))) -} - -/// An inserted-rows filter holding the refresh token plus whatever the -/// tee collected. -fn refresh_filter(keys: &Arc>) -> Result { - let mut keys = keys.lock().map_err(|_| Error::Runtime { - message: "the provenance key filter was poisoned mid-refresh".into(), - })?; - keys.insert(KeyValue::UInt64(REFRESH_TOKEN_ID)) - .map_err(|e| Error::Runtime { - message: format!("failed to mark the refresh's filter: {e}"), - })?; - Ok(keys.build()) -} - -/// Reconciled ids past this fall back to the streamed rebuild, whose memory -/// does not grow with the delta. -fn eviction_rebuild_cap() -> usize { - #[cfg(test)] - if let Some(cap) = tests::eviction_cap_override() { - return cap; - } - 4 * 1024 * 1024 -} - -/// Provenance ids per staged eviction: the delete predicate carries one -/// literal per id, so a whole delta at once is unbounded. Each chunk costs a -/// pass over the view's provenance column, which is why it is not smaller. -const EVICTION_CHUNK: usize = 64 * 1024; - -/// Accumulates the view's removals from bounded chunks of provenance ids into -/// the one set of fragment changes the refresh publishes. -struct Eviction { - /// The view as the chunks staged so far leave it. A chunk's delete has to - /// see the deletion vectors the earlier ones wrote, or it stages a - /// fragment that drops them. - snapshot: Dataset, - chunk: usize, - updated: HashMap, - removed: Vec, - pending: Vec, - staged: bool, - /// Largest the buffer ever got, which is the bound under test. - #[cfg(test)] - peak: usize, -} - -impl Eviction { - fn new(view_ds: &Dataset, chunk: usize) -> Self { - Self { - snapshot: view_ds.clone(), - chunk, - updated: HashMap::new(), - removed: Vec::new(), - pending: Vec::with_capacity(chunk), - staged: false, - #[cfg(test)] - peak: 0, - } - } - - /// Fill a chunk at a time. Taking the batch whole and splitting it would - /// hold a fragment's worth of ids and recopy the tail per chunk. - async fn push(&mut self, ids: &UInt64Array) -> Result<()> { - for id in ids.values() { - self.pending.push(*id); - #[cfg(test)] - { - self.peak = self.peak.max(self.pending.len()); - } - if self.pending.len() == self.chunk { - self.flush().await?; - } - } - Ok(()) - } - - /// Stage what is buffered, keeping the buffer's allocation. - async fn flush(&mut self) -> Result<()> { - let mut chunk = std::mem::take(&mut self.pending); - let staged = self.stage(&chunk).await; - chunk.clear(); - self.pending = chunk; - staged - } - - /// The staged fragment changes, or `None` where nothing was evicted. - async fn finish(mut self) -> Result, Vec)>> { - if !self.pending.is_empty() { - self.flush().await?; - } - if !self.staged { - return Ok(None); - } - let mut updated: Vec = self.updated.into_values().collect(); - updated.sort_unstable_by_key(|f| f.id); - Ok(Some((updated, self.removed))) - } - - async fn stage(&mut self, ids: &[u64]) -> Result<()> { - let (updated, removed) = stage_eviction(&self.snapshot, ids).await?; - self.snapshot = advance(&self.snapshot, &updated); - for fragment in updated { - self.updated.insert(fragment.id, fragment); - } - self.removed.extend(removed); - self.staged = true; - Ok(()) - } -} - -/// The view as staged removals leave it, without committing them. Fragments -/// are replaced in place so the fragment bitmap stays in step; an emptied one -/// is left alone, since the delta names each provenance id only once. -fn advance(view_ds: &Dataset, updated: &[Fragment]) -> Dataset { - if updated.is_empty() { - return view_ds.clone(); - } - let by_id: HashMap = updated.iter().map(|f| (f.id, f)).collect(); - let fragments = view_ds - .manifest - .fragments - .iter() - .map(|f| by_id.get(&f.id).map_or_else(|| f.clone(), |u| (*u).clone())) - .collect(); - let mut manifest = view_ds.manifest.as_ref().clone(); - manifest.fragments = Arc::new(fragments); - let mut snapshot = view_ds.clone(); - snapshot.manifest = Arc::new(manifest); - snapshot -} - -/// The fragment changes that remove the view's rows for `ids`, staged rather -/// than committed so they can ride in the refresh's single data commit. -async fn stage_eviction(view_ds: &Dataset, ids: &[u64]) -> Result<(Vec, Vec)> { - // An expression rather than SQL text: the id list is a value here, not a - // predicate string that grows with the delta and has to be parsed. - let predicate = col(SOURCE_ROW_ID_COLUMN).in_list( - ids.iter() - .map(|id| lit(ScalarValue::UInt64(Some(*id)))) - .collect(), - false, - ); - let staged = DeleteBuilder::from_expr(Arc::new(view_ds.clone()), predicate) - .execute_uncommitted() - .await?; - let Operation::Delete { - updated_fragments, - deleted_fragment_ids, - .. - } = staged.transaction.operation - else { - return Err(Error::Runtime { - message: "expected a delete when staging the view's evictions".into(), - }); - }; - Ok((updated_fragments, deleted_fragment_ids)) -} - -fn source_row_id_field_id(view_ds: &Dataset) -> Result { - view_ds - .schema() - .field(SOURCE_ROW_ID_COLUMN) - .map(|f| f.id) - .ok_or_else(|| Error::Runtime { - message: format!("the view has no '{SOURCE_ROW_ID_COLUMN}' column"), - }) -} - -/// Tee the provenance ids of everything written into `keys`. -fn collect_source_row_ids( - stream: SendableRecordBatchStream, - keys: Arc>, - written: Arc, -) -> SendableRecordBatchStream { - let schema = stream.schema(); - let mapped = stream.map(move |batch| { - let batch = batch?; - let column = batch - .column_by_name(SOURCE_ROW_ID_COLUMN) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "'{SOURCE_ROW_ID_COLUMN}' is missing from the rows being written" - )) - })? - .as_primitive_opt::() - .ok_or_else(|| { - DataFusionError::Internal(format!("'{SOURCE_ROW_ID_COLUMN}' is not a uint64")) - })?; - let mut keys = keys - .lock() - .map_err(|_| DataFusionError::Internal("provenance key filter poisoned".into()))?; - for id in column.values() { - keys.insert(KeyValue::UInt64(*id)) - .map_err(|e| DataFusionError::Internal(e.to_string()))?; - } - written.fetch_add(batch.num_rows() as u64, Ordering::Relaxed); - Ok(batch) - }); - Box::pin(RecordBatchStreamAdapter::new(schema, mapped)) -} - #[cfg(test)] -pub(crate) mod tests { - - /// Park a refresh between planning and publication so a test can move - /// the view underneath it. Inert unless [`DRIFT_TARGET`] names this view. - pub(in crate::materialized_view) async fn hold_before_publish(uri: &str) { - { - let mut target = DRIFT_TARGET.lock().unwrap(); - if target.as_deref() != Some(uri) { - return; - } - // Take it: memory:// uris are relative and repeat across tests, so - // leaving it armed would park an unrelated refresh forever. - *target = None; - } - DRIFT_PLANNED.notify_one(); - DRIFT_RELEASED.notified().await; - } - - /// The rendezvous below is one global pair, so the cases that use it run - /// one at a time rather than trading each other's signals. - pub(super) static EVICTION_CAP: StdMutex> = StdMutex::new(None); - pub(super) fn eviction_cap_override() -> Option { - *EVICTION_CAP.lock().unwrap() - } - - pub(in crate::materialized_view) static DRIFT_LOCK: tokio::sync::Mutex<()> = - tokio::sync::Mutex::const_new(()); - pub(in crate::materialized_view) static DRIFT_TARGET: StdMutex> = - StdMutex::new(None); - pub(in crate::materialized_view) static DRIFT_PLANNED: tokio::sync::Notify = - tokio::sync::Notify::const_new(); - pub(in crate::materialized_view) static DRIFT_RELEASED: tokio::sync::Notify = - tokio::sync::Notify::const_new(); - - /// Block until every participant in a cross-process race has planned and - /// staged its write, so the commits they then attempt genuinely contend - /// rather than depending on the scheduler to overlap them. Inert unless - /// `MV_RACE_SYNC` names a directory shared by the participants. - pub(super) fn hold_until_peers_planned() { - let (Ok(dir), Ok(tag), Ok(peers)) = ( - std::env::var("MV_RACE_SYNC"), - std::env::var("MV_RACE_TAG"), - std::env::var("MV_RACE_PEERS"), - ) else { - return; - }; - let dir = std::path::PathBuf::from(dir); - let peers: usize = peers.parse().unwrap(); - std::fs::write(dir.join(format!("planned-{tag}")), b"1").unwrap(); - let deadline = std::time::Instant::now() + std::time::Duration::from_secs(120); - while planned_count(&dir) < peers { - assert!( - std::time::Instant::now() < deadline, - "peers never reached the commit boundary" - ); - std::thread::sleep(std::time::Duration::from_millis(2)); - } - } - - fn planned_count(dir: &std::path::Path) -> usize { - std::fs::read_dir(dir) - .map(|entries| { - entries - .filter_map(|e| e.ok()) - .filter(|e| e.file_name().to_string_lossy().starts_with("planned-")) - .count() - }) - .unwrap_or(0) - } - use arrow_array::{Int32Array, record_batch}; - use futures::TryStreamExt; - use lance::dataset::NewColumnTransform; - use lance_file::version::LanceFileVersion; +mod tests { + use arrow_array::record_batch; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use super::*; use crate::connect; - use crate::connection::Connection; - use crate::index::Index; - use crate::index::scalar::BTreeIndexBuilder; - use crate::materialized_view::MaterializedView; - use crate::query::{ExecutableQuery, QueryBase, Select}; - use crate::table::{CompactionOptions, OptimizeAction}; - - async fn db_with_source(values: Vec) -> (Connection, Table) { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, values)).unwrap(); - let table = conn - .create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - (conn, table) - } - - async fn doubled_view(conn: &Connection) -> MaterializedView { - conn.create_materialized_view("doubled", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap() - } - - /// A refreshed doubled view over a source holding `values`. - async fn refreshed_doubled(values: Vec) -> (Connection, Table, MaterializedView) { - let (conn, source) = db_with_source(values).await; - let view = doubled_view(&conn).await; - view.refresh().execute().await.unwrap(); - (conn, source, view) - } - - /// The per-batch expansion behind an `expanded_select` view: every - /// element of the list column becomes a row, the other columns repeat - /// for each, and an empty or null list contributes no rows at all -- - /// which is exactly "zero rows out" for a table-valued function. - /// Four documents with a `list>` column named - /// `c`: doc 1 has two chunks, doc 2 none, doc 3 a null list, doc 4 one. - fn chunked_batch() -> RecordBatch { - use arrow_array::builder::{Int32Builder, ListBuilder, StringBuilder, StructBuilder}; - use arrow_array::{ArrayRef, Int64Array}; - use arrow_schema::{DataType, Field, Fields}; - - let element_fields = Fields::from(vec![ - Field::new("chunk", DataType::Utf8, true), - Field::new("ordinal", DataType::Int32, true), - ]); - let mut list = ListBuilder::new(StructBuilder::new( - element_fields, - vec![ - Box::new(StringBuilder::new()), - Box::new(Int32Builder::new()), - ], - )); - for chunks in [Some(vec!["a", "b"]), Some(vec![]), None, Some(vec!["c"])] { - match chunks { - Some(chunks) => { - for (i, c) in chunks.iter().enumerate() { - let s = list.values(); - s.field_builder::(0).unwrap().append_value(c); - s.field_builder::(1) - .unwrap() - .append_value(i as i32); - s.append(true); - } - list.append(true); - } - None => list.append(false), - } - } - let meta = arrow_array::StructArray::from(vec![( - Arc::new(Field::new("title", DataType::Utf8, true)), - Arc::new(arrow_array::StringArray::from(vec!["t1", "t2", "t3", "t4"])) as ArrayRef, - )]); - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef, - ), - ("meta", Arc::new(meta) as ArrayRef), - ("c", Arc::new(list.finish()) as ArrayRef), - ]) - .unwrap() - } - - #[test] - fn unnest_repeats_siblings_per_element_and_drops_empty_lists() { - use arrow_array::{StringArray, StructArray}; - use arrow_schema::{DataType, Field, Fields}; - - let batch = chunked_batch(); - let element_fields = Fields::from(vec![ - Field::new("chunk", DataType::Utf8, true), - Field::new("ordinal", DataType::Int32, true), - ]); - let out = unnest_batch(&batch, "c").unwrap(); - assert_eq!(out.num_rows(), 3, "{out:?}"); - let ids: Vec = out["id"] - .as_primitive::() - .values() - .to_vec(); - assert_eq!(ids, [1, 1, 4]); - let element = out["c"].as_any().downcast_ref::().unwrap(); - let chunks: Vec<&str> = element - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect(); - assert_eq!(chunks, ["a", "b", "c"]); - let ordinals: Vec = element - .column(1) - .as_primitive::() - .values() - .to_vec(); - assert_eq!(ordinals, [0, 1, 0]); - // the element column is now the struct itself, not a list of it - assert_eq!( - out.schema().field_with_name("c").unwrap().data_type(), - &DataType::Struct(element_fields) - ); - } - - /// A Function in FROM position is refreshed from its staging table: the - /// query on the view stays the one the user wrote, the staging binding - /// says which table and list column refresh reads, and rows come out one - /// per element as with UNNEST. Without a staging, a local database - /// refuses the query rather than guess. - #[tokio::test] - async fn a_function_in_from_position_refreshes_from_its_staging() { - use arrow_array::StringArray; - - let conn = connect("memory://").execute().await.unwrap(); - let staging = conn - .create_table("docs__chunk", chunked_batch()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let query = - "SELECT id AS doc, e.chunk AS text, e.ordinal FROM docs, chunk(meta.title, 2) AS e"; - let definition = MaterializedViewDefinition::from_sql(query).unwrap(); - - let err = crate::materialized_view::prepare_definition(&staging, definition.clone()) - .await - .unwrap_err(); - assert!( - matches!(&err, Error::InvalidInput { message } if message.contains("reads 'docs'")), - "{err:?}" - ); - - let view = crate::materialized_view::prepare_staged_definition(&staging, definition, "c") - .await - .unwrap() - .create("chunks") - .await - .unwrap(); - assert_eq!(view.definition().to_sql(), query); - let metadata = view.table().schema().await.unwrap().metadata().clone(); - assert_eq!( - crate::materialized_view::read_staging(&metadata).unwrap(), - Some(crate::materialized_view::StagingBinding { - table: "docs__chunk".into(), - namespace: Vec::new(), - column: "c".into(), - }) - ); - - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "ordinal").await, [0, 0, 1]); - let batches = view - .table() - .query() - .select(Select::columns(&["text"])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let out = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); - let texts: Vec<&str> = out["text"] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect(); - assert_eq!(texts, ["a", "b", "c"]); - - // The reopened view reads the same logical query and refreshes - // incrementally from the staging. - staging.add(chunked_batch()).execute().await.unwrap(); - let reopened = conn.open_materialized_view("chunks").await.unwrap(); - assert_eq!(reopened.definition().to_sql(), query); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(reopened.table(), "ordinal").await, [0, 0, 0, 0, 1, 1]); - } - - /// An expanded view materializes one row per list element, with the - /// projections reading the element through the alias and the other - /// source columns repeated alongside; sources with no elements yield - /// no rows. The lineage is the list column, so a change to it is what - /// drives incremental refresh. - #[tokio::test] - async fn an_expanded_view_materializes_one_row_per_element() { - use arrow_array::{Int64Array, StringArray}; - - let conn = connect("memory://").execute().await.unwrap(); - let source = conn - .create_table("docs", chunked_batch()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let mut view = crate::materialized_view::prepare_definition( - &source, - MaterializedViewDefinition::from_sql( - "SELECT id AS doc, meta.title AS title, e.ordinal + 1 AS nth \ - FROM docs, UNNEST(c) AS e WHERE e.ordinal < 5", - ) - .unwrap(), - ) - .await - .unwrap(); - // A computed column's input read through the alias is an element - // field; the recorded source input is the list column. - let text = view.input_column("e.chunk").unwrap(); - assert!(view.definition.lateral.is_some()); - let view = view.create("chunks").await.unwrap(); - view.refresh().execute().await.unwrap(); - - let batches = view - .table() - .query() - .select(Select::columns(&["doc", &text, "title"])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let out = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); - let docs: Vec = out["doc"] - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec(); - let strings = |column: &str| -> Vec { - out[column] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .map(str::to_string) - .collect() - }; - assert_eq!(docs, [1, 1, 4]); - assert_eq!(strings(&text), ["a", "b", "c"]); - assert_eq!(strings("title"), ["t1", "t1", "t4"]); - assert_eq!(read(view.table(), "nth").await, [1, 1, 2]); - - // Appended documents expand incrementally; the existing rows stay. - source.add(chunked_batch()).execute().await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "nth").await, [1, 1, 1, 1, 2, 2]); - } - - async fn read(table: &Table, column: &str) -> Vec { - let batches = table - .query() - .select(Select::columns(&[column])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let mut values: Vec = batches - .iter() - .flat_map(|batch| { - batch[column] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .flatten() - .collect::>() - }) - .collect(); - values.sort(); - values - } - - async fn append(table: &Table, values: Vec) { - let batch = record_batch!(("x", Int32, values)).unwrap(); - table.add(batch).execute().await.unwrap(); - } #[tokio::test] - async fn test_first_refresh_materializes_the_view() { - let (conn, _) = db_with_source(vec![1, 2, 3]).await; - let view = doubled_view(&conn).await; - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 3); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - - // The watermark survives on the stored schema, not just the handle. - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - let again = reopened.refresh().execute().await.unwrap(); - assert_eq!(again.mode, RefreshMode::NoOp); - assert_eq!(again.rows_written, 0); - } - - #[tokio::test] - async fn test_filter_selects_the_source_rows() { - let (conn, _) = db_with_source(vec![1, 20, 3, 40]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") + async fn raced_first_full_refresh_lands_one_result() { + let directory = tempfile::tempdir().unwrap(); + let connection = connect(directory.path().to_str().unwrap()) .execute() .await .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "x").await, vec![20, 40]); - } - - #[tokio::test] - async fn test_mixed_case_filter_is_canonicalized_for_lineage_and_refresh() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("id", Int32, [1, 2, 3]), - ("PartyAbbrev", Utf8, ["D", "R", "D"]) - ) - .unwrap(); - conn.create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) + let source_batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); + connection + .create_table("src", source_batch.clone()) .execute() .await .unwrap(); - conn.create_materialized_view("democrats", "src") - .with_no_data(true) - .select([("id", "id")]) - .only_if(r#""PartyAbbrev" = 'D'"#) - .execute() - .await - .unwrap(); - - // Reopen from schema metadata so these assertions cover the stored - // predicate and lineage, not only the declaration-time handle. - let view = conn.open_materialized_view("democrats").await.unwrap(); - assert_eq!( - view.definition().filter.as_deref(), - Some("`PartyAbbrev` = 'D'") - ); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "id").await, vec![1, 3]); - } - - /// A legacy layout whose query means what the canonical one means is - /// rewritten in the current layout on the next refresh, without a - /// rebuild: the rows it certified are the rows the query produces. - #[tokio::test] - async fn a_legacy_layout_with_the_same_meaning_is_rewritten_without_a_rebuild() { - let (conn, source, view) = refreshed_doubled(vec![1]).await; - let legacy = serde_json::json!({ - "kind": "select", - "source_table": "src", - "projections": [ - {"output": "x", "expression": "`x`"}, - {"output": "twice", "expression": "x*2"}, - ], - "inputs": ["x"], - }) - .to_string(); - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let predicted = dataset.version().version + 1; - dataset - .update_schema_metadata([ - (DEFINITION_META_KEY.to_string(), Some(legacy)), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]) - .await - .unwrap(); - native.dataset.update(dataset); - - append(&source, vec![2]).await; - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(reopened.table(), "twice").await, vec![2, 4]); - let stored: serde_json::Value = serde_json::from_str( - &reopened.table().schema().await.unwrap().metadata()[DEFINITION_META_KEY], - ) - .unwrap(); - assert_eq!( - stored, - serde_json::json!({ - "kind": "query", - "format": 1, - "query": "SELECT x, x * 2 AS twice FROM src", - }) - ); - } - - #[tokio::test] - async fn test_legacy_raw_filter_rebuilds_and_persists_canonical_definition() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!( - ("id", Int32, [1, 2, 3]), - ("PartyAbbrev", Utf8, ["D", "R", "D"]) - ) - .unwrap(); - conn.create_table("legacy_src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("legacy_view", "legacy_src") - .with_no_data(true) - .select([("id", "id")]) - .only_if(r#""PartyAbbrev" = 'X'"#) - .execute() - .await - .unwrap(); - assert_eq!(view.refresh().execute().await.unwrap().rows_written, 0); - - // Model a definition and up-to-date watermark written before filter - // canonicalization was applied to materialized views. - let legacy = serde_json::json!({ - "kind": "select", - "source_table": "legacy_src", - "projections": [{"output": "id", "expression": "id"}], - "filter": r#""PartyAbbrev" = 'D'"#, - "inputs": ["id"], - }) - .to_string(); - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let predicted = dataset.version().version + 1; - dataset - .update_schema_metadata([ - (DEFINITION_META_KEY.to_string(), Some(legacy)), - ( - VIEW_VERSION_META_KEY.to_string(), - Some(predicted.to_string()), - ), - ]) - .await - .unwrap(); - native.dataset.update(dataset); - - let reopened = conn.open_materialized_view("legacy_view").await.unwrap(); - let result = reopened.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(read(reopened.table(), "id").await, vec![1, 3]); - - // A fresh handle proves the migration was stored alongside the new - // watermark and therefore happens only once. - let migrated = conn.open_materialized_view("legacy_view").await.unwrap(); - assert_eq!( - migrated.definition().filter.as_deref(), - Some("`PartyAbbrev` = 'D'") - ); - assert_eq!( - migrated.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!(read(migrated.table(), "id").await, vec![1, 3]); - } - - #[tokio::test] - async fn test_append_refreshes_incrementally() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - append(&source, vec![5]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 10]); - } - - #[tokio::test] - async fn test_incremental_applies_the_filter() { - let (conn, source) = db_with_source(vec![1, 20]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![3, 30]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "x").await, vec![20, 30]); - } - - /// The watermark has to advance even when no appended row matches, or the - /// same fragments would be rescanned by every later refresh. - #[tokio::test] - async fn test_incremental_with_nothing_matching_advances_the_watermark() { - let (conn, source) = db_with_source(vec![20]).await; - let view = conn - .create_materialized_view("big", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 10") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![1, 2]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - - let again = view.refresh().execute().await.unwrap(); - assert_eq!(again.mode, RefreshMode::NoOp); - } - - /// Unlike a computed column, a view reflects source mutation: an update - /// rebuilds rather than going stale. - #[tokio::test] - async fn test_update_replaces_the_rows_it_changed() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - // An update changes rows in place, so the view replaces exactly - // those rows and leaves the rest of what it holds alone. - source - .update() - .column("x", "20") - .only_if("x = 2") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1, "only the changed row is recomputed"); - assert_eq!(read(view.table(), "twice").await, vec![2, 6, 40]); - } - - /// Legacy storage cannot serve the row-version columns update discovery - /// scans; appends stay incremental, updates rebuild rather than fail. - #[tokio::test] - async fn test_a_legacy_storage_source_is_reconciled() { - let conn = connect("memory://").execute().await.unwrap(); - let batch = record_batch!(("x", Int32, [1, 2, 3])).unwrap(); - let source = conn - .create_table("legacy_src", batch) - .write_options(crate::table::WriteOptions { - lance_write_params: Some(lance::dataset::WriteParams { - enable_stable_row_ids: true, - data_storage_version: Some(LanceFileVersion::Legacy), - ..Default::default() - }), - }) - .execute() - .await - .unwrap(); - let view = conn - .create_materialized_view("legacy_doubled", "legacy_src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - source - .add(record_batch!(("x", Int32, [4])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6, 8]); - - source - .update() - .column("x", "20") - .only_if("x = 2") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 6, 8, 40]); - } - - #[tokio::test] - async fn test_delete_evicts_the_view_rows_it_removed() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - // A delete removes source rows without changing the ones that - // remain, so the view evicts exactly those rows and keeps the rest - // rather than recomputing every row it already held. - source.delete("x = 2").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 6]); - - // A delete and an append in one span: both are applied. - source.delete("x = 1").await.unwrap(); - source - .add(record_batch!(("x", Int32, vec![4])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![6, 8]); - } - - /// merge_insert commits an `Update`, and its by-source arm removes rows - /// rather than changing them, so the classifier must treat that - /// transaction form as a source of deletions. - #[tokio::test] - async fn test_merge_insert_by_source_delete_evicts_the_view_rows() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - let batch = record_batch!(("x", Int32, vec![1, 3])).unwrap(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = source.merge_insert(&["x"]); - merge.when_not_matched_by_source_delete(None); - merge.execute(Box::new(reader)).await.unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 6]); - } - - /// A refresh may only certify the generation it planned from. A write - /// that lands between planning and publication is drift the refresh did - /// not account for, so it aborts rather than stamp it as materialized. - #[tokio::test(flavor = "multi_thread")] - async fn test_refresh_aborts_rather_than_certify_view_drift() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("drifting_view", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - let certified = source_version_of(view.table()).await; - - // The source loses a row, so the next refresh plans an eviction. - source.delete("x = 2").await.unwrap(); - - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = tokio::spawn(async move { view.refresh().execute().await }); - - // Move the view once the refresh has planned against it. - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - let drifted = conn.open_table("drifting_view").execute().await.unwrap(); - drifted.delete("twice = 6").await.unwrap(); - DRIFT_RELEASED.notify_one(); - - // Publishing removals and additions as one change makes the drift a - // conflict lance itself rejects; a pure append, which touches no - // existing fragment, still relies on the generation check. - let err = refreshing.await.unwrap().unwrap_err(); - let message = err.to_string(); - assert!( - message.contains("raced this refresh") || message.contains("preempted by concurrent"), - "got {err:?}" - ); - // The watermark still names the generation that was actually proven. - assert_eq!(source_version_of(&drifted).await, certified); - } - - async fn source_version_of(table: &Table) -> Option { - table - .schema() - .await - .unwrap() - .metadata() - .get(SOURCE_VERSION_META_KEY) - .cloned() - } - - /// A transaction file carries placeholder ids for the fragments it - /// creates. Into an empty source those collide with the ids the commit - /// assigns, so treating them as already-materialized drops the very - /// first rows while the watermark still advances past them. - #[tokio::test] - async fn test_merge_into_an_empty_source_is_materialized() { - let (_conn, source, view) = refreshed_doubled(vec![]).await; - assert_eq!(read(view.table(), "twice").await, Vec::::new()); - - let batch = record_batch!(("x", Int32, vec![1, 2])).unwrap(); - let reader = arrow_array::RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); - let mut merge = source.merge_insert(&["x"]); - merge - .when_matched_update_all(None) - .when_not_matched_insert_all(); - merge.execute(Box::new(reader)).await.unwrap(); - - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - // A second refresh must not double them either. - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// A refresh publishes what it removes and what it adds as one change, - /// so an update never exposes the view without the rows it is replacing. - #[tokio::test(flavor = "multi_thread")] - async fn test_an_update_is_never_visible_as_a_gap() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("atomic_view", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - source - .update() - .column("x", "x + 10") - .execute() - .await - .unwrap(); - - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = tokio::spawn(async move { view.refresh().execute().await }); - - // Read the view while the refresh is staged but not yet published. - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - let midway = conn.open_table("atomic_view").execute().await.unwrap(); - assert_eq!( - read(&midway, "twice").await, - vec![2, 4, 6], - "the pre-refresh rows must still be there in full" - ); - DRIFT_RELEASED.notify_one(); - - refreshing.await.unwrap().unwrap(); - let after = conn.open_table("atomic_view").execute().await.unwrap(); - assert_eq!(read(&after, "twice").await, vec![22, 24, 26]); - } - - async fn provenance_by_x(table: &Table) -> HashMap { - let batches = table - .query() - .select(Select::columns(&["x", SOURCE_ROW_ID_COLUMN])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - batches - .iter() - .flat_map(|batch| { - let xs = batch["x"].as_any().downcast_ref::().unwrap(); - let ids = batch[SOURCE_ROW_ID_COLUMN].as_primitive::(); - (0..batch.num_rows()) - .map(|i| (xs.value(i), ids.value(i))) - .collect::>() - }) - .collect() - } - - /// One delta batch is as large as a fragment's deletions, so it arrives - /// well past a chunk: it has to be drained a chunk at a time without ever - /// holding the batch, and the passes' deletion vectors have to accumulate - /// rather than replace each other. - #[tokio::test] - async fn test_a_chunked_eviction_removes_every_row_it_names() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2, 3, 4, 5, 6]).await; - - let native = view.table().as_native().unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - let provenance = provenance_by_x(view.table()).await; - - let mut eviction = Eviction::new(&view_ds, 2); - let batch = UInt64Array::from_iter_values([1, 2, 3, 4].map(|x| provenance[&x])); - eviction.push(&batch).await.unwrap(); - assert_eq!( - eviction.peak, 2, - "a batch past the chunk was buffered whole" - ); - let staged = eviction.finish().await.unwrap(); - assert!(staged.is_some(), "four ids over a chunk of two stage twice"); - publish(&view_ds, staged, Vec::new(), None, None) - .await - .unwrap(); - native.dataset.reload().await.unwrap(); - - assert_eq!(read(view.table(), "x").await, vec![5, 6]); - } - - /// Two first rebuilds materialize the same source rows; the loser must - /// key-conflict at commit and land nothing, or the view doubles. - #[tokio::test] - async fn test_a_raced_first_rebuild_lands_nothing() { - let _guard = DRIFT_LOCK.lock().await; - let (conn, _) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("raced_rebuild", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - let native = view.table().as_native().unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - let uri = view_ds.uri().to_string(); - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let racing = tokio::spawn(async move { view.refresh().execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the rebuild never reached the publication boundary"); - - // The other process's first rebuild, reduced to its commit. Its - // source rows are DISJOINT from ours -- the sentinel alone must make - // the two rebuilds conflict. - let batch = record_batch!( - ("x", Int32, [8, 9]), - ("twice", Int32, [16, 18]), - ("__source_row_id", UInt64, [7u64, 8]) - ) - .unwrap(); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - schema - .fields() - .iter() - .map(|f| batch.column_by_name(f.name()).unwrap().clone()) - .collect(), - ) - .unwrap(); - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - futures::stream::iter([Ok(batch)]), - )); - let keys = Arc::new(StdMutex::new(KeyExistenceFilterBuilder::new(vec![ - source_row_id_field_id(&view_ds).unwrap(), - ]))); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await - .unwrap(); - let Operation::Append { fragments } = write_txn.operation else { - panic!("expected an append"); - }; - let filter = { - let mut keys = keys.lock().unwrap(); - keys.insert(KeyValue::UInt64(super::REFRESH_TOKEN_ID)) - .unwrap(); - keys.build() - }; - CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - view_ds.version().version, - Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments: Vec::new(), - new_fragments: fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - )) - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = racing.await.unwrap().unwrap_err(); - // The loser lands nothing: only the winner's rows remain. - let raced = conn.open_table("raced_rebuild").execute().await.unwrap(); - assert_eq!( - read(&raced, "twice").await, - vec![16, 18], - "the losing rebuild must not union with the winner ({err})" - ); - } - - /// An incremental refresh racing any other refresh must land nothing, - /// even when their planned rows are disjoint: the shared token makes the - /// commits conflict. - #[tokio::test] - async fn test_a_raced_incremental_lands_nothing() { - let _guard = DRIFT_LOCK.lock().await; - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("raced_incremental", "src") - .with_no_data(true) - .select([("x", "x"), ("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - append(&source, vec![4]).await; - - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let view_ds = native.dataset.get().await.unwrap().as_ref().clone(); - *DRIFT_TARGET.lock().unwrap() = Some(view_ds.uri().to_string()); - let racing = tokio::spawn(async move { view.refresh().execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("the refresh never reached the publication boundary"); - - // The concurrent refresh, reduced to its commit: disjoint rows, the - // shared token alone must collide. - let batch = record_batch!( - ("x", Int32, [9]), - ("twice", Int32, [18]), - ("__source_row_id", UInt64, [8u64]) - ) - .unwrap(); - let schema = Arc::new(ArrowSchema::from(view_ds.schema())); - let batch = RecordBatch::try_new( - schema.clone(), - schema - .fields() - .iter() - .map(|f| batch.column_by_name(f.name()).unwrap().clone()) - .collect(), - ) - .unwrap(); - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - schema.clone(), - futures::stream::iter([Ok(batch)]), - )); - let keys = empty_keys(&view_ds).unwrap(); - let stream = collect_source_row_ids(stream, keys.clone(), Arc::new(AtomicU64::new(0))); - let write_txn = InsertBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .with_params(&WriteParams { - mode: WriteMode::Append, - ..Default::default() - }) - .execute_uncommitted_stream(stream) - .await - .unwrap(); - let Operation::Append { fragments } = write_txn.operation else { - panic!("expected an append"); - }; - let filter = refresh_filter(&keys).unwrap(); - CommitBuilder::new(WriteDestination::Dataset(Arc::new(view_ds.clone()))) - .execute(Transaction::new( - view_ds.version().version, - Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments: Vec::new(), - new_fragments: fragments, - fields_modified: Vec::new(), - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: None, - inserted_rows_filter: Some(filter), - updated_fragment_offsets: None, - }, - None, - )) - .await - .unwrap(); - DRIFT_RELEASED.notify_one(); - - let err = racing.await.unwrap().unwrap_err(); - let raced = conn - .open_table("raced_incremental") - .execute() - .await - .unwrap(); - assert_eq!( - read(&raced, "twice").await, - vec![2, 4, 6, 18], - "the losing increment must not land its rows ({err})" - ); - } - - /// Past the eviction cap the refresh falls back to the streamed rebuild, - /// and the result is identical either way. - #[tokio::test] - async fn test_oversized_delta_falls_back_to_rebuild() { - let (conn, source) = db_with_source((1..=20).collect()).await; - let view = doubled_view(&conn).await; - view.refresh().execute().await.unwrap(); - - source.delete("x <= 10").await.unwrap(); - *tests::EVICTION_CAP.lock().unwrap() = Some(4); - let result = view.refresh().execute().await; - *tests::EVICTION_CAP.lock().unwrap() = None; - let result = result.unwrap(); - assert_eq!( - result.mode, - RefreshMode::Rebuild, - "ten evictions, cap of four" - ); - assert_eq!(read(view.table(), "x").await, (11..=20).collect::>()); - - // Under the cap the same shape stays incremental. - source.delete("x = 11").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "x").await, (12..=20).collect::>()); - } - - /// A cap of zero is a view that holds nothing, not a view without a cap. - #[tokio::test] - async fn test_zero_limit_holds_no_rows() { - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("empty", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(0) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 0); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - - // Still nothing after the source grows, and the watermark advances. - append(&source, vec![4]).await; - view.refresh().execute().await.unwrap(); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// An update rewrites a whole fragment. When it lands on one appended - /// since the watermark, that fragment also holds rows the update never - /// touched -- rows the recompute does not cover and the append set no - /// longer reaches. - #[tokio::test] - async fn test_update_touching_a_new_fragment_keeps_its_untouched_rows() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - // One fragment, appended after the watermark, holding both a row the - // update will change and a row it will not. - append(&source, vec![3, 40]).await; - source - .update() - .column("x", "99") - .only_if("x = 3") - .execute() - .await - .unwrap(); - - view.refresh().execute().await.unwrap(); - assert_eq!( - read(view.table(), "twice").await, - vec![2, 4, 80, 198], - "a row appended into the updated fragment went missing" - ); - } - - async fn compact(source: &Table) { - source - .optimize(OptimizeAction::Compact { - options: CompactionOptions::default(), - remap_options: None, - }) - .await - .unwrap(); - } - - /// A compaction rearranges rows without changing them, so it costs the - /// view nothing: the watermark advances and no row is recomputed. - #[tokio::test] - async fn test_compaction_alone_refreshes_incrementally() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - append(&source, vec![3]).await; - view.refresh().execute().await.unwrap(); - - compact(&source).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Rows appended after a compaction are separable through the transaction - /// log: only the appended fragments are computed. - #[tokio::test] - async fn test_append_after_compaction_stays_incremental() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - append(&source, vec![2]).await; - view.refresh().execute().await.unwrap(); - - compact(&source).await; - append(&source, vec![3]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - } - - /// An append swallowed by a later compaction cannot be told apart from - /// the rows the view already holds, so the refresh rebuilds -- once -- - /// rather than duplicate or drop. - #[tokio::test] - async fn test_append_folded_into_compaction_rebuilds() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2]).await; - compact(&source).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// The create-time gate holds across a drop-and-recreate of the source - /// under the same name. - #[tokio::test] - async fn test_refresh_refuses_a_recreated_source_without_stable_row_ids() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - - conn.drop_table("src", &[]).await.unwrap(); - let batch = record_batch!(("x", Int32, [9])).unwrap(); - conn.create_table("src", batch).execute().await.unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!( - matches!(err, Error::InvalidInput { message } if message.contains("stable row ids")) - ); - } - - /// A change to a column the view does not read is not a reason to - /// rebuild: the exact signature check is scoped to the view's inputs. - #[tokio::test] - async fn test_unrelated_column_change_does_not_rebuild() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2]).await; - - source - .add_columns() - .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( - vec![arrow_schema::Field::new( - "unrelated", - arrow_schema::DataType::Int32, - true, - )], - )))) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 0); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - #[tokio::test] - async fn test_full_forces_a_rebuild() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2]).await; - let result = view.refresh().full(true).execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// A cap and incremental reconciliation do not compose: rows skipped at - /// the cap fall behind the watermark, so room freed later cannot be - /// refilled from any delta. A capped view rebuilds instead, which its - /// own cap keeps cheap. - #[tokio::test] - async fn test_limited_view_rebuilds_rather_than_reconcile() { - let (conn, source) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("capped", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(2) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(read(view.table(), "x").await, vec![1, 2]); - - append(&source, vec![3, 4]).await; - source - .update() - .column("x", "11") - .only_if("x = 1") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - let held = read(view.table(), "x").await; - assert_eq!(held.len(), 2, "the cap holds: {held:?}"); - let selectable = read(&source, "x").await; - assert!( - held.iter().all(|x| selectable.contains(x)), - "{held:?} is not a subset of {selectable:?}" - ); - } - - #[tokio::test] - async fn test_limit_caps_the_view() { - let (conn, source) = db_with_source(vec![1, 2, 3]).await; - let view = conn - .create_materialized_view("capped", "src") - .with_no_data(true) - .select([("x", "x")]) - .limit(4) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 3); - - // The cap counts already-held rows, so only one appended row lands. - append(&source, vec![4, 5, 6]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(view.table().count_rows(None).await.unwrap(), 4); - - // At the cap, later appends only move the watermark. - append(&source, vec![7]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// The whole point of the fragment-swap commit: a rebuild must never - /// leave the view without its index definitions. - #[tokio::test] - async fn test_rebuild_retains_indexes() { - let (_conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - - view.table() - .create_index(&["twice"], Index::BTree(BTreeIndexBuilder::default())) - .execute() - .await - .unwrap(); - assert_eq!(view.table().list_indices().await.unwrap().len(), 1); - - source - .update() - .column("x", "x + 10") - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(view.table().list_indices().await.unwrap().len(), 1); - assert_eq!(read(view.table(), "twice").await, vec![22, 24, 26]); - - // The swapped-in rows are reachable through an indexed query. - let batches = view - .table() - .query() - .only_if("twice = 24") - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 1); - } - - #[tokio::test] - async fn test_rebuild_of_an_empty_result_is_an_empty_view() { - let (conn, _) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("none", "src") - .with_no_data(true) - .select([("x", "x")]) - .only_if("x > 100") - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 0); - assert_eq!(view.table().count_rows(None).await.unwrap(), 0); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Provenance: every view row records the source row that produced it. - #[tokio::test] - async fn test_source_row_ids_are_recorded() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2, 3]).await; - - let batches = view - .table() - .query() - .select(Select::columns(&[SOURCE_ROW_ID_COLUMN])) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let total: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total, 3); - for batch in &batches { - assert_eq!(batch[SOURCE_ROW_ID_COLUMN].null_count(), 0); - } - } - - #[tokio::test] - async fn test_dropping_a_source_input_fails_the_refresh() { - let (conn, source) = db_with_source(vec![1]).await; - let view = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .select([("twice", "x * 2")]) - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - - append(&source, vec![2]).await; - source - .add_columns() - .transform(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( - vec![arrow_schema::Field::new( - "y", - arrow_schema::DataType::Int32, - true, - )], - )))) - .execute() - .await - .unwrap(); - source.drop_columns(&["x"]).await.unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!( - matches!(&err, Error::Schema { message } - if message.contains("dropped or renamed") && message.contains("No field named x")), - "{err:?}" - ); - } - - /// A pinned refresh materializes the source as of `version`; catching up - /// to the appends beyond it stays incremental. - #[tokio::test] - async fn test_pinned_refresh_and_catch_up() { - let (conn, source) = db_with_source(vec![1]).await; - let view = doubled_view(&conn).await; - let pinned = source.version().await.unwrap(); - - append(&source, vec![2]).await; - let result = view - .refresh() - .source_version(pinned) - .execute() - .await - .unwrap(); - assert_eq!(result.source_version, pinned); - assert_eq!(read(view.table(), "twice").await, vec![2]); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// Views chain: a view's stable row ids and provenance column make it a - /// source like any other, and the default projection takes its declared - /// columns without copying its provenance. - #[tokio::test] - async fn test_a_view_can_source_another_view() { - let (conn, source) = db_with_source(vec![1, 2, 30]).await; - let first = doubled_view(&conn).await; - first.refresh().execute().await.unwrap(); - - let second = conn - .create_materialized_view("second", "doubled") - .with_no_data(true) - .only_if("twice > 10") - .execute() - .await - .unwrap(); - assert!( - second - .definition() - .projections - .iter() - .all(|p| p.output != SOURCE_ROW_ID_COLUMN) - ); - let result = second.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 1); - assert_eq!(read(second.table(), "twice").await, vec![60]); - - append(&source, vec![50]).await; - first.refresh().execute().await.unwrap(); - let result = second.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(read(second.table(), "twice").await, vec![60, 100]); - } - - /// The watermark speaks only for the state a refresh left behind: a - /// direct write to the view is drift, and the next refresh rebuilds - /// rather than preserving it as current. - #[tokio::test] - async fn test_direct_view_mutation_forces_a_rebuild() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - view.table().delete("x = 1").await.unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// A dropped and recreated source reuses version numbers but never their - /// timestamps; the watermark must not vouch for the replacement's rows. - #[tokio::test] - async fn test_source_recreation_forces_a_rebuild() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - - conn.drop_table("src", &[]).await.unwrap(); - let batch = record_batch!(("x", Int32, [7])).unwrap(); - conn.create_table("src", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![14]); - } - - /// A refresh bound to an incarnation refuses a view dropped and recreated - /// since, even under the same name and definition; the recreated view's - /// own token is accepted, and the token survives a refresh's stamp. - #[tokio::test] - async fn test_refresh_refuses_a_recreated_view_incarnation() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - let token = view.incarnation().unwrap().to_string(); - view.refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap(); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert_eq!(reopened.incarnation(), Some(token.as_str())); - - conn.drop_table("doubled", &[]).await.unwrap(); - let recreated = doubled_view(&conn).await; - assert_ne!(recreated.incarnation(), Some(token.as_str())); - - let err = recreated - .refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("dropped and recreated"), "{err}"); - assert_eq!(read(recreated.table(), "twice").await, Vec::::new()); - - recreated - .refresh() - .expect_incarnation(recreated.incarnation().unwrap()) - .execute() - .await - .unwrap(); - assert_eq!(read(recreated.table(), "twice").await, vec![2]); - } - - /// A cloned declaration creates two physical tables; each gets its own - /// token. - #[tokio::test] - async fn test_cloned_declaration_mints_a_fresh_incarnation_per_create() { - let (conn, source) = db_with_source(vec![1]).await; - let prepared = crate::materialized_view::prepare_declaration( - &source, - Some(&[("x".into(), "x".into()), ("twice".into(), "x * 2".into())]), - None, - None, - ) - .await - .unwrap(); - let replacement = prepared.clone(); - let first = prepared.create("cloned").await.unwrap(); - let first_token = first.incarnation().unwrap().to_string(); - - conn.drop_table("cloned", &[]).await.unwrap(); - let second = replacement.create("cloned").await.unwrap(); - assert_ne!(second.incarnation(), Some(first_token.as_str())); - } - - /// A recreation that lands after planning but before publication is - /// caught by the pre-commit read: the stale refresh fails and the - /// replacement stays empty under its own token. - #[tokio::test(flavor = "multi_thread")] - async fn test_bound_refresh_cannot_publish_into_a_raced_recreation() { - let _serial = DRIFT_LOCK.lock().await; - let (conn, _) = db_with_source(vec![1]).await; - let view = doubled_view(&conn).await; - let token = view.incarnation().unwrap().to_string(); - let uri = view - .table() - .as_native() - .unwrap() - .dataset - .get() - .await - .unwrap() - .uri() - .to_string(); - - *DRIFT_TARGET.lock().unwrap() = Some(uri); - let refreshing = - tokio::spawn(async move { view.refresh().expect_incarnation(token).execute().await }); - tokio::time::timeout(std::time::Duration::from_secs(30), DRIFT_PLANNED.notified()) - .await - .expect("refresh never reached publication"); - - conn.drop_table("doubled", &[]).await.unwrap(); - let replacement = doubled_view(&conn).await; - let replacement_token = replacement.incarnation().unwrap().to_string(); - DRIFT_RELEASED.notify_one(); - - let result = refreshing.await.unwrap(); - assert!(result.is_err(), "the stale refresh unexpectedly succeeded"); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert_eq!(reopened.incarnation(), Some(replacement_token.as_str())); - assert_eq!(read(reopened.table(), "twice").await, Vec::::new()); - } - - /// Replacing the schema metadata wholesale drops the token. A refresh - /// bound to the old token is refused for that reason, not as a - /// recreation; an unbound refresh mints the view a fresh one. - #[tokio::test] - async fn test_a_view_whose_metadata_was_replaced_starts_a_new_incarnation() { - let (conn, _, view) = refreshed_doubled(vec![1]).await; - let token = view.incarnation().unwrap().to_string(); - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(view.definition()).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - assert_eq!( - conn.open_materialized_view("doubled") - .await - .unwrap() - .incarnation(), - None - ); - - let err = view - .refresh() - .expect_incarnation(&token) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("no incarnation token"), "{err}"); - - view.refresh().execute().await.unwrap(); - let reopened = conn.open_materialized_view("doubled").await.unwrap(); - assert!(reopened.incarnation().is_some()); - assert_ne!(reopened.incarnation(), Some(token.as_str())); - } - - /// In-process refreshes of one view serialize: the loser of the race - /// observes the winner's watermark instead of appending the same rows. - #[tokio::test(flavor = "multi_thread")] - async fn test_concurrent_refreshes_do_not_duplicate() { - let (_conn, source, view) = refreshed_doubled(vec![1]).await; - - append(&source, vec![2, 3]).await; - let (a, b) = tokio::join!(view.refresh().execute(), view.refresh().execute()); - let (a, b) = (a.unwrap(), b.unwrap()); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6]); - let modes = [a.mode, b.mode]; - assert!(modes.contains(&RefreshMode::Incremental)); - assert!(modes.contains(&RefreshMode::NoOp)); - } - - /// A second handle's lazy cache must not defeat the lock: after another - /// handle's refresh commits, the stale handle plans from the reloaded - /// state and no-ops instead of appending the same fragments again. - #[tokio::test] - async fn test_a_second_handle_does_not_double_append() { - let (conn, source, view) = refreshed_doubled(vec![1, 2, 3]).await; - let stale = conn.open_materialized_view("doubled").await.unwrap(); - - append(&source, vec![4]).await; - view.refresh().execute().await.unwrap(); - - let result = stale.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::NoOp); - assert_eq!(read(view.table(), "twice").await, vec![2, 4, 6, 8]); - } - - /// A commit racing between a refresh's data commit and its stamp must - /// not be certified as the refresh's generation: the stamp aborts, and - /// the next refresh rebuilds from the drifted state. - #[tokio::test] - async fn test_stamp_aborts_on_a_racing_commit() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let view_native = view.table().as_native().unwrap(); - let stale = view_native.dataset.get().await.unwrap().as_ref().clone(); - view.table().delete("x = 1").await.unwrap(); - - let err = stamp_watermark(view_native, stale, 99, 99, None, None).await; - assert!(err.is_err()); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![2, 4]); - } - - /// What refresh executes and what it stamps must be one generation: a - /// replaced definition wins over whatever a stale handle cached. - #[tokio::test] - async fn test_refresh_uses_the_latest_persisted_definition() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let replacement = crate::materialized_view::MaterializedViewDefinition { - vector_source: None, - source_table: "src".into(), - source_namespace: Vec::new(), - projections: vec![ - crate::materialized_view::ViewProjection { - output: "x".into(), - expression: "x".into(), - }, - crate::materialized_view::ViewProjection { - output: "twice".into(), - expression: "x * 3".into(), - }, - ], - filter: None, - group_by: Vec::new(), - limit: None, - lateral: None, - }; - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(&replacement).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(read(view.table(), "twice").await, vec![3, 6]); - } - - /// A persisted definition that does not produce this view's schema must - /// not refresh at all, let alone be certified. - #[tokio::test] - async fn test_definition_view_schema_mismatch_is_refused() { - let (_conn, _, view) = refreshed_doubled(vec![1, 2]).await; - - let narrower = crate::materialized_view::MaterializedViewDefinition { - vector_source: None, - source_table: "src".into(), - source_namespace: Vec::new(), - projections: vec![crate::materialized_view::ViewProjection { - output: "x".into(), - expression: "x".into(), - }], - filter: None, - group_by: Vec::new(), - limit: None, - lateral: None, - }; - let mut metadata = HashMap::new(); - metadata.insert( - crate::materialized_view::DEFINITION_META_KEY.to_string(), - crate::materialized_view::definition_to_metadata(&narrower).unwrap(), - ); - view.table() - .as_native() - .unwrap() - .replace_schema_metadata(metadata) - .await - .unwrap(); - - let err = view.refresh().execute().await.unwrap_err(); - assert!(matches!(err, Error::Schema { message } if message.contains("does not produce")),); - } - - /// An overlay replaces cell values without changing any file path; the - /// signature must see it, scoped to the columns the view reads like - /// data files are. - #[test] - fn test_fragment_signature_sees_overlays() { - use lance_file::version::ConcreteFileVersion; - use lance_table::format::DataFile; - use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; - - let base = Fragment::new(7); - let mut file = DataFile::new_unstarted("f0.lance", ConcreteFileVersion::V2_1); - file.fields = vec![0, 1].into(); - let mut with_file = base.clone(); - with_file.files.push(file.clone()); - - let overlay = |field: i32| { - let mut data_file = DataFile::new_unstarted("o0.lance", ConcreteFileVersion::V2_1); - data_file.fields = vec![field].into(); - DataOverlayFile { - data_file, - coverage: OverlayCoverage::PerField(Vec::new()), - committed_version: 9, - } - }; - let relevant: HashSet = [0].into_iter().collect(); - - let mut overlaid_relevant = with_file.clone(); - overlaid_relevant.overlays.push(overlay(0)); - assert_ne!( - fragment_signature(&with_file, &relevant), - fragment_signature(&overlaid_relevant, &relevant), - ); - - let mut overlaid_unrelated = with_file.clone(); - overlaid_unrelated.overlays.push(overlay(5)); - assert_eq!( - fragment_signature(&with_file, &relevant), - fragment_signature(&overlaid_unrelated, &relevant), - ); - } - - /// An output whose name needs quoting flows through as a projection - /// alias, never spliced into SQL text. - #[tokio::test] - async fn test_output_names_needing_quotes() { - let (conn, _) = db_with_source(vec![1, 2]).await; - let view = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .select([("double value", "x * 2")]) - .execute() - .await - .unwrap(); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.rows_written, 2); - assert_eq!(read(view.table(), "double value").await, vec![2, 4]); - } - - /// MemWAL tiers are visible to reads but not to the refresh scan, so LSM - /// state disqualifies every participant: the source at create, either - /// side at refresh, and the view can never accept a spec. Retained - /// un-compacted rows (the catch-up flag outlives unset) count as state. - #[tokio::test] - async fn lsm_state_disqualifies_source_and_view() { - use crate::table::LsmWriteSpec; - use arrow_array::RecordBatchIterator; - - let tmp_dir = tempfile::tempdir().unwrap(); - let conn = connect(tmp_dir.path().to_str().unwrap()) - .execute() - .await - .unwrap(); - // Hand-rolled: the LSM primary key must be non-nullable, which - // record_batch! cannot express. - let schema = Arc::new(ArrowSchema::new(vec![ - arrow_schema::Field::new("id", arrow_schema::DataType::Int64, false), - arrow_schema::Field::new("x", arrow_schema::DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(arrow_array::Int64Array::from(vec![1, 2])) as _, - Arc::new(Int32Array::from(vec![1, 2])) as _, - ], - ) - .unwrap(); - let table = conn - .create_table("src", batch.clone()) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - table.set_unenforced_primary_key(["id"]).await.unwrap(); - table - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap(); - - // An active-LSM source is refused at create. - let err = conn - .create_materialized_view("v", "src") - .with_no_data(true) - .execute() - .await - .unwrap_err(); - assert!(err.to_string().contains("un-compacted"), "{err}"); - - // Unset with nothing written clears the state: the view creates, - // and a spec can never be installed over it. - table.unset_lsm_write_spec().await.unwrap(); - let view = conn - .create_materialized_view("v", "src") + let view = connection + .create_materialized_view("copy", "src") .with_no_data(true) .execute() .await .unwrap(); - let err = view - .table() - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap_err(); - assert!(err.to_string().contains("materialized view"), "{err}"); - - // A source that acquires a spec after create fails refresh. - table - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err(); - assert!(err.to_string().contains("source table 'src'"), "{err}"); - - // Retained rows outlive unset: write through the WAL, unset, and - // refresh still refuses. - let mut merge = table.merge_insert(&["id"]); - merge - .when_matched_update_all(None) - .when_not_matched_insert_all() - .use_lsm(true); - merge - .execute(Box::new(RecordBatchIterator::new( - vec![Ok(batch.clone())], - batch.schema(), - ))) - .await - .unwrap(); - table.unset_lsm_write_spec().await.unwrap(); - let err = view.refresh().execute().await.unwrap_err(); - assert!(err.to_string().contains("source table 'src'"), "{err}"); - } - - /// A view with a computed column, declared over `people` and refreshed. - async fn refreshed_computed_view(conn: &Connection) -> MaterializedView { - use crate::materialized_view::tests::{computed_field, people, test_binding}; - let source = people(conn).await; - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[ - ("id".to_string(), "id".to_string()), - ("name".to_string(), "name".to_string()), - ]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns( - vec![(2, computed_field("emb", "fb_1", "name"))], - &[test_binding("fb_1", "name", "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - view - } - - async fn unfilled(view: &MaterializedView) -> usize { - view.table() - .count_rows(Some("emb IS NULL".to_string())) - .await - .unwrap() - } - - async fn append_people(conn: &Connection, ids: Vec, names: Vec<&str>) { - let batch = record_batch!(("id", Int32, ids), ("name", Utf8, names)).unwrap(); - conn.open_table("people") - .execute() - .await - .unwrap() - .add(batch) - .execute() - .await - .unwrap(); - } - - /// Commit the fill job's shape on the view: a column rewrite of - /// `fields`, touching no rows. The data is left as it is; what matters - /// here is how the next refresh classifies the commit. - async fn commit_column_rewrite(view: &MaterializedView, fields: &[&str]) { - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let dataset = native.dataset.get().await.unwrap().as_ref().clone(); - let fields_modified = fields - .iter() - .map(|name| dataset.schema().field(name).unwrap().id as u32) - .collect(); - let updated_fragments = dataset - .get_fragments() - .iter() - .map(|fragment| fragment.metadata().clone()) - .collect(); - let operation = Operation::Update { - removed_fragment_ids: Vec::new(), - updated_fragments, - new_fragments: Vec::new(), - fields_modified, - compacted_sstables: Vec::new(), - fields_for_preserving_frag_bitmap: Vec::new(), - update_mode: Some(UpdateMode::RewriteColumns), - inserted_rows_filter: None, - updated_fragment_offsets: None, - }; - let read_version = dataset.version().version; - CommitBuilder::new(WriteDestination::Dataset(Arc::new(dataset))) - .execute(Transaction::new(read_version, operation, None)) - .await - .unwrap(); - } - - /// Refresh never computes a computed column: every row it writes, on a - /// rebuild, an append and a rewrite, carries NULL there, and the - /// declaration survives all three. - #[tokio::test] - async fn test_computed_columns_are_written_null_and_kept() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - assert_eq!(unfilled(&view).await, 3); - - append_people(&conn, vec![4, 5], vec!["d", "e"]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(unfilled(&view).await, 5); - - conn.open_table("people") - .execute() - .await - .unwrap() - .update() - .column("name", "'z'") - .only_if("id = 1") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!(unfilled(&view).await, 5); - assert_eq!(read(view.table(), "id").await, vec![1, 2, 3, 4, 5]); - let schema = view.table().schema().await.unwrap(); - assert!( - crate::table::computed_columns::function_bindings(&schema) - .unwrap() - .iter() - .any(|b| b.binding_id() == "fb_1"), - "the binding envelope was lost" - ); - assert!( - computed_column_from_field(schema.field_with_name("emb").unwrap()).is_some(), - "the declaration was lost" - ); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// Field metadata on `field` only, the commit shape of the freshness - /// stamp `refresh_column` leaves after its fill. - async fn commit_field_metadata(view: &MaterializedView, field: &str, key: &str) { - let native = view.table().as_native().unwrap(); - native.dataset.reload().await.unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .update_field_metadata() - .update(field, [(key.to_string(), "{}".to_string())]) - .unwrap() - .await - .unwrap(); - } - - /// The stamp is metadata on the computed column and rewrites nothing - /// refresh certifies, so it is not drift; the same commit shape on a - /// projected column is, like any other write to it. - #[tokio::test] - async fn test_a_freshness_stamp_is_not_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_field_metadata( - &view, - "emb", - crate::table::computed_columns::SOURCE_SIGNATURE_META_KEY, - ) - .await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - - commit_field_metadata(&view, "id", "probe").await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - } - - /// The fill job's commit rewrites only computed columns. It is the one - /// commit on a view that is not drift: the next refresh carries on from - /// its watermark instead of rebuilding, which would null what the fill - /// just wrote. - #[tokio::test] - async fn test_a_computed_column_fill_is_not_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_column_rewrite(&view, &["emb"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - - commit_column_rewrite(&view, &["emb"]).await; - append_people(&conn, vec![4], vec!["d"]).await; - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Incremental); - assert_eq!(result.rows_written, 1); - assert_eq!(read(view.table(), "id").await, vec![1, 2, 3, 4]); - } - - /// A column rewrite that reaches a projected column is drift like any - /// other write: refresh certifies those columns and must recompute them. - #[tokio::test] - async fn test_a_rewrite_of_a_projected_column_is_drift() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - - commit_column_rewrite(&view, &["emb", "name"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - } - - /// The declaration contract is checked before any refresh mutation: a - /// missing binding envelope and a column that lost its declaration both - /// fail closed. - #[tokio::test] - async fn test_a_broken_declaration_is_refused_before_refresh() { - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .update_schema_metadata(vec![( - crate::table::computed_columns::FUNCTION_BINDINGS_META_KEY.to_string(), - None, - )]) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err().to_string(); - assert!(err.contains("references missing binding 'fb_1'"), "{err}"); - - let conn = connect("memory://").execute().await.unwrap(); - let view = refreshed_computed_view(&conn).await; - let native = view.table().as_native().unwrap(); - let mut dataset = native.dataset.get().await.unwrap().as_ref().clone(); - dataset - .replace_field_metadata(vec![( - dataset.schema().field("emb").unwrap().id as u32, - HashMap::new(), - )]) - .await - .unwrap(); - let err = view.refresh().execute().await.unwrap_err().to_string(); - assert!(err.contains("does not match binding 'fb_1'"), "{err}"); - } - - /// An input the view does not project is materialized on every refresh - /// path, before the provenance column, with the source's values. - #[tokio::test] - async fn test_internal_inputs_are_materialized_and_refreshed() { - use crate::materialized_view::tests::{computed_field, strict_people, test_binding}; - let conn = connect("memory://").execute().await.unwrap(); - let source = strict_people(&conn).await; - let mut prepared = crate::materialized_view::prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap(); - let input = prepared.input_column("name").unwrap(); - let view = prepared - .with_computed_columns( - vec![(1, computed_field("emb", "fb_1", &input))], - &[test_binding("fb_1", &input, "emb")], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let names: Vec = view + let result_batch = + RecordBatch::try_new(schema.clone(), source_batch.columns().to_vec()).unwrap(); + let planned = view .table() - .schema() - .await - .unwrap() - .fields() - .iter() - .map(|f| f.name().clone()) - .collect(); - assert_eq!(names, ["id", "emb", "__input_name", SOURCE_ROW_ID_COLUMN]); - - let unfilled_inputs = || async { - view.table() - .count_rows(Some("__input_name IS NULL".to_string())) - .await - .unwrap() - }; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - assert_eq!(view.table().count_rows(None).await.unwrap(), 3); - assert_eq!(unfilled_inputs().await, 0); - - let more = arrow_array::RecordBatch::try_new( - source.schema().await.unwrap(), - vec![ - Arc::new(Int32Array::from(vec![4])), - Arc::new(arrow_array::StringArray::from(vec!["d"])), - ], - ) - .unwrap(); - source.add(more).execute().await.unwrap(); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Incremental - ); - assert_eq!(unfilled_inputs().await, 0); - assert_eq!( - view.table() - .count_rows(Some("__input_name = 'd'".to_string())) - .await - .unwrap(), - 1 - ); - - source - .update() - .column("name", "'z'") - .only_if("id = 1") - .execute() - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - view.table() - .count_rows(Some("__input_name = 'z'".to_string())) - .await - .unwrap(), - 1 - ); - assert_eq!( - unfilled(&view).await, - 4, - "rewritten and new rows are unfilled" - ); - } - - /// A SQL declaration is filled by `refresh_column` on the view, which - /// commits a data replacement and then its freshness stamp; the next - /// refresh continues from its watermark and keeps what the fill wrote, - /// and only rows the view added since come back unfilled. - #[tokio::test] - async fn test_a_sql_fill_is_not_drift() { - use crate::materialized_view::tests::{people, sql_field}; - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[("id".to_string(), "id".to_string())]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns( - vec![( - 1, - sql_field("next", arrow_schema::DataType::Int32, "id + 1", r#"["id"]"#), - )], - &[], - ) - .unwrap() - .create("v") - .await - .unwrap(); - let filled = || async { - view.table() - .count_rows(Some("next = id + 1".to_string())) - .await - .unwrap() - }; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Rebuild - ); - assert_eq!( - view.table() - .refresh_column("next") - .await - .unwrap() - .rows_filled, - 3 - ); - assert_eq!(filled().await, 3); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!(filled().await, 3); - - append_people(&conn, vec![4], vec!["d"]).await; - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::Incremental - ); - assert_eq!(filled().await, 3); - assert_eq!( - view.table() - .refresh_column("next") - .await - .unwrap() - .rows_filled, - 1 - ); - assert_eq!(filled().await, 4); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - } - - /// A fill of a nested computed column writes its child fields; that is - /// still a fill, not drift. - #[tokio::test] - async fn test_a_nested_sql_fill_is_not_drift() { - use crate::materialized_view::tests::{people, sql_field}; - let conn = connect("memory://").execute().await.unwrap(); - let source = people(&conn).await; - let payload = sql_field( - "payload", - arrow_schema::DataType::Struct( - vec![arrow_schema::Field::new( - "value", - arrow_schema::DataType::Utf8, - true, - )] - .into(), - ), - "named_struct('value', name)", - r#"["name"]"#, - ); - let view = crate::materialized_view::prepare_declaration( - &source, - Some(&[("name".to_string(), "name".to_string())]), - None, - None, - ) - .await - .unwrap() - .with_computed_columns(vec![(1, payload)], &[]) - .unwrap() - .create("v") - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - view.table() - .refresh_column("payload") - .await - .unwrap() - .rows_filled, - 3 - ); - assert_eq!( - view.refresh().execute().await.unwrap().mode, - RefreshMode::NoOp - ); - assert_eq!( - view.table() - .count_rows(Some("payload.value = name".to_string())) - .await - .unwrap(), - 3 - ); - } - - async fn grouped_source(conn: &Connection) -> Table { - conn.create_table( - "src", - record_batch!(("k", Int32, [1, 1, 2, 3]), ("x", Int32, [10, 11, 20, 30])).unwrap(), - ) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap() - } - - pub(in crate::materialized_view) async fn declare( - source: &Table, - name: &str, - sql: &str, - ) -> Result { - crate::materialized_view::prepare_definition( - source, - MaterializedViewDefinition::from_sql(sql)?, - ) - .await? - .create(name) - .await - } - - /// Rows of `table` as `columns` rendered to text, sorted. - pub(in crate::materialized_view) async fn rows(table: &Table, columns: &[&str]) -> Vec { - let batches = table - .query() - .select(Select::columns(columns)) - .execute() - .await - .unwrap() - .try_collect::>() - .await - .unwrap(); - let mut rows = Vec::new(); - for batch in &batches { - let formatters: Vec<_> = batch - .columns() - .iter() - .map(|c| { - arrow_cast::display::ArrayFormatter::try_new(c.as_ref(), &Default::default()) - .unwrap() - }) - .collect(); - for row in 0..batch.num_rows() { - let cells: Vec = formatters - .iter() - .map(|f| f.value(row).to_string()) - .collect(); - rows.push(cells.join(" ")); - } - } - rows.sort(); - rows - } - - #[tokio::test] - async fn a_grouped_view_holds_one_row_per_group() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - let view = declare( - &source, - "groups", - "SELECT k, array_agg(x) AS xs, sum(x) AS total FROM src WHERE x < 30 GROUP BY k", - ) - .await - .unwrap(); - assert_eq!( - view.definition().to_sql(), - "SELECT k, array_agg(x) AS xs, sum(x) AS total FROM src WHERE x < 30 GROUP BY k" - ); - - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(result.rows_written, 2); - assert_eq!(rows(view.table(), &["k", "total"]).await, ["1 21", "2 20"]); - let unchanged = view.refresh().execute().await.unwrap(); - assert_eq!(unchanged.mode, RefreshMode::NoOp); - - // Any source change recomputes every group. - source - .add(record_batch!(("k", Int32, [2]), ("x", Int32, [5])).unwrap()) - .execute() - .await - .unwrap(); - let result = view.refresh().execute().await.unwrap(); - assert_eq!(result.mode, RefreshMode::Rebuild); - assert_eq!(rows(view.table(), &["k", "total"]).await, ["1 21", "2 25"]); - - // A view over the groups expands them again. - let expanded = declare( - view.table(), - "elements", - "SELECT k, e AS x FROM groups, UNNEST(xs) AS e", - ) - .await - .unwrap(); - expanded.refresh().execute().await.unwrap(); - assert_eq!( - rows(expanded.table(), &["k", "x"]).await, - ["1 10", "1 11", "2 20", "2 5"] - ); - } - - /// Provenance is one source row per group, so the view's row-id keys - /// stay unique. - #[tokio::test] - async fn a_grouped_view_records_one_source_row_per_group() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - let view = declare( - &source, - "groups", - "SELECT k, count(*) AS n FROM src GROUP BY k", - ) - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - let ids = rows(view.table(), &[SOURCE_ROW_ID_COLUMN]).await; - assert_eq!(ids.len(), 3); - assert_eq!(ids.iter().collect::>().len(), 3, "{ids:?}"); - } - - #[tokio::test] - async fn a_grouped_view_is_planned_at_declaration() { - let conn = connect("memory://").execute().await.unwrap(); - let source = grouped_source(&conn).await; - for sql in [ - // x is neither a key nor aggregated - "SELECT k, x FROM src GROUP BY k", - "SELECT k, random() AS r FROM src GROUP BY k", - "SELECT k, count(*) AS n FROM src GROUP BY missing", - "SELECT k, count(*) AS __source_row_id FROM src GROUP BY k", - "SELECT k, count(*) AS k FROM src GROUP BY k", - ] { - assert!(declare(&source, "v", sql).await.is_err(), "{sql}"); - } - let mut grouped = crate::materialized_view::prepare_definition( - &source, - MaterializedViewDefinition::from_sql("SELECT k, count(*) AS n FROM src GROUP BY k") - .unwrap(), - ) - .await - .unwrap(); - assert!(grouped.input_column("x").is_err()); - } - - /// Twenty nonzero vectors in two clusters far apart, ids 0-9 and 10-19. - /// Nonzero so every vector has a cosine direction. - pub(in crate::materialized_view) async fn clustered_source(conn: &Connection) -> Table { - use arrow_array::types::Float32Type; - use arrow_array::{FixedSizeListArray, Int32Array}; - - let vectors = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| { - let base = if i < 10 { 0.0 } else { 100.0 }; - Some(vec![Some(base + 1.0 + i as f32 * 0.1), Some(base)]) - }), - 2, - ); - let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as arrow_array::ArrayRef, - ), - ("vec", Arc::new(vectors) as arrow_array::ArrayRef), - ]) - .unwrap(); - conn.create_table("images", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap() - } - - #[tokio::test] - async fn ivf_partition_groups_by_the_index_on_the_column() { - use crate::index::vector::IvfFlatIndexBuilder; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - let unindexed = declare(&source, "buckets", BUCKETS).await.unwrap_err(); - assert!( - unindexed.to_string().contains("IVF vector index"), - "{unindexed}" - ); - - source - .create_index( - &["vec"], - Index::IvfFlat(IvfFlatIndexBuilder::default().num_partitions(2)), - ) - .execute() - .await - .unwrap(); - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["n", "lo", "hi"]).await, - ["10 0 9", "10 10 19"] - ); - let buckets = rows(view.table(), &["bucket"]).await; - assert_eq!(buckets, ["0", "1"]); - } - - #[tokio::test] - async fn ivf_partition_takes_a_column() { - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - for sql in [ - "SELECT ivf_partition(id) AS b, count(*) AS n FROM images GROUP BY ivf_partition(id)", - "SELECT ivf_partition(vec, vec) AS b, count(*) AS n FROM images \ - GROUP BY ivf_partition(vec, vec)", - ] { - assert!(declare(&source, "v", sql).await.is_err(), "{sql}"); - } - } - - /// A cosine index assigns by L2 over unit vectors, as lance's own index - /// path does; forwarding cosine to the assigner panics. - #[tokio::test] - async fn ivf_partition_supports_cosine_indices() { - use crate::index::vector::IvfFlatIndexBuilder; - - const BUCKETS: &str = "SELECT ivf_partition(vec) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(vec)"; - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - let view = declare(&source, "cosine_buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["n", "lo", "hi"]).await, - ["10 0 9", "10 10 19"] - ); - } - - /// Segments of one index trained separately carry different centroids; - /// grouping by one of them would bucket the other segments' rows wrongly. - #[tokio::test] - async fn ivf_partition_refuses_segments_with_different_models() { - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - // A second fragment far from the first, so each segment trains its own centroids. - let far = FixedSizeListArray::from_iter_primitive::( - (0..20).map(|i| Some(vec![Some(-500.0 - i as f32), Some(-500.0)])), - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(20..40)) as ArrayRef, - ), - ("vec", Arc::new(far) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - let dataset = source.as_native().unwrap().dataset.get().await.unwrap(); - let mut dataset = dataset.as_ref().clone(); - let params = VectorIndexParams::ivf_flat(2, lance_linalg::distance::DistanceType::L2); - let mut segments = Vec::new(); - for fragment in dataset.get_fragments() { - let segment = dataset - .create_index_builder(&["vec"], IndexType::Vector, ¶ms) - .name("shared".to_string()) - .fragments(vec![fragment.id() as u32]) - .execute_uncommitted() - .await - .unwrap(); - segments.push(segment); - } - dataset - .commit_existing_index_segments("shared", "vec", segments) - .await - .unwrap(); - assert_eq!(dataset.load_indices().await.unwrap().len(), 2); - - let source = conn.open_table("images").execute().await.unwrap(); - let error = declare( - &source, - "buckets", - "SELECT ivf_partition(vec) AS b, count(*) AS n FROM images GROUP BY ivf_partition(vec)", - ) - .await - .unwrap_err(); - assert!(error.to_string().contains("trained separately"), "{error}"); - } - - /// A vector the assigner cannot place has no bucket: it is grouped under - /// NULL, never under partition 0. - #[tokio::test] - async fn ivf_partition_leaves_an_unassignable_vector_unbucketed() { - use crate::index::vector::IvfFlatIndexBuilder; - use arrow_array::types::Float32Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - - let conn = connect("memory://").execute().await.unwrap(); - let source = clustered_source(&conn).await; - source - .create_index( - &["vec"], - Index::IvfFlat( - IvfFlatIndexBuilder::default() - .num_partitions(2) - .distance_type(crate::DistanceType::Cosine), - ), - ) - .execute() - .await - .unwrap(); - // The zero vector has no direction, so cosine cannot place it. - let zero = FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(0.0), Some(0.0)])], - 2, - ); - source - .add( - RecordBatch::try_from_iter(vec![ - ("id", Arc::new(Int32Array::from(vec![20])) as ArrayRef), - ("vec", Arc::new(zero) as ArrayRef), - ]) - .unwrap(), - ) - .execute() - .await - .unwrap(); - let view = declare( - &source, - "buckets", - "SELECT ivf_partition(vec) AS bucket, count(*) AS n, min(id) AS lo, max(id) AS hi \ - FROM images GROUP BY ivf_partition(vec)", - ) - .await - .unwrap(); - view.refresh().execute().await.unwrap(); - // ArrayFormatter renders NULL as the empty string; partition numbers - // are the index's own and not asserted. - let groups = rows(view.table(), &["bucket", "n", "lo", "hi"]).await; - let (unbucketed, bucketed): (Vec<_>, Vec<_>) = - groups.iter().partition(|row| row.starts_with(' ')); - assert_eq!(unbucketed, [" 1 20 20"], "{groups:?}"); - let mut clusters: Vec<&str> = bucketed - .iter() - .map(|row| row.split_once(' ').unwrap().1) - .collect(); - clusters.sort(); - assert_eq!(clusters, ["10 0 9", "10 10 19"], "{groups:?}"); - } - - /// Byte vectors group by a Hamming index, the phash case. The centroids - /// are given, as lance's own tests do: two-means over a handful of - /// hashes can converge with one partition empty. - #[tokio::test] - async fn ivf_partition_groups_byte_vectors_by_hamming() { - use arrow_array::types::UInt8Type; - use arrow_array::{ArrayRef, FixedSizeListArray, Int32Array}; - use lance::index::DatasetIndexExt; - use lance::index::vector::VectorIndexParams; - use lance_index::IndexType; - use lance_index::vector::ivf::IvfBuildParams; - - let hashes = FixedSizeListArray::from_iter_primitive::( - (0..20u8).map(|i| { - let base = if i < 10 { 0x00 } else { 0xff }; - Some(vec![ - Some(base ^ (i % 4)), - Some(base), - Some(base), - Some(base), - ]) - }), - 4, - ); - let batch = RecordBatch::try_from_iter(vec![ - ( - "id", - Arc::new(Int32Array::from_iter_values(0..20)) as ArrayRef, - ), - ("phash", Arc::new(hashes) as ArrayRef), - ]) - .unwrap(); - let conn = connect("memory://").execute().await.unwrap(); - let source = conn - .create_table("images", batch) - .write_options(crate::materialized_view::tests::stable_row_ids()) - .execute() - .await - .unwrap(); - let centroids = FixedSizeListArray::from_iter_primitive::( - [Some(vec![Some(0); 4]), Some(vec![Some(0xff); 4])], - 4, - ); - let params = VectorIndexParams::with_ivf_flat_params( - lance_linalg::distance::DistanceType::Hamming, - IvfBuildParams::try_with_centroids(2, Arc::new(centroids)).unwrap(), - ); - let mut dataset = source .as_native() .unwrap() .dataset @@ -4806,19 +607,63 @@ pub(crate) mod tests { .unwrap() .as_ref() .clone(); - dataset - .create_index(&["phash"], IndexType::Vector, None, ¶ms, true) + assert!(planned.get_fragments().is_empty()); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let loser_batch = result_batch.clone(); + let blocked = futures::stream::once(async move { + entered_tx.send(()).unwrap(); + release_rx.await.unwrap(); + Ok::<_, DataFusionError>(loser_batch) + }); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema.clone(), blocked)); + let loser_table = view.table().clone(); + let loser = + tokio::spawn( + async move { replace_materialized_view_fragments(&loser_table, stream).await }, + ); + entered_rx.await.unwrap(); + + let winner_stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter([Ok(result_batch)]), + )); + let planned = Arc::new(planned); + let staged = InsertBuilder::new(WriteDestination::Dataset(planned.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted_stream(winner_stream) .await .unwrap(); - let source = conn.open_table("images").execute().await.unwrap(); + let Operation::Append { fragments } = staged.operation else { + panic!("expected append staging"); + }; + let filter = refresh_conflict_filter(&planned).unwrap(); + CommitBuilder::new(WriteDestination::Dataset(planned.clone())) + .execute(Transaction::new( + planned.version().version, + Operation::Update { + removed_fragment_ids: Vec::new(), + updated_fragments: Vec::new(), + new_fragments: fragments, + fields_modified: Vec::new(), + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: Vec::new(), + update_mode: None, + inserted_rows_filter: Some(filter), + updated_fragment_offsets: None, + }, + None, + )) + .await + .unwrap(); + release_tx.send(()).unwrap(); - const BUCKETS: &str = "SELECT ivf_partition(phash) AS bucket, count(*) AS n, \ - min(id) AS lo, max(id) AS hi FROM images GROUP BY ivf_partition(phash)"; - let view = declare(&source, "buckets", BUCKETS).await.unwrap(); - view.refresh().execute().await.unwrap(); - assert_eq!( - rows(view.table(), &["bucket", "n", "lo", "hi"]).await, - ["0 10 0 9", "1 10 10 19"] - ); + loser.await.unwrap().unwrap_err(); + let current = connection.open_table("copy").execute().await.unwrap(); + assert_eq!(current.count_rows(None).await.unwrap(), 3); } } diff --git a/rust/lancedb/src/materialized_view/vector_dedup.rs b/rust/lancedb/src/materialized_view/vector_dedup.rs deleted file mode 100644 index dc2759433..000000000 --- a/rust/lancedb/src/materialized_view/vector_dedup.rs +++ /dev/null @@ -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) -> 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, - policy_version: u32, -} -impl VectorDedupPlan { - pub(super) fn dependencies(&self) -> u32 { - self.native.snapshot.units - } - fn identity(&self) -> Result { - 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, - pub(super) fragments: Vec, - pub(super) rows: u64, -} - -pub(super) async fn plan_refresh(view: &Table, pinned: Option) -> Result { - 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::>(); - 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)> { - 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::>() - != 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::() - .expect("native row id"); - let b = batch - .column(1) - .as_any() - .downcast_ref::() - .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::() - .expect("sorted id"); - let b = batch - .column(1) - .as_any() - .downcast_ref::() - .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> { - let row_index = get_row_id_index(source).await?; - let mut masks: BTreeMap = 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> { - 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::>>() - .ok_or_else(|| invalid("missing dedup dependency receipt")) -} - -async fn load_masks( - view: &Dataset, - source_fragment: u32, - dependencies: &[&WrittenDedupUnit], -) -> Result { - 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 { - 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, - expected: Option<&str>, -) -> Result { - 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::>(); - 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::>(), [1, high + 1]); - assert_eq!(representatives.iter().collect::>(), [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::>(); - 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)); - } -} diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 27606cdfe..4f1c3f56c 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -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 Database for RemoteDatabase { &self, request: CreateMaterializedViewRequest, ) -> Result { - 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 Database for RemoteDatabase { 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 Database for RemoteDatabase { 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 Database for RemoteDatabase { // 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 Database for RemoteDatabase { 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::(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 { + 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] diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 489158922..e8ddcf6ff 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -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( + &self, + statement: String, + namespace: Vec, + finish: F, + ) -> Job + where + T: Clone + Serialize + DeserializeOwned + Send + Sync + 'static, + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + { + let client = self.clone(); + let (query_tx, query_rx) = watch::channel(None::>); + 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 { let query = self .queries diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index d20ba495f..0d7efc9c4 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -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, transient_poll_failures: Arc, headers: Arc>>, + gate_refresh_release: Arc, 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::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::(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 { + 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(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 898681fa7..cd9e9ab64 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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 { enum TrackedJobResult { None, RefreshColumn, - MaterializedView, } #[async_trait] @@ -305,11 +305,6 @@ impl crate::job::JobHandle for FreshnessJob { }) }) }), - TrackedJobResult::MaterializedView => result.value().and_then(|value| { - serde_json::from_value::(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, interval: Option, @@ -496,6 +495,7 @@ pub struct RemoteTable { namespace: Vec, identifier: String, server_version: ServerVersion, + sql_client: Option, version: Arc>>, location: RwLock>, @@ -654,12 +654,13 @@ impl RemoteTable { Ok(job_id) } - pub fn new( + pub(super) fn new_with_sql_client( client: RestfulLanceDbClient, name: String, namespace: Vec, identifier: String, server_version: ServerVersion, + sql_client: Option, ) -> Self { Self { client, @@ -667,6 +668,7 @@ impl RemoteTable { 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 RemoteTable { 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 BaseTable for RemoteTable { &self.identifier } async fn materialized_view_info(&self) -> Result { - #[derive(Deserialize)] - struct Projection { - output_column: String, - expression: String, - } - #[derive(Deserialize)] struct DescribeMaterializedViewResponse { - source_table: String, - #[serde(default)] - source_namespace: Vec, - #[serde(default)] - projections: Vec, - #[serde(default)] - filter: Option, - #[serde(default)] - limit: Option, - /// The defining query, once the server describes a view by it; - /// takes precedence over the structured fields. - #[serde(default)] - query: Option, - #[serde(default)] - incarnation: Option, + query: String, } let request = self.client.post(&format!( @@ -2139,79 +2125,64 @@ impl BaseTable for RemoteTable { 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, - expected_incarnation: Option<&str>, ) -> Result> { 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> { 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] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index ea6e512fd..cf7fcc145 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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; - /// 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, - _expected_incarnation: Option<&str>, ) -> Result> { Err(Error::NotSupported { message: "remote materialized-view refresh is not supported on this table type".into(), diff --git a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs index 0d4f23d92..000addbe1 100644 --- a/rust/lancedb/src/table/datafusion/udtf/vector_query.rs +++ b/rust/lancedb/src/table/datafusion/udtf/vector_query.rs @@ -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::() + .expect("native row id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .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::() + .expect("sorted id"); + let b = batch + .column(1) + .as_any() + .downcast_ref::() + .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( diff --git a/rust/lancedb/tests/duplicate_pairs.rs b/rust/lancedb/tests/duplicate_pairs.rs index f7d7ed3a2..63f1c9f5a 100644 --- a/rust/lancedb/tests/duplicate_pairs.rs +++ b/rust/lancedb/tests/duplicate_pairs.rs @@ -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::>() - .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::>() - .await?; - let mut actual_ids = kept - .iter() - .flat_map(|b| { - b.column(0) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect::>(); - 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::>() - .await?; - let mut ids = Vec::new(); - for batch in batches { - let row_ids = batch - .column_by_name("id") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let images = batch - .column_by_name("image") - .unwrap() - .as_any() - .downcast_ref::() - .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(()) }