rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-13 14:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: accept defining queries in Python materialized views (#4980)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/python/python.md
+14
View File
@@ -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`:
docs/web/feature-engineering/materialized-views.mdx
+15 -13
View File
@@ -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):
<CodeGroup>
```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
```
</CodeGroup>
`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
python/python/lancedb/_lancedb.pyi
+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
python/python/lancedb/db.py
+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,
python/python/lancedb/materialized_view.py
+10
View File
@@ -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
python/python/lancedb/namespace.py
+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,
python/python/lancedb/remote/db.py
+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,
python/python/tests/test_materialized_views.py
+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()
python/src/connection.rs
+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);
}
rust/lancedb/src/materialized_view.rs
+165 -34
View File
@@ -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<String>,
source: String,
input: MaterializedViewInput,
source_namespace: Vec<String>,
projections: Vec<(String, String)>,
projections: Option<Vec<(String, String)>>,
filter: Option<String>,
limit: Option<u64>,
with_no_data: bool,
@@ -1589,6 +1590,11 @@ pub struct CreateMaterializedViewBuilder {
auto_refresh_interval_seconds: Option<f64>,
}
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<Item = (impl Into<String>, impl Into<String>)>,
) -> 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<String> {
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::<Vec<_>>()
.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::<Vec<_>>()
.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<String>,
source: impl Into<String>,
) -> 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<String>,
query: impl Into<String>,
) -> 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([(
rust/lancedb/src/remote/sql_test.rs
+55
View File
@@ -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] {