mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat: add materialized view refresh resource budgets (#4902)
Expose aggregate CPU, memory, and GPU budgets plus GPU-class placement
on materialized-view refreshes in Python and Rust. Forward options
through SQL and reject invalid or unsupported Embedded requests instead
of silently dropping constraints.
```python
job = view.refresh_async(
cpu_millis=8000,
memory_bytes=16 * 1024**3,
gpu_count=2,
gpu_class="h100",
)
job.wait()
```
```sql
REFRESH MATERIALIZED VIEW curated_training_set
WITH (cpu_millis = 8000, memory_bytes = 17179869184,
gpu_count = 2, gpu_class = 'h100');
```
The same options apply to blocking refreshes, async Python connections,
and the Rust refresh builder. CPU/memory budgets require all three
resource fields; GPU-only caps and class-only placement are also
supported. These are execution-wide caps and require a supporting
Enterprise server with an isolated execution backend. A cap does not
increase refresh parallelism; current MV refreshes use at most one GPU.
Updates API documentation and SQL reference, with Python sync/async and
Rust remote-submission regression coverage. Formatting and Python lint
pass; local builds and tests are intentionally skipped.
The Linux CI build and example steps explicitly select the LanceDB
package, matching the feature-test step. This avoids rebuilding and
retaining a second dependency graph after a workspace-wide build, which
exhausted runner disk space on both this branch and main.
Related: [ENT-3386](https://linear.app/lancedb/issue/ENT-3386)
This commit is contained in:
@@ -240,18 +240,11 @@ jobs:
|
||||
sudo mkswap "$swapfile"
|
||||
sudo swapon "$swapfile"
|
||||
- name: Build
|
||||
run: cargo build --profile ci --all-features --tests --locked --examples
|
||||
- name: Reclaim completed workspace build artifacts
|
||||
run: |
|
||||
# Package-level tests resolve a different feature graph from the workspace
|
||||
# build. Reclaim its lancedb outputs before linking the test graph.
|
||||
df -h .
|
||||
cargo clean --profile ci -p lancedb
|
||||
df -h .
|
||||
run: cargo build --profile ci -p lancedb --all-features --tests --locked --examples
|
||||
- name: Run feature tests
|
||||
run: CARGO_ARGS="--profile ci" make -C ./lancedb feature-tests
|
||||
- name: Run examples
|
||||
run: cargo run --profile ci --all-features --example simple --locked
|
||||
run: cargo run --profile ci -p lancedb --all-features --example simple --locked
|
||||
- name: Run blob v2 docs example
|
||||
# Python and TypeScript blob docs tests install preview packages.
|
||||
# Those builds are not guaranteed to remain available for more than
|
||||
|
||||
@@ -59,6 +59,65 @@ Use `view.table.count_rows()` for the view's row count at any time; on
|
||||
LanceDB Embedded it always matches `rows_written` from the most recent
|
||||
refresh.
|
||||
|
||||
{/*
|
||||
Draft: publish this section after the Python/Rust resource-budget SDK support
|
||||
is released. TypeScript parity is tracked separately in issue #4903.
|
||||
|
||||
## Resource budgets {#resources}
|
||||
|
||||
For training-data curation, bound the resources used to recompute a view:
|
||||
|
||||
<CodeGroup>
|
||||
```python Python icon="python"
|
||||
job = view.refresh_async(
|
||||
cpu_millis=8000,
|
||||
memory_bytes=16 * 1024**3,
|
||||
gpu_count=2,
|
||||
gpu_class="h100",
|
||||
)
|
||||
job.wait()
|
||||
```
|
||||
|
||||
```rust Rust icon="rust"
|
||||
use lancedb::materialized_view::MaterializedView;
|
||||
|
||||
async fn refresh(view: &MaterializedView) -> lancedb::Result<()> {
|
||||
let job = view
|
||||
.refresh()
|
||||
.cpu_millis(8000)
|
||||
.memory_bytes(16_u64 << 30)
|
||||
.gpu_count(2)
|
||||
.gpu_class("h100")
|
||||
.execute_async()
|
||||
.await?;
|
||||
job.wait().await?;
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
The same options are available on `refresh`, asynchronous Python connections,
|
||||
and the Rust refresh builder (`cpu_millis`, `memory_bytes`, `gpu_count`, and
|
||||
`gpu_class`). CPU is in millicores (8000 means eight cores), memory is bytes,
|
||||
and GPU count is an aggregate concurrent cap. CPU or memory requires all three
|
||||
budget fields. Use zero GPUs for a CPU-only view, or specify a GPU cap alone.
|
||||
|
||||
The cap covers the refresh execution and its Function processes together.
|
||||
It does not allocate that budget separately to each Function. A GPU class
|
||||
selects placement for this execution; discover configured classes with
|
||||
`db.list_gpu_classes()`. A cap does not promise more parallelism: current MV
|
||||
refreshes use at most one GPU. Omitted settings keep existing defaults.
|
||||
|
||||
These controls require an Enterprise server with an isolated backend that can
|
||||
enforce them; Embedded and unsupported execution backends reject them. Earlier
|
||||
Enterprise servers reject the new SQL options. See the
|
||||
[SQL resource options](/sql/dml/refresh-materialized-view#resources).
|
||||
|
||||
TypeScript does not expose these options yet; SDK parity is tracked in
|
||||
[issue #4903](https://github.com/lancedb/lancedb/issues/4903).
|
||||
|
||||
*/}
|
||||
|
||||
## Automatic refresh {#auto-refresh}
|
||||
|
||||
A materialized view can refresh itself on a schedule instead of waiting
|
||||
|
||||
@@ -14,7 +14,8 @@ A LanceDB extension with no DataFusion equivalent.
|
||||
<CodeGroup>
|
||||
```sql SQL icon="database"
|
||||
REFRESH MATERIALIZED VIEW <name>
|
||||
[WITH (full = <boolean>)]
|
||||
[WITH (full = <boolean>, cpu_millis = <integer>, memory_bytes = <integer>,
|
||||
gpu_count = <integer>, gpu_class = '<class>')]
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
@@ -36,6 +37,39 @@ regardless of what a local SDK's `refresh()` method accepts -- do not
|
||||
assume the two option sets match.
|
||||
</Note>
|
||||
|
||||
## Resource budgets {#resources}
|
||||
|
||||
Set budgets for one refresh execution:
|
||||
|
||||
```sql
|
||||
REFRESH MATERIALIZED VIEW curated_training_set WITH (
|
||||
cpu_millis = 8000,
|
||||
memory_bytes = 17179869184,
|
||||
gpu_count = 2,
|
||||
gpu_class = 'h100'
|
||||
);
|
||||
```
|
||||
|
||||
`cpu_millis` is CPU capacity in millicores (1000 = one core), not CPU time.
|
||||
`memory_bytes` caps memory in bytes. `gpu_count` caps concurrent GPUs for the
|
||||
refresh. These are aggregate caps across its SQL execution and Function
|
||||
processes; they are not multiplied by the number of Functions in the query.
|
||||
|
||||
CPU or memory requires all three budget fields. CPU and memory must be positive;
|
||||
`gpu_count = 0` is allowed for CPU-only views. A GPU budget alone is also accepted.
|
||||
A cap does not reserve all of its GPUs or increase refresh parallelism. Current
|
||||
MV refreshes use at most one GPU, shared by their GPU Function calls.
|
||||
|
||||
`gpu_class` selects an operator-configured class for this execution, leaving the
|
||||
Function registration and saved MV placement unchanged. Use `SHOW GPU CLASSES`
|
||||
to discover configured classes. Omitted budgets and class preserve defaults.
|
||||
Options can be combined with `full = true`.
|
||||
|
||||
This requires an Enterprise server and an isolated execution backend that can
|
||||
enforce resource budgets. Unsupported backends reject constrained refreshes.
|
||||
Older servers reject the new options. Saved defaults and automatic refresh
|
||||
configuration are unchanged.
|
||||
|
||||
## Watching a refresh {#watching}
|
||||
|
||||
A refresh submits work that outlives the call. Track it with
|
||||
|
||||
@@ -654,10 +654,24 @@ class Table:
|
||||
limit: Optional[int] = None,
|
||||
) -> FunctionErrors: ...
|
||||
async def refresh_materialized_view(
|
||||
self, source_version: Optional[int] = None, *, full: bool = False
|
||||
self,
|
||||
source_version: Optional[int] = None,
|
||||
*,
|
||||
full: bool = False,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = None,
|
||||
) -> RefreshMaterializedViewResult: ...
|
||||
async def refresh_materialized_view_async(
|
||||
self, source_version: Optional[int] = None, *, full: bool = False
|
||||
self,
|
||||
source_version: Optional[int] = None,
|
||||
*,
|
||||
full: bool = False,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = None,
|
||||
) -> Job: ...
|
||||
async def materialized_view_definition(self) -> str: ...
|
||||
async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ...
|
||||
|
||||
@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .functions import _validate_gpu_class_name
|
||||
from .background_loop import LOOP
|
||||
from .job import AsyncJob, Job, _typed_job
|
||||
|
||||
@@ -156,6 +157,34 @@ def normalize_select(select: SelectArg) -> Optional[List[Tuple[str, str]]]:
|
||||
return normalized
|
||||
|
||||
|
||||
def _refresh_execution_options(cpu_millis, memory_bytes, gpu_count, gpu_class):
|
||||
values = {
|
||||
"cpu_millis": cpu_millis,
|
||||
"memory_bytes": memory_bytes,
|
||||
"gpu_count": gpu_count,
|
||||
}
|
||||
for name, value in values.items():
|
||||
maximum = 2**32 - 1 if name == "gpu_count" else 2**64 - 1
|
||||
minimum = 0 if name == "gpu_count" else 1
|
||||
if value is not None and (
|
||||
type(value) is not int or not minimum <= value <= maximum
|
||||
):
|
||||
raise ValueError(
|
||||
f"{name} must be an integer between {minimum} and {maximum}"
|
||||
)
|
||||
if (cpu_millis is not None or memory_bytes is not None) and any(
|
||||
value is None for value in values.values()
|
||||
):
|
||||
raise ValueError(
|
||||
"cpu_millis, memory_bytes, and gpu_count must be specified together"
|
||||
)
|
||||
if gpu_class is not None:
|
||||
values["gpu_class"] = _validate_gpu_class_name(gpu_class)
|
||||
if gpu_count == 0:
|
||||
raise ValueError("gpu_class requires a positive GPU budget")
|
||||
return {name: value for name, value in values.items() if value is not None}
|
||||
|
||||
|
||||
class AsyncMaterializedView:
|
||||
"""A handle on a materialized view: its table plus its definition.
|
||||
|
||||
@@ -185,7 +214,14 @@ class AsyncMaterializedView:
|
||||
return _definition_from_json(raw)
|
||||
|
||||
async def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self,
|
||||
*,
|
||||
full: bool = False,
|
||||
source_version: Optional[int] = None,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = None,
|
||||
) -> "RefreshMaterializedViewResult":
|
||||
"""Recompute the view from its source.
|
||||
|
||||
@@ -194,16 +230,35 @@ class AsyncMaterializedView:
|
||||
and ``source_version`` selects a source version instead of the latest.
|
||||
LanceDB Enterprise SQL refreshes do not support source-version pinning.
|
||||
|
||||
``cpu_millis`` (1000 = one core), ``memory_bytes``, and ``gpu_count``
|
||||
cap aggregate resources for this execution, including its Function
|
||||
processes. CPU or memory requires all three budget fields; GPU count
|
||||
may be zero for CPU-only views. ``gpu_class`` selects an operator-defined
|
||||
GPU class for this execution only. Omitted options preserve defaults.
|
||||
These options require an Enterprise backend that enforces isolation;
|
||||
Embedded refreshes reject them. Other servers may reject the new options.
|
||||
|
||||
LanceDB Enterprise results leave ``mode``, ``rows_written``, and
|
||||
``source_version`` as ``None`` because the SQL job adapter does not expose
|
||||
these metrics.
|
||||
"""
|
||||
return await self._table._inner.refresh_materialized_view(
|
||||
source_version=source_version, full=full
|
||||
source_version=source_version,
|
||||
full=full,
|
||||
**_refresh_execution_options(
|
||||
cpu_millis, memory_bytes, gpu_count, gpu_class
|
||||
),
|
||||
)
|
||||
|
||||
async def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self,
|
||||
*,
|
||||
full: bool = False,
|
||||
source_version: Optional[int] = None,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = None,
|
||||
) -> "AsyncJob[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
@@ -221,7 +276,11 @@ class AsyncMaterializedView:
|
||||
|
||||
return _typed_job(
|
||||
await self._table._inner.refresh_materialized_view_async(
|
||||
source_version=source_version, full=full
|
||||
source_version=source_version,
|
||||
full=full,
|
||||
**_refresh_execution_options(
|
||||
cpu_millis, memory_bytes, gpu_count, gpu_class
|
||||
),
|
||||
),
|
||||
RefreshMaterializedViewResult.from_json,
|
||||
)
|
||||
@@ -253,14 +312,36 @@ class MaterializedView:
|
||||
return LOOP.run(self._async.definition())
|
||||
|
||||
def refresh(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self,
|
||||
*,
|
||||
full: bool = False,
|
||||
source_version: Optional[int] = None,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = 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(
|
||||
full=full,
|
||||
source_version=source_version,
|
||||
**_refresh_execution_options(
|
||||
cpu_millis, memory_bytes, gpu_count, gpu_class
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def refresh_async(
|
||||
self, *, full: bool = False, source_version: Optional[int] = None
|
||||
self,
|
||||
*,
|
||||
full: bool = False,
|
||||
source_version: Optional[int] = None,
|
||||
cpu_millis: Optional[int] = None,
|
||||
memory_bytes: Optional[int] = None,
|
||||
gpu_count: Optional[int] = None,
|
||||
gpu_class: Optional[str] = None,
|
||||
) -> "Job[RefreshMaterializedViewResult]":
|
||||
"""Submit a refresh and return its job without waiting.
|
||||
|
||||
@@ -269,6 +350,12 @@ class MaterializedView:
|
||||
"""
|
||||
return Job(
|
||||
LOOP.run(
|
||||
self._async.refresh_async(full=full, source_version=source_version)
|
||||
self._async.refresh_async(
|
||||
full=full,
|
||||
source_version=source_version,
|
||||
**_refresh_execution_options(
|
||||
cpu_millis, memory_bytes, gpu_count, gpu_class
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -63,7 +63,17 @@ def mock_remote_materialized_view_refresh():
|
||||
@pytest.mark.parametrize("sync", [False, True])
|
||||
@pytest.mark.parametrize("background", [False, True])
|
||||
@pytest.mark.parametrize("full", [False, True])
|
||||
async def test_remote_refresh_reports_unknown_metrics(sync, background, full):
|
||||
@pytest.mark.parametrize(
|
||||
"resources",
|
||||
[
|
||||
{},
|
||||
{"gpu_count": 2, "gpu_class": "h100"},
|
||||
{"cpu_millis": 8000, "memory_bytes": 16 * 1024**3, "gpu_count": 0},
|
||||
],
|
||||
)
|
||||
async def test_remote_refresh_reports_unknown_metrics(
|
||||
sync, background, full, resources
|
||||
):
|
||||
flight = pytest.importorskip("pyarrow.flight")
|
||||
|
||||
class RefreshServer(flight.FlightServerBase):
|
||||
@@ -90,7 +100,7 @@ async def test_remote_refresh_reports_unknown_metrics(sync, background, full):
|
||||
mock_remote_materialized_view_refresh() as (host, requests),
|
||||
RefreshServer() as server,
|
||||
):
|
||||
options = {"full": True} if full else {}
|
||||
options = {"full": full, **resources}
|
||||
connect = lancedb.connect if sync else lancedb.connect_async
|
||||
db = connect(
|
||||
"db://dev",
|
||||
@@ -140,7 +150,28 @@ async def test_remote_refresh_reports_unknown_metrics(sync, background, full):
|
||||
)
|
||||
(refresh,) = refreshes
|
||||
assert b'REFRESH MATERIALIZED VIEW "view"' in refresh
|
||||
assert (b"WITH (full = true)" in refresh) is full
|
||||
assert (b"full = true" in refresh) is full
|
||||
for key, value in resources.items():
|
||||
expected = (
|
||||
f"{key} = '{value}'" if key == "gpu_class" else f"{key} = {value}"
|
||||
)
|
||||
assert expected.encode() in refresh
|
||||
previous = len(server.commands)
|
||||
for invalid in [
|
||||
{"cpu_millis": 0},
|
||||
{"memory_bytes": 1},
|
||||
{"gpu_count": -1},
|
||||
{"gpu_count": True},
|
||||
{"gpu_count": 2**32},
|
||||
{"gpu_class": "H200:8"},
|
||||
{"gpu_count": 0, "gpu_class": "h100"},
|
||||
]:
|
||||
with pytest.raises(ValueError):
|
||||
if sync:
|
||||
view.refresh(**invalid)
|
||||
else:
|
||||
await view.refresh(**invalid)
|
||||
assert len(server.commands) == previous
|
||||
assert not any("count_rows" in path for path in requests)
|
||||
|
||||
|
||||
|
||||
+36
-2
@@ -2309,11 +2309,16 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (source_version=None, *, full=false))]
|
||||
#[pyo3(signature = (source_version=None, *, full=false, cpu_millis=None, memory_bytes=None, gpu_count=None, gpu_class=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn refresh_materialized_view(
|
||||
self_: PyRef<'_, Self>,
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
cpu_millis: Option<u64>,
|
||||
memory_bytes: Option<u64>,
|
||||
gpu_count: Option<u32>,
|
||||
gpu_class: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -2321,6 +2326,18 @@ impl Table {
|
||||
.await
|
||||
.infer_error()?;
|
||||
let mut builder = view.refresh().full(full);
|
||||
if let Some(value) = cpu_millis {
|
||||
builder = builder.cpu_millis(value);
|
||||
}
|
||||
if let Some(value) = memory_bytes {
|
||||
builder = builder.memory_bytes(value);
|
||||
}
|
||||
if let Some(value) = gpu_count {
|
||||
builder = builder.gpu_count(value);
|
||||
}
|
||||
if let Some(value) = gpu_class {
|
||||
builder = builder.gpu_class(value);
|
||||
}
|
||||
if let Some(version) = source_version {
|
||||
builder = builder.source_version(version);
|
||||
}
|
||||
@@ -2329,11 +2346,16 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (source_version=None, *, full=false))]
|
||||
#[pyo3(signature = (source_version=None, *, full=false, cpu_millis=None, memory_bytes=None, gpu_count=None, gpu_class=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn refresh_materialized_view_async(
|
||||
self_: PyRef<'_, Self>,
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
cpu_millis: Option<u64>,
|
||||
memory_bytes: Option<u64>,
|
||||
gpu_count: Option<u32>,
|
||||
gpu_class: Option<String>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
@@ -2341,6 +2363,18 @@ impl Table {
|
||||
.await
|
||||
.infer_error()?;
|
||||
let mut builder = view.refresh().full(full);
|
||||
if let Some(value) = cpu_millis {
|
||||
builder = builder.cpu_millis(value);
|
||||
}
|
||||
if let Some(value) = memory_bytes {
|
||||
builder = builder.memory_bytes(value);
|
||||
}
|
||||
if let Some(value) = gpu_count {
|
||||
builder = builder.gpu_count(value);
|
||||
}
|
||||
if let Some(value) = gpu_class {
|
||||
builder = builder.gpu_class(value);
|
||||
}
|
||||
if let Some(version) = source_version {
|
||||
builder = builder.source_version(version);
|
||||
}
|
||||
|
||||
@@ -1863,18 +1863,119 @@ impl MaterializedView {
|
||||
view: self.clone(),
|
||||
source_version: None,
|
||||
full: false,
|
||||
options: Default::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-execution resource caps and GPU placement for Enterprise MV refreshes.
|
||||
///
|
||||
/// CPU is in millicores (1000 = one core), memory is bytes, and GPUs are a
|
||||
/// simultaneous cap across the refresh, not a per-Function allocation. CPU or
|
||||
/// memory requires all three budget fields. Embedded refreshes reject options.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
#[non_exhaustive]
|
||||
pub struct RefreshMaterializedViewOptions {
|
||||
pub cpu_millis: Option<u64>,
|
||||
pub memory_bytes: Option<u64>,
|
||||
pub gpu_count: Option<u32>,
|
||||
pub gpu_class: Option<String>,
|
||||
}
|
||||
|
||||
impl RefreshMaterializedViewOptions {
|
||||
pub(crate) fn is_empty(&self) -> bool {
|
||||
self.cpu_millis.is_none()
|
||||
&& self.memory_bytes.is_none()
|
||||
&& self.gpu_count.is_none()
|
||||
&& self.gpu_class.is_none()
|
||||
}
|
||||
|
||||
pub(crate) fn validate(&self) -> Result<()> {
|
||||
let invalid = |message: &str| Error::InvalidInput {
|
||||
message: message.into(),
|
||||
};
|
||||
if self.cpu_millis == Some(0) || self.memory_bytes == Some(0) {
|
||||
return Err(invalid(
|
||||
"cpu_millis and memory_bytes must be positive integers",
|
||||
));
|
||||
}
|
||||
if (self.cpu_millis.is_some() || self.memory_bytes.is_some())
|
||||
&& !(self.cpu_millis.is_some()
|
||||
&& self.memory_bytes.is_some()
|
||||
&& self.gpu_count.is_some())
|
||||
{
|
||||
return Err(invalid(
|
||||
"cpu_millis, memory_bytes, and gpu_count must be specified together",
|
||||
));
|
||||
}
|
||||
if let Some(class) = &self.gpu_class {
|
||||
crate::function::validate_gpu_class_name(class)?;
|
||||
if self.gpu_count == Some(0) {
|
||||
return Err(invalid("gpu_class requires a positive GPU budget"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn sql_clause(&self, full: bool) -> Result<String> {
|
||||
self.validate()?;
|
||||
let mut values = Vec::new();
|
||||
if full {
|
||||
values.push("full = true".to_string());
|
||||
}
|
||||
if let Some(cpu) = self.cpu_millis {
|
||||
values.push(format!("cpu_millis = {cpu}"));
|
||||
}
|
||||
if let Some(memory) = self.memory_bytes {
|
||||
values.push(format!("memory_bytes = {memory}"));
|
||||
}
|
||||
if let Some(gpu) = self.gpu_count {
|
||||
values.push(format!("gpu_count = {gpu}"));
|
||||
}
|
||||
if let Some(class) = &self.gpu_class {
|
||||
values.push(format!("gpu_class = '{class}'"));
|
||||
}
|
||||
Ok(if values.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" WITH ({})", values.join(", "))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a refresh. Created by [`MaterializedView::refresh`].
|
||||
pub struct RefreshMaterializedViewBuilder {
|
||||
view: MaterializedView,
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
options: RefreshMaterializedViewOptions,
|
||||
}
|
||||
|
||||
impl RefreshMaterializedViewBuilder {
|
||||
/// Cap total CPU capacity in millicores; requires memory and GPU budgets.
|
||||
pub fn cpu_millis(mut self, cpu_millis: u64) -> Self {
|
||||
self.options.cpu_millis = Some(cpu_millis);
|
||||
self
|
||||
}
|
||||
|
||||
/// Cap total memory in bytes; requires CPU and GPU budgets.
|
||||
pub fn memory_bytes(mut self, memory_bytes: u64) -> Self {
|
||||
self.options.memory_bytes = Some(memory_bytes);
|
||||
self
|
||||
}
|
||||
|
||||
/// Cap concurrent GPUs for this refresh. Zero is CPU-only.
|
||||
pub fn gpu_count(mut self, gpu_count: u32) -> Self {
|
||||
self.options.gpu_count = Some(gpu_count);
|
||||
self
|
||||
}
|
||||
|
||||
/// Select an operator-configured GPU class for this execution only.
|
||||
pub fn gpu_class(mut self, gpu_class: impl Into<String>) -> Self {
|
||||
self.options.gpu_class = Some(gpu_class.into());
|
||||
self
|
||||
}
|
||||
|
||||
/// Force a complete rebuild instead of letting the server choose
|
||||
/// the refresh mode. Defaults to `false`. LanceDB Embedded views always rebuild.
|
||||
///
|
||||
@@ -1900,14 +2001,22 @@ impl RefreshMaterializedViewBuilder {
|
||||
|
||||
/// Submit the refresh and return a job that settles with its result.
|
||||
pub async fn execute_async(self) -> Result<Job<RefreshMaterializedViewResult>> {
|
||||
self.options.validate()?;
|
||||
if self.view.table.as_native().is_none() {
|
||||
return self
|
||||
.view
|
||||
.table
|
||||
.base_table()
|
||||
.refresh_materialized_view_async(self.source_version, self.full)
|
||||
.refresh_materialized_view_async_with_options(
|
||||
self.source_version,
|
||||
self.full,
|
||||
self.options,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
if !self.options.is_empty() {
|
||||
return Err(Error::NotSupported { message: "resource budgets and GPU placement require an Enterprise materialized-view refresh".into() });
|
||||
}
|
||||
crate::require_experimental_functions()?;
|
||||
Ok(Job::spawned(tokio::spawn(async move {
|
||||
refresh::execute_refresh(&self.view.table, self.source_version).await
|
||||
@@ -1916,6 +2025,10 @@ impl RefreshMaterializedViewBuilder {
|
||||
|
||||
/// Refresh the view, waiting for the job to finish.
|
||||
pub async fn execute(self) -> Result<RefreshMaterializedViewResult> {
|
||||
self.options.validate()?;
|
||||
if self.view.table.as_native().is_some() && !self.options.is_empty() {
|
||||
return Err(Error::NotSupported { message: "resource budgets and GPU placement require an Enterprise materialized-view refresh".into() });
|
||||
}
|
||||
if self.view.table.as_native().is_some() {
|
||||
return refresh::execute_refresh(&self.view.table, self.source_version).await;
|
||||
}
|
||||
@@ -2151,6 +2264,15 @@ mod tests {
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
view.refresh().gpu_count(0).execute().await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
assert!(matches!(
|
||||
view.refresh().gpu_class("h100").execute_async().await,
|
||||
Err(Error::NotSupported { .. })
|
||||
));
|
||||
|
||||
assert_eq!(view.table().count_rows(None).await.unwrap(), 2);
|
||||
assert!(
|
||||
view.table()
|
||||
|
||||
@@ -1329,7 +1329,30 @@ async fn remote_refresh_omits_unknown_metrics_and_forwards_full() {
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
for result in [first, second] {
|
||||
let third = view
|
||||
.refresh()
|
||||
.full(true)
|
||||
.cpu_millis(8000)
|
||||
.memory_bytes(16 << 30)
|
||||
.gpu_count(2)
|
||||
.gpu_class("h100")
|
||||
.execute_async()
|
||||
.await
|
||||
.unwrap()
|
||||
.wait()
|
||||
.await
|
||||
.unwrap();
|
||||
for invalid in [
|
||||
view.refresh().cpu_millis(1),
|
||||
view.refresh().gpu_class("H200:8"),
|
||||
view.refresh().gpu_count(0).gpu_class("h100"),
|
||||
] {
|
||||
assert!(matches!(
|
||||
invalid.execute_async().await,
|
||||
Err(Error::InvalidInput { .. })
|
||||
));
|
||||
}
|
||||
for result in [first, second, third] {
|
||||
assert_eq!(result.mode, None);
|
||||
assert_eq!(result.rows_written, None);
|
||||
assert_eq!(result.source_version, None);
|
||||
@@ -1350,6 +1373,7 @@ async fn remote_refresh_omits_unknown_metrics_and_forwards_full() {
|
||||
"DESCRIBE MATERIALIZED VIEW \"view\"",
|
||||
"REFRESH MATERIALIZED VIEW \"view\"",
|
||||
"REFRESH MATERIALIZED VIEW \"view\" WITH (full = true)",
|
||||
"REFRESH MATERIALIZED VIEW \"view\" WITH (full = true, cpu_millis = 8000, memory_bytes = 17179869184, gpu_count = 2, gpu_class = 'h100')",
|
||||
]
|
||||
);
|
||||
assert!(
|
||||
|
||||
@@ -2379,6 +2379,17 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
) -> Result<Job<RefreshMaterializedViewResult>> {
|
||||
self.refresh_materialized_view_async_with_options(source_version, full, Default::default())
|
||||
.await
|
||||
}
|
||||
|
||||
async fn refresh_materialized_view_async_with_options(
|
||||
&self,
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
options: crate::materialized_view::RefreshMaterializedViewOptions,
|
||||
) -> Result<Job<RefreshMaterializedViewResult>> {
|
||||
let with_clause = options.sql_clause(full)?;
|
||||
self.check_mutable().await?;
|
||||
if source_version.is_some() {
|
||||
return Err(Error::NotSupported {
|
||||
@@ -2392,7 +2403,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
|
||||
let statement = format!(
|
||||
"REFRESH MATERIALIZED VIEW {}{}",
|
||||
quote_sql_identifier(&self.name),
|
||||
if full { " WITH (full = true)" } else { "" },
|
||||
with_clause,
|
||||
);
|
||||
let namespace = self.namespace.clone();
|
||||
let client = self.client.clone();
|
||||
|
||||
@@ -596,6 +596,25 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
|
||||
.append_to("remote materialized-view refresh is not supported on this table type"),
|
||||
})
|
||||
}
|
||||
/// Submit an Enterprise refresh with aggregate resource caps and GPU placement.
|
||||
#[doc(hidden)]
|
||||
async fn refresh_materialized_view_async_with_options(
|
||||
&self,
|
||||
source_version: Option<u64>,
|
||||
full: bool,
|
||||
options: crate::materialized_view::RefreshMaterializedViewOptions,
|
||||
) -> Result<Job<crate::materialized_view::RefreshMaterializedViewResult>> {
|
||||
options.validate()?;
|
||||
if !options.is_empty() {
|
||||
return Err(Error::NotSupported {
|
||||
message:
|
||||
"this table cannot enforce materialized-view resource budgets or GPU placement"
|
||||
.into(),
|
||||
});
|
||||
}
|
||||
self.refresh_materialized_view_async(source_version, full)
|
||||
.await
|
||||
}
|
||||
/// Create a read-only handle pinned to the table's current active revision.
|
||||
///
|
||||
/// The returned handle is independent from later refreshes or checkouts on
|
||||
|
||||
Reference in New Issue
Block a user