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.
This commit is contained in:
Jack Ye
2026-11-12 14:57:00 -08:00
committed by GitHub
parent f863ea5e08
commit 9302287bed
11 changed files with 536 additions and 87 deletions
+6 -2
View File
@@ -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
+46 -14
View File
@@ -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,
@@ -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
+12 -4
View File
@@ -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,
+6 -2
View File
@@ -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,
+163 -12
View File
@@ -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()
+44 -6
View File
@@ -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<String>,
query: Option<String>,
has_clauses: bool,
) -> PyResult<CreateMaterializedViewBuilder> {
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<LanceConnection>,
@@ -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<String>,
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<f64>,
query: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
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<String>,
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<f64>,
query: Option<String>,
) -> PyResult<Bound<'_, PyAny>> {
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);
}