From 9302287bed205bc017979681fa31b9b8e6ae450f Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 12 Nov 2026 14:57:00 -0800 Subject: [PATCH] fix: accept defining queries in Python materialized views (#4980) Python materialized-view creation now accepts a complete defining SQL query across sync, async, remote, and namespace connections, including creation jobs: ```python view = db.create_materialized_view( "events_by_kind", query="SELECT kind, count(*) AS n FROM events GROUP BY kind", ) ``` `query` is mutually exclusive with `source`, `select`, `where`, and `limit`. Existing source-based calls remain supported. Adds Rust `Connection::create_materialized_view_from_query`: embedded connections plan through `MaterializedViewDefinition::from_sql`, while remote connections submit the full query for server-side parsing. Includes regression coverage for grouping, UNNEST, query forwarding, creation jobs, and argument validation. --- docs/src/python/python.md | 14 ++ .../materialized-views.mdx | 28 +-- python/python/lancedb/_lancedb.pyi | 8 +- python/python/lancedb/db.py | 60 ++++-- python/python/lancedb/materialized_view.py | 10 + python/python/lancedb/namespace.py | 16 +- python/python/lancedb/remote/db.py | 8 +- .../python/tests/test_materialized_views.py | 175 +++++++++++++-- python/src/connection.rs | 50 ++++- rust/lancedb/src/materialized_view.rs | 199 +++++++++++++++--- rust/lancedb/src/remote/sql_test.rs | 55 +++++ 11 files changed, 536 insertions(+), 87 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 058f2529a..4ce325b2f 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -627,6 +627,20 @@ remain available, as do LanceDB Enterprise operations. The same feature requirement applies to LanceDB Embedded SQL computed-column declaration and refresh; ordinary eager `add_columns` transforms remain available. +Pass a complete defining SQL query with `query=` to use grouping, Functions +in `FROM`, or `UNNEST`: + +```python +view = db.create_materialized_view( + "events_by_kind", + query="SELECT kind, count(*) AS n FROM events GROUP BY kind", +) +``` + +`query` is mutually exclusive with `source`, `select`, `where`, and `limit`. +It is also accepted by asynchronous connections and `create_materialized_view_async`. +The backend SQL parser determines which query shapes are supported. + In LanceDB Enterprise, creation accepts `auto_refresh_enabled` and `auto_refresh_interval_seconds`: diff --git a/docs/web/feature-engineering/materialized-views.mdx b/docs/web/feature-engineering/materialized-views.mdx index 0f9494c9f..617eb7525 100644 --- a/docs/web/feature-engineering/materialized-views.mdx +++ b/docs/web/feature-engineering/materialized-views.mdx @@ -20,19 +20,19 @@ In SQL this is [`CREATE MATERIALIZED VIEW`](/sql/ddl/materialized-view#create). ## Create a view {#create} A materialized view is defined by a SQL query, so it is created with one. -From Python, run the statement on an [Enterprise connection](/basics/catalog#enterprise) -with [`execute_query`](/basics/sql#sql-from-python); the SQL is the same -either way: +From Python, pass the defining query to `create_materialized_view` on an +[Enterprise connection](/basics/catalog#enterprise): ```python Python icon="python" -db.execute_query(""" - CREATE MATERIALIZED VIEW frames_captioned AS - SELECT frame_id, video_id, ts, caption - FROM frames - WHERE selected = true -""").read_all() -view = db.open_materialized_view("frames_captioned") +view = db.create_materialized_view( + "frames_captioned", + query=""" + SELECT frame_id, video_id, ts, caption + FROM frames + WHERE selected = true + """, +) ``` ```sql SQL icon="database" @@ -43,9 +43,11 @@ WHERE selected = true ``` -`execute_query` returns once the result stream is readable, not once the -statement has finished, so `.read_all()` waits out the full creation before -the next line opens it. +`create_materialized_view` waits for creation and initial population, then +returns the view handle. Use `create_materialized_view_async` to submit a job +and wait for it separately. Both methods accept `query=` on synchronous and +asynchronous connections. `query` cannot be combined with `source`, `select`, +`where`, or `limit`; put those clauses in the query itself. The query defines the view: here, the `frames` table's `frame_id`, `video_id`, `ts` and `caption` columns, filtered to the frames selected for diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index ffdd68e5f..5897ddd0c 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -349,24 +349,28 @@ class Connection(object): async def create_materialized_view( self, name: str, - source: str, + source: Optional[str] = None, projections: Optional[List[Tuple[str, str]]] = None, filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, auto_refresh_enabled: Optional[bool] = None, auto_refresh_interval_seconds: Optional[float] = None, + *, + query: Optional[str] = None, ) -> Table: ... async def create_materialized_view_async( self, name: str, - source: str, + source: Optional[str] = None, projections: Optional[List[Tuple[str, str]]] = None, filter: Optional[str] = None, limit: Optional[int] = None, with_no_data: bool = False, auto_refresh_enabled: Optional[bool] = None, auto_refresh_interval_seconds: Optional[float] = None, + *, + query: Optional[str] = None, ) -> Job: ... def list_materialized_views( self, *, page_token: Optional[str] = None, page_limit: Optional[int] = None diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index b4e5c6182..7d690abe3 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -79,6 +79,7 @@ from .materialized_view import ( MaterializedViewSource, SelectArg, normalize_select, + _validate_create_materialized_view, ) from .secrets import ( EnvVarSecret, @@ -698,8 +699,9 @@ class DBConnection(EnforceOverrides): def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -722,8 +724,12 @@ class DBConnection(EnforceOverrides): ---------- name: str The name of the view. - source: str - The name of the source table. + source: str, optional + The name of the source table. Required when ``query`` is omitted. + query: str, optional + The complete defining SQL query, including any GROUP BY or FROM + Function. Mutually exclusive with ``source``, ``select``, ``where`` + and ``limit``. Supported query shapes depend on the backend. 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 @@ -745,6 +751,13 @@ class DBConnection(EnforceOverrides): Returns ------- MaterializedView + + Examples + -------- + >>> view = db.create_materialized_view( # doctest: +SKIP + ... "events_by_kind", + ... query="SELECT kind, count(*) AS n FROM events GROUP BY kind", + ... ) """ raise NotImplementedError( _append_doc_url( @@ -756,8 +769,9 @@ class DBConnection(EnforceOverrides): def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -2108,8 +2122,9 @@ class LanceDBConnection(DBConnection): def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -2144,6 +2159,7 @@ class LanceDBConnection(DBConnection): self._conn.create_materialized_view( name, source, + query=query, select=select, where=where, limit=limit, @@ -2158,8 +2174,9 @@ class LanceDBConnection(DBConnection): def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -2171,6 +2188,7 @@ class LanceDBConnection(DBConnection): self._conn.create_materialized_view_async( name, source, + query=query, select=select, where=where, limit=limit, @@ -3390,8 +3408,9 @@ class AsyncConnection(object): async def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -3407,9 +3426,15 @@ class AsyncConnection(object): See [DBConnection.create_materialized_view][lancedb.DBConnection.create_materialized_view]. """ - source, native_options = await self._materialized_view_source( - source, select, where, limit - ) + _validate_create_materialized_view(source, query, select, where, limit) + if source is not None: + source, native_options = await self._materialized_view_source( + source, select, where, limit + ) + else: + native_options = {} + if query is not None: + native_options["query"] = query inner = await self._inner.create_materialized_view( name, source, @@ -3426,8 +3451,9 @@ class AsyncConnection(object): async def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -3441,9 +3467,15 @@ class AsyncConnection(object): Enterprise, creation uses SQL; ``job_type`` is ``"query"`` and ``id`` identifies that query on this connection. """ - source, native_options = await self._materialized_view_source( - source, select, where, limit - ) + _validate_create_materialized_view(source, query, select, where, limit) + if source is not None: + source, native_options = await self._materialized_view_source( + source, select, where, limit + ) + else: + native_options = {} + if query is not None: + native_options["query"] = query inner = await self._inner.create_materialized_view_async( name, source, diff --git a/python/python/lancedb/materialized_view.py b/python/python/lancedb/materialized_view.py index 7b8c8c6a4..dc04f89f8 100644 --- a/python/python/lancedb/materialized_view.py +++ b/python/python/lancedb/materialized_view.py @@ -134,6 +134,16 @@ def _quote_identifier(name: str) -> str: return f"`{escaped}`" +def _validate_create_materialized_view(source, query, select, where, limit): + if query is not None: + if any(value is not None for value in (source, select, where, limit)): + raise ValueError( + "query cannot be combined with source, select, where or limit" + ) + elif source is None: + raise ValueError("either source or query must be provided") + + def normalize_select(select: SelectArg) -> Optional[List[Tuple[str, str]]]: """``select`` items may be a column name, an ``(alias, expression)`` pair, or a dict of the same. A bare name projects itself and is quoted, so any diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index aa9025641..72dc78c8b 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -696,8 +696,9 @@ class LanceNamespaceDBConnection(DBConnection): def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: "SelectArg" = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -715,6 +716,7 @@ class LanceNamespaceDBConnection(DBConnection): self._inner.create_materialized_view( name, source, + query=query, select=select, where=where, limit=limit, @@ -730,8 +732,9 @@ class LanceNamespaceDBConnection(DBConnection): def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: "SelectArg" = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -743,6 +746,7 @@ class LanceNamespaceDBConnection(DBConnection): self._inner.create_materialized_view_async( name, source, + query=query, select=select, where=where, limit=limit, @@ -1339,8 +1343,9 @@ class AsyncLanceNamespaceDBConnection: async def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: "SelectArg" = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -1352,6 +1357,7 @@ class AsyncLanceNamespaceDBConnection: view = await self._inner.create_materialized_view( name, source, + query=query, select=select, where=where, limit=limit, @@ -1366,8 +1372,9 @@ class AsyncLanceNamespaceDBConnection: async def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: "SelectArg" = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -1379,6 +1386,7 @@ class AsyncLanceNamespaceDBConnection: return await self._inner.create_materialized_view_async( name, source, + query=query, select=select, where=where, limit=limit, diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 911766674..db3ab554c 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -792,8 +792,9 @@ class RemoteDBConnection(DBConnection): def create_materialized_view( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -807,6 +808,7 @@ class RemoteDBConnection(DBConnection): self._conn.create_materialized_view( name, source, + query=query, select=select, where=where, limit=limit, @@ -828,8 +830,9 @@ class RemoteDBConnection(DBConnection): def create_materialized_view_async( self, name: str, - source: MaterializedViewSource, + source: Optional[MaterializedViewSource] = None, *, + query: Optional[str] = None, select: SelectArg = None, where: Optional[str] = None, limit: Optional[int] = None, @@ -841,6 +844,7 @@ class RemoteDBConnection(DBConnection): self._conn.create_materialized_view_async( name, source, + query=query, select=select, where=where, limit=limit, diff --git a/python/python/tests/test_materialized_views.py b/python/python/tests/test_materialized_views.py index bf964102b..1ea0c1cb9 100644 --- a/python/python/tests/test_materialized_views.py +++ b/python/python/tests/test_materialized_views.py @@ -10,6 +10,7 @@ import lancedb import pytest import pyarrow as pa from lancedb.db import AsyncConnection +from lancedb.exceptions import JobFailedError from lancedb.materialized_view import MaterializedViewDefinition from lancedb.remote.db import RemoteDBConnection from utils import start_server_thread @@ -175,7 +176,8 @@ async def test_remote_refresh_reports_unknown_metrics( assert not any("count_rows" in path for path in requests) -def test_sync_remote_create_uses_public_async_connection(): +@pytest.mark.parametrize("query", [None, "SELECT name FROM people WHERE age >= 18"]) +def test_sync_remote_create_uses_public_async_connection(query): calls = [] class StubAsyncTable: @@ -190,6 +192,7 @@ def test_sync_remote_create_uses_public_async_connection(): name, source, *, + query=None, select=None, where=None, limit=None, @@ -201,6 +204,7 @@ def test_sync_remote_create_uses_public_async_connection(): ( name, source, + query, select, where, limit, @@ -219,18 +223,24 @@ def test_sync_remote_create_uses_public_async_connection(): db.db_name = "example" db.serialize = lambda: "{}" - view = db.create_materialized_view( - "adults", - "people", - select=["name"], - where="age >= 18", - limit=10, - with_no_data=True, - auto_refresh_enabled=True, - auto_refresh_interval_seconds=0.5, - ) + if query is None: + view = db.create_materialized_view( + "adults", + "people", + select=["name"], + where="age >= 18", + limit=10, + with_no_data=True, + auto_refresh_enabled=True, + auto_refresh_interval_seconds=0.5, + ) + assert calls == [ + ("adults", "people", None, ["name"], "age >= 18", 10, True, True, 0.5) + ] + else: + view = db.create_materialized_view("adults", query=query) + assert calls == [("adults", None, query, None, None, None, False, None, None)] assert view.name == "adults" - assert calls == [("adults", "people", ["name"], "age >= 18", 10, True, True, 0.5)] db.drop_materialized_view("adults", namespace_path=["analytics"]) assert calls[-1] == ("drop", "adults", ["analytics"]) @@ -790,3 +800,144 @@ def test_embedded_explicitly_disabled_refresh_is_supported(tmp_path): db = make_db(tmp_path) view = db.create_materialized_view("mv", "people", auto_refresh_enabled=False) assert view.table.count_rows() == 3 + + +@pytest.mark.experimental_functions +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("namespace", [False, True]) +@pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize( + "query, expected", + [ + ( + "SELECT kind, count(*) AS n FROM events GROUP BY kind", + [{"kind": "a", "n": 2}, {"kind": "b", "n": 1}], + ), + ( + "SELECT id, tag FROM events, UNNEST(tags) AS tag", + [{"id": 1, "tag": "x"}, {"id": 1, "tag": "y"}, {"id": 3, "tag": "z"}], + ), + ], +) +async def test_query_defined_materialized_view( + tmp_path, sync, namespace, background, query, expected +): + if namespace: + connect = lancedb.connect_namespace if sync else lancedb.connect_namespace_async + db = connect("dir", {"root": str(tmp_path)}, storage_options=STABLE_ROW_IDS) + elif sync: + db = lancedb.connect(tmp_path, storage_options=STABLE_ROW_IDS) + else: + db = await lancedb.connect_async(tmp_path, storage_options=STABLE_ROW_IDS) + rows = [ + {"id": 1, "kind": "a", "tags": ["x", "y"]}, + {"id": 2, "kind": "a", "tags": []}, + {"id": 3, "kind": "b", "tags": ["z"]}, + ] + if sync: + db.create_table("events", rows, storage_options=STABLE_ROW_IDS) + if background: + db.create_materialized_view_async("result", query=query).wait() + view = db.open_materialized_view("result") + else: + view = db.create_materialized_view("result", query=query) + actual = view.table.to_arrow().to_pylist() + assert view.definition.query == query + view.refresh() + refreshed = view.table.to_arrow().to_pylist() + else: + await db.create_table("events", rows, storage_options=STABLE_ROW_IDS) + if background: + job = await db.create_materialized_view_async("result", query=query) + await job.wait() + view = await db.open_materialized_view("result") + else: + view = await db.create_materialized_view("result", query=query) + actual = (await view.table.to_arrow()).to_pylist() + assert (await view.definition()).query == query + await view.refresh() + refreshed = (await view.table.to_arrow()).to_pylist() + assert sorted(actual, key=str) == sorted(expected, key=str) + assert sorted(refreshed, key=str) == sorted(expected, key=str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("background", [False, True]) +@pytest.mark.parametrize( + "clauses", + [{"source": "people"}, {"select": []}, {"select": {}}, {"where": ""}, {"limit": 0}], +) +async def test_query_rejects_source_clauses_before_submission( + sync, background, clauses +): + connect = lancedb.connect if sync else lancedb.connect_async + db = connect("db://dev", api_key="fake", host_override="http://127.0.0.1:1") + if not sync: + db = await db + create = ( + db.create_materialized_view_async if background else db.create_materialized_view + ) + with pytest.raises(ValueError, match="query cannot be combined"): + if sync: + create("invalid", query="SELECT * FROM people", **clauses) + else: + await create("invalid", query="SELECT * FROM people", **clauses) + with pytest.raises(ValueError, match="either source or query"): + if sync: + create("missing") + else: + await create("missing") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize( + "query", + [ + "SELECT kind, count(*) AS n FROM events GROUP BY kind", + "SELECT id, c.text FROM docs, chunk(body) AS c", + "SELECT id, tag FROM events, UNNEST(tags) AS tag", + "SELECT a.id FROM a JOIN b ON a.id = b.id", + ], +) +async def test_remote_create_forwards_defining_query(sync, query): + native = CapturingCreateConnection() + db = AsyncConnection(native) + if sync: + remote = RemoteDBConnection.__new__(RemoteDBConnection) + remote._conn = db + job = remote.create_materialized_view_async("result", query=query) + job.wait() + else: + job = await db.create_materialized_view_async("result", query=query) + await job.wait() + assert native.calls[0][0:2] == ("result", None) + assert native.calls[0][2]["query"] == query + + +@pytest.mark.experimental_functions +@pytest.mark.parametrize("background", [False, True]) +def test_query_with_no_data_and_invalid_query(tmp_path, background): + db = make_db(tmp_path) + query = "SELECT name FROM people WHERE age >= 18" + if background: + db.create_materialized_view_async( + "adults", query=query, with_no_data=True + ).wait() + view = db.open_materialized_view("adults") + else: + view = db.create_materialized_view("adults", query=query, with_no_data=True) + assert view.table.count_rows() == 0 + view.refresh() + assert view.table.count_rows() == 2 + create = ( + db.create_materialized_view_async if background else db.create_materialized_view + ) + error = JobFailedError if background else ValueError + with pytest.raises(error, match="unsupported clause"): + result = create("invalid", query="SELECT * FROM people ORDER BY age") + if background: + result.wait() + assert "invalid" not in db.table_names() diff --git a/python/src/connection.rs b/python/src/connection.rs index 2bf01a6f7..a32c00be1 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -26,6 +26,7 @@ use lancedb::{ database::{CreateTableMode, Database, ReadConsistency}, function::{GpuClass as LanceGpuClass, RefreshColumnResult}, graph::PropertyGraphDefinition, + materialized_view::CreateMaterializedViewBuilder, }; use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, @@ -34,6 +35,27 @@ use pyo3::{ types::{PyAnyMethods, PyDict, PyDictMethods, PyList}, }; +fn materialized_view_builder( + connection: &LanceConnection, + name: String, + source: Option, + query: Option, + has_clauses: bool, +) -> PyResult { + match (source, query) { + (None, Some(query)) if !has_clauses => { + Ok(connection.create_materialized_view_from_query(name, query)) + } + (Some(source), None) => Ok(connection.create_materialized_view(name, source)), + (None, None) => Err(PyValueError::new_err( + "either source or query must be provided", + )), + _ => Err(PyValueError::new_err( + "query cannot be combined with source, select, where or limit", + )), + } +} + #[pyclass] pub struct Connection { inner: Option, @@ -483,21 +505,29 @@ impl Connection { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] + #[pyo3(signature = (name, source=None, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None, *, query=None))] pub fn create_materialized_view( self_: PyRef<'_, Self>, name: String, - source: String, + source: Option, projections: Option>, filter: Option, limit: Option, with_no_data: bool, auto_refresh_enabled: Option, auto_refresh_interval_seconds: Option, + query: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); + let builder = materialized_view_builder( + &inner, + name, + source, + query, + projections.is_some() || filter.is_some() || limit.is_some(), + )?; future_into_py(self_.py(), async move { - let mut builder = inner.create_materialized_view(name, source); + let mut builder = builder; if let Some(enabled) = auto_refresh_enabled { builder = builder.auto_refresh_enabled(enabled); } @@ -520,21 +550,29 @@ impl Connection { } #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (name, source, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))] + #[pyo3(signature = (name, source=None, projections=None, filter=None, limit=None, with_no_data=false, auto_refresh_enabled=None, auto_refresh_interval_seconds=None, *, query=None))] pub fn create_materialized_view_async( self_: PyRef<'_, Self>, name: String, - source: String, + source: Option, projections: Option>, filter: Option, limit: Option, with_no_data: bool, auto_refresh_enabled: Option, auto_refresh_interval_seconds: Option, + query: Option, ) -> PyResult> { let inner = self_.get_inner()?.clone(); + let builder = materialized_view_builder( + &inner, + name, + source, + query, + projections.is_some() || filter.is_some() || limit.is_some(), + )?; future_into_py(self_.py(), async move { - let mut builder = inner.create_materialized_view(name, source); + let mut builder = builder; if let Some(enabled) = auto_refresh_enabled { builder = builder.auto_refresh_enabled(enabled); } diff --git a/rust/lancedb/src/materialized_view.rs b/rust/lancedb/src/materialized_view.rs index bd7fbb942..9920cd893 100644 --- a/rust/lancedb/src/materialized_view.rs +++ b/rust/lancedb/src/materialized_view.rs @@ -1574,14 +1574,15 @@ async fn prepare_with( } /// Builds a materialized view. Created by -/// [`Connection::create_materialized_view`]. +/// [`Connection::create_materialized_view`] or +/// [`Connection::create_materialized_view_from_query`]. pub struct CreateMaterializedViewBuilder { connection: Connection, name: String, namespace: Vec, - source: String, + input: MaterializedViewInput, source_namespace: Vec, - projections: Vec<(String, String)>, + projections: Option>, filter: Option, limit: Option, with_no_data: bool, @@ -1589,6 +1590,11 @@ pub struct CreateMaterializedViewBuilder { auto_refresh_interval_seconds: Option, } +enum MaterializedViewInput { + Source(String), + Query(String), +} + fn source_sql_identifier( namespace_path: &[String], source: &str, @@ -1609,14 +1615,14 @@ fn source_sql_identifier( } impl CreateMaterializedViewBuilder { - pub(crate) fn new(connection: Connection, name: String, source: String) -> Self { + fn new(connection: Connection, name: String, input: MaterializedViewInput) -> Self { Self { connection, name, namespace: Vec::new(), - source, + input, source_namespace: Vec::new(), - projections: Vec::new(), + projections: None, filter: None, limit: None, with_no_data: false, @@ -1644,10 +1650,12 @@ impl CreateMaterializedViewBuilder { mut self, columns: impl IntoIterator, impl Into)>, ) -> Self { - self.projections = columns - .into_iter() - .map(|(output, expression)| (output.into(), expression.into())) - .collect(); + self.projections = Some( + columns + .into_iter() + .map(|(output, expression)| (output.into(), expression.into())) + .collect(), + ); self } @@ -1697,22 +1705,38 @@ impl CreateMaterializedViewBuilder { } fn query(&self) -> Result { + if let MaterializedViewInput::Query(query) = &self.input { + if !self.source_namespace.is_empty() + || self.projections.is_some() + || self.filter.is_some() + || self.limit.is_some() + { + return Err(Error::InvalidInput { + message: "a defining query cannot be combined with source_namespace, select, only_if or limit".into(), + }); + } + return Ok(query.clone()); + } + let MaterializedViewInput::Source(source) = &self.input else { + unreachable!() + }; fn quote(name: &str) -> String { format!("\"{}\"", name.replace('"', "\"\"")) } - let projection = if self.projections.is_empty() { - "*".to_string() - } else { - self.projections - .iter() - .map(|(output, expression)| format!("{expression} AS {}", quote(output))) - .collect::>() - .join(", ") - }; + let projection = + if let Some(projections) = self.projections.as_ref().filter(|p| !p.is_empty()) { + projections + .iter() + .map(|(output, expression)| format!("{expression} AS {}", quote(output))) + .collect::>() + .join(", ") + } else { + "*".to_string() + }; let source = source_sql_identifier( &self.source_namespace, - &self.source, + source, self.connection.uri().starts_with("db://"), ); let mut query = format!("SELECT {projection} FROM {source}"); @@ -1775,19 +1799,33 @@ impl CreateMaterializedViewBuilder { .to_string(), }); } - let source = self - .connection - .open_table(&self.source) - .namespace(self.source_namespace.clone()) - .execute() - .await?; - let prepared = prepare_declaration( - &source, - (!self.projections.is_empty()).then_some(self.projections.as_slice()), - self.filter.as_deref(), - self.limit, - ) - .await?; + let prepared = match &self.input { + MaterializedViewInput::Source(name) => { + let source = self + .connection + .open_table(name) + .namespace(self.source_namespace.clone()) + .execute() + .await?; + prepare_declaration( + &source, + self.projections.as_deref().filter(|p| !p.is_empty()), + self.filter.as_deref(), + self.limit, + ) + .await? + } + MaterializedViewInput::Query(_) => { + let definition = MaterializedViewDefinition::from_sql(&self.query()?)?; + let source = self + .connection + .open_table(&definition.source_table) + .namespace(definition.source_namespace.clone()) + .execute() + .await?; + prepare_definition(&source, definition).await? + } + }; let view = prepared.create_in(&self.namespace, &self.name).await?; if !self.with_no_data { view.refresh().execute().await?; @@ -2067,7 +2105,41 @@ impl Connection { name: impl Into, source: impl Into, ) -> CreateMaterializedViewBuilder { - CreateMaterializedViewBuilder::new(self.clone(), name.into(), source.into()) + CreateMaterializedViewBuilder::new( + self.clone(), + name.into(), + MaterializedViewInput::Source(source.into()), + ) + } + + /// Define a materialized view using a complete SQL query. + /// + /// Embedded connections parse the query with + /// [`MaterializedViewDefinition::from_sql`]. Remote connections send the + /// query to the server, which determines the supported SQL syntax. + /// Put projections, predicates and limits in the query rather than using + /// the builder's source-specific methods. + /// + /// ```no_run + /// # #![recursion_limit = "256"] + /// # async fn create(db: &lancedb::Connection) -> lancedb::Result<()> { + /// let view = db.create_materialized_view_from_query( + /// "events_by_kind", + /// "SELECT kind, count(*) AS n FROM events GROUP BY kind", + /// ).execute().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn create_materialized_view_from_query( + &self, + name: impl Into, + query: impl Into, + ) -> CreateMaterializedViewBuilder { + CreateMaterializedViewBuilder::new( + self.clone(), + name.into(), + MaterializedViewInput::Query(query.into()), + ) } /// Open the materialized view named `name`. @@ -2229,6 +2301,65 @@ mod tests { ); } + #[tokio::test] + async fn query_builder_rejects_source_clauses() { + let conn = connect("memory://").execute().await.unwrap(); + let builder = || conn.create_materialized_view_from_query("mv", "SELECT * FROM src"); + for invalid in [ + builder().select([("x", "x")]), + builder().select(Vec::<(String, String)>::new()), + builder().only_if("x > 0"), + builder().limit(0), + builder().source_namespace(vec!["ns".into()]), + ] { + assert!(matches!(invalid.query(), Err(Error::InvalidInput { .. }))); + } + } + + #[tokio::test] + #[cfg_attr( + not(feature = "experimental-functions"), + ignore = "requires experimental-functions" + )] + async fn create_from_query_supports_grouping_and_background_population() { + let conn = connect("memory://").execute().await.unwrap(); + conn.create_table( + "src", + record_batch!(("kind", Utf8, ["a", "a", "b"])).unwrap(), + ) + .execute() + .await + .unwrap(); + let query = "SELECT kind, count(*) AS n FROM src GROUP BY kind"; + let view = conn + .create_materialized_view_from_query("groups", query) + .with_no_data(true) + .execute() + .await + .unwrap(); + assert_eq!(view.definition_sql(), query); + assert_eq!(view.table().count_rows(None).await.unwrap(), 0); + view.refresh().execute().await.unwrap(); + assert_eq!(view.table().count_rows(None).await.unwrap(), 2); + conn.create_materialized_view_from_query("background", query) + .execute_async() + .await + .unwrap() + .wait() + .await + .unwrap(); + assert_eq!( + conn.open_materialized_view("background") + .await + .unwrap() + .table() + .count_rows(None) + .await + .unwrap(), + 2 + ); + } + #[test] fn sql_definition_defaults_share_one_metadata_value() { let metadata = HashMap::from([( diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index eb994f7dd..d25c72830 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -1137,6 +1137,61 @@ async fn catalog_connections_use_explicit_sql_endpoint_and_database_scope() { server.await.unwrap().unwrap(); } +#[tokio::test] +async fn remote_query_builder_sends_the_complete_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 { + adaptive: true, + ..Default::default() + }; + let statements = service.statements.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 connection = crate::connect("db://analytics") + .region("us-east-1") + .api_key("test-key") + .host_override("http://127.0.0.1:1") + .sql_host_override(&format!("grpc://{address}")) + .execute() + .await + .unwrap(); + for query in [ + "SELECT kind, count(*) AS n FROM events GROUP BY kind", + "SELECT id, c.text FROM docs, chunk(body) AS c", + "SELECT id, tag FROM events, UNNEST(tags) AS tag", + "SELECT a.id FROM a JOIN b ON a.id = b.id", + ] { + connection + .create_materialized_view_from_query("result", query) + .auto_refresh_enabled(false) + .execute_async() + .await + .unwrap() + .wait() + .await + .unwrap(); + assert_eq!( + statements.lock().unwrap().last().unwrap(), + &format!( + "CREATE MATERIALIZED VIEW \"result\" WITH (auto_refresh_enabled = false) AS {query}" + ) + ); + } + shutdown_tx.send(()).unwrap(); + server.await.unwrap().unwrap(); +} + #[tokio::test] async fn remote_create_uses_sql_and_preserves_options_and_namespace() { for defer in [false, true] {