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

feat: persist automatic Function refresh resource budgets (#4953)

Browse Source 
## Depends on merged #4951

Phase 1 merged as `4d22c0ce2b6825c880c818e737b079db494b4053`, so this PR
now targets `main`. It reuses the canonical
`FunctionWorkerResourceBudget` introduced there rather than copying the
type.

For Function `add_columns` declarations, Python and Rust now persist
`auto_refresh_resource_budget` in computed-column metadata and send it
as `auto_refresh_resource_budget` in the remote request. CPU
milliseconds and memory bytes are required positive integers. Omitting
`gpu_count` is accepted as zero and the canonical outgoing JSON still
contains `"gpu_count": 0`. Unknown fields, missing CPU or memory, wrong
types, zero CPU or memory, negative values, and overflow are rejected.

The same budget is applied to every public output binding of a
multi-output Function declaration. Existing metadata is retained, and
SQL-expression declarations or local execution reject the budget instead
of silently ignoring it. Node is unchanged because its current
`addColumns` surface only accepts SQL computed expressions and has no
Function declaration path.

This change does not implement fleet call concurrency limits or enforce
the budget in the runtime. The Enterprise server must support the nested
Function worker resource-budget contract before these values can act as
execution limits.

## Validation

- Python extension build passed.
- Remote Function declaration and local unsupported-path rejection each
passed.
- Canonical metadata round-trip passed, including omitted GPU decoding
to zero and explicit `gpu_count: 0` on encode.
- Function multi-output remote wire test passed with `--features
remote`.
- Repository-scoped Python type check passed with 0 errors.
- All runtime source blobs match the native-tested source. The
post-native deltas are one documentation example line and an existing
routing-test stub update for the new optional argument; the affected
parameterized test passed 3/3 against the existing extension build. No
private test runners are part of this PR.
This commit is contained in:
Xuanwo
2026-11-11 11:23:17 +08:00
committed by GitHub
parent c9d176ae7e
commit 8029428c71
13 changed files with 427 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/web/feature-engineering/add-computed-column.mdx
+43 -3
View File
@@ -60,9 +60,49 @@ columns through `col(...)`, keyed by the output column name. It commits a
new, empty table version the same way a SQL expression does, so the usual
[versioning](/management/table#versioning) applies.
Both a Function application and a SQL expression accept the same
`auto_refresh_enabled` and `auto_refresh_interval_seconds` options, and both
refresh and toggle the same way; see [Auto Refresh](/enterprise/auto-refresh).
Both a Function application and a SQL expression accept
`auto_refresh_enabled` and `auto_refresh_interval_seconds`, and both refresh
and toggle the same way; see [Auto Refresh](/enterprise/auto-refresh).
Python and Rust Function declarations can also set
`auto_refresh_resource_budget` to give their automatic refreshes an aggregate
CPU, memory, and optional GPU budget. CPU and memory are required. The nested
`gpu_count` defaults to zero, so omitting it requests CPU-only workers:
<CodeGroup>
```python Python icon="python"
table.add_columns(
{"caption": version(jpeg=col("jpeg"))},
auto_refresh_resource_budget={
"cpu_millis": 8000,
"memory_bytes": 16 * 1024**3,
},
)
```
```rust Rust icon="rust"
table
.add_columns()
.function(function_application)
.computed_options(ComputedColumnOptions {
auto_refresh_enabled: true,
auto_refresh_interval_seconds: None,
auto_refresh_resource_budget: Some(
FunctionWorkerResourceBudget::new(8000, 17_179_869_184),
),
})
.execute()
.await?;
```
</CodeGroup>
This setting applies to every output in the Function binding. SQL-expression
columns do not use Function workers, so they reject
`auto_refresh_resource_budget` instead of accepting an option that has no
effect. The separate `resource_budget` refresh option configures one explicit
refresh and does not change this automatic-refresh declaration. The Node.js
SDK does not currently expose a Function declaration through `addColumns`, so
this option is not added to an unsupported surface.
## Fill a column with a SQL expression {#expressions}
docs/web/feature-engineering/refresh-column.mdx
+5 -3
View File
@@ -219,9 +219,11 @@ for changing only the aggregate GPU budget. It cannot be combined with the
complete resource budget. Omitting both preserves the server's existing
defaults.
These examples configure only the explicit refresh shown. They do not set the
budget for an automatic refresh scheduled after `add_columns`; that remains a
separate declaration and scheduler capability.
These examples configure only the explicit refresh shown. A registered
Function declaration can separately set `auto_refresh_resource_budget` on
`add_columns`; that declaration controls the budget for its automatic
refreshes. SQL-expression columns reject the automatic resource-budget option
because they do not use Function workers.
`batch_size` caps rows in one Function call; LanceDB adapts actual call and
work sizes. For Function columns, the refresh defaults to
python/python/lancedb/_lancedb.pyi
+2
View File
@@ -615,6 +615,7 @@ class Table:
columns: list[tuple[str, str]],
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[int] = None,
auto_refresh_resource_budget: Optional[dict[str, int]] = None,
) -> AddColumnsResult: ...
async def add_function_columns(
self,
@@ -622,6 +623,7 @@ class Table:
output_name: Optional[str],
auto_refresh_enabled: Optional[bool] = None,
auto_refresh_interval_seconds: Optional[int] = None,
auto_refresh_resource_budget: Optional[dict[str, int]] = None,
) -> AddColumnsResult: ...
async def refresh_column(
self,
python/python/lancedb/remote/table.py
+2
View File
@@ -1127,6 +1127,7 @@ class RemoteTable(Table):
computed: Dict[str, str] | None = None,
auto_refresh_enabled: bool | None = None,
auto_refresh_interval_seconds: int | None = None,
auto_refresh_resource_budget: Optional[Dict[str, int]] = None,
) -> AddColumnsResult:
return LOOP.run(
self._table.add_columns(
@@ -1134,6 +1135,7 @@ class RemoteTable(Table):
computed=computed,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
auto_refresh_resource_budget=auto_refresh_resource_budget,
)
)
python/python/lancedb/table.py
+19
View File
@@ -3085,6 +3085,7 @@ class Table(ABC):
computed: Dict[str, str] | None = None,
auto_refresh_enabled: bool | None = None,
auto_refresh_interval_seconds: int | None = None,
auto_refresh_resource_budget: Optional[Dict[str, int]] = None,
):
"""
Add new columns with defined values.
@@ -3153,6 +3154,12 @@ class Table(ABC):
When omitted, the server determines when to refresh. Applies to every
computed output in this call.
Scheduling options require computed or Function columns.
auto_refresh_resource_budget: dict[str, int], optional
Aggregate CPU, memory, and GPU budget for automatic refreshes of a
registered Function declaration. ``cpu_millis`` and
``memory_bytes`` are required; ``gpu_count`` defaults to zero.
This option cannot be used with SQL-expression columns. LanceDB
Enterprise only.
Returns
-------
@@ -5468,6 +5475,7 @@ class LanceTable(Table):
computed: Dict[str, str] | None = None,
auto_refresh_enabled: bool | None = None,
auto_refresh_interval_seconds: int | None = None,
auto_refresh_resource_budget: Optional[Dict[str, int]] = None,
) -> AddColumnsResult:
return LOOP.run(
self._table.add_columns(
@@ -5475,6 +5483,7 @@ class LanceTable(Table):
computed=computed,
auto_refresh_enabled=auto_refresh_enabled,
auto_refresh_interval_seconds=auto_refresh_interval_seconds,
auto_refresh_resource_budget=auto_refresh_resource_budget,
)
)
@@ -7758,6 +7767,7 @@ class AsyncTable:
computed: dict[str, str] | None = None,
auto_refresh_enabled: bool | None = None,
auto_refresh_interval_seconds: int | None = None,
auto_refresh_resource_budget: Optional[Dict[str, int]] = None,
) -> AddColumnsResult:
"""
Add new columns with defined values.
@@ -7824,6 +7834,12 @@ class AsyncTable:
not force recomputation. Applies to every declared computed output.
These options cannot be used with one-time transforms or ordinary
all-null columns.
auto_refresh_resource_budget: dict[str, int], optional
Aggregate CPU, memory, and GPU budget for automatic refreshes of
registered Function outputs. ``cpu_millis`` and ``memory_bytes``
are required; ``gpu_count`` defaults to zero. SQL-expression
columns do not use Function workers and reject this option.
LanceDB Enterprise only.
Examples
--------
@@ -7852,6 +7868,7 @@ class AsyncTable:
options_supplied = (
auto_refresh_enabled is not None
or auto_refresh_interval_seconds is not None
or auto_refresh_resource_budget is not None
)
function_application = None
function_output_name = None
@@ -7880,6 +7897,7 @@ class AsyncTable:
function_output_name,
auto_refresh_enabled,
auto_refresh_interval_seconds,
auto_refresh_resource_budget,
)
if isinstance(transforms, pa.Field):
@@ -7900,6 +7918,7 @@ class AsyncTable:
list(computed.items()),
auto_refresh_enabled,
auto_refresh_interval_seconds,
auto_refresh_resource_budget,
)
if options_supplied:
raise ValueError(
python/python/tests/test_first_class_function_slice1.py
+4 -2
View File
@@ -459,6 +459,7 @@ class _FunctionDeclarationInner:
output_name,
auto_refresh_enabled,
auto_refresh_interval_seconds,
auto_refresh_resource_budget,
):
self.calls.append(
(
@@ -466,6 +467,7 @@ class _FunctionDeclarationInner:
output_name,
auto_refresh_enabled,
auto_refresh_interval_seconds,
auto_refresh_resource_budget,
)
)
return "declared"
@@ -491,7 +493,7 @@ async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomica
)
assert result == "declared"
assert inner.calls[-1][1] == "features"
assert inner.calls[-1][2:] == (None, None)
assert inner.calls[-1][2:] == (None, None, None)
bare = application._copy(update={"columns": {}}).rename(
columns={"normalized_text": "search_text"}
@@ -502,7 +504,7 @@ async def test_add_columns_routes_struct_as_one_and_multi_output_binding_atomica
assert result == "declared"
assert inner.calls[-1][1] is None
assert inner.calls[-1][0]["columns"] == {"normalized_text": "search_text"}
assert inner.calls[-1][2:] == (False, 300)
assert inner.calls[-1][2:] == (False, 300, None)
@pytest.mark.asyncio
python/python/tests/test_remote_db.py
+87
View File
@@ -36,6 +36,7 @@ from lancedb.graph import (
PropertyGraph,
PropertyGraphDescription,
)
from lancedb.functions import FunctionApplication
from lancedb.query import (
AsyncQuery,
ColumnOrdering,
@@ -5831,6 +5832,92 @@ def test_remote_add_columns_reader_not_supported():
assert all("/add_columns/" not in path for path in requests)
@pytest.mark.asyncio
async def test_remote_add_function_columns_forwards_auto_refresh_resource_budget():
bodies = []
application = FunctionApplication.from_json(
json.dumps(
{
"function": {
"name": "embed",
"version": "1",
"object_id": "fixture",
"location": "memory:///fixture",
"manifest_digest": (
"sha256:7e22f815b6648e14f093a3979a8e5a2082"
"fa773ebe1ec84b135cae7e84d6f8e6"
),
},
"inputs": [
{
"parameter": "text",
"kind": "column",
"value": {"path": "text"},
}
],
"output": {
"kind": "scalar",
"arrow_type": "float32",
"nullable": False,
},
}
)
)
def handler(request):
if request.path == "/v1/table/test/describe/":
response = {
"version": 1,
"schema": {
"fields": [
{
"name": "text",
"type": {"type": "string"},
"nullable": True,
}
]
},
}
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(json.dumps(response).encode())
elif request.path == "/v1/table/test/add_columns/":
content_len = int(request.headers.get("Content-Length", 0))
bodies.append(json.loads(request.rfile.read(content_len)))
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"version": 2}')
else:
request.send_response(404)
request.end_headers()
budget = {"cpu_millis": 2_000, "memory_bytes": 4 * 1024**3}
with mock_lancedb_connection(handler) as db:
table = db.open_table("test")
table.add_columns(
{"embedding": application},
auto_refresh_resource_budget=budget,
)
async with mock_lancedb_connection_async(handler) as db:
table = await db.open_table("test")
await table.add_columns(
{"embedding": application},
auto_refresh_resource_budget=budget,
)
assert len(bodies) == 2
for body in bodies:
assert body["auto_refresh_enabled"] is True
assert body["auto_refresh_resource_budget"] == {
**budget,
"gpu_count": 0,
}
assert body["new_columns"] == [{"name": "embedding", "all_null": True}]
@pytest.mark.asyncio
async def test_async_property_graph_handles():
calls = []
python/python/tests/test_table.py
+15
View File
@@ -6255,6 +6255,21 @@ def test_auto_refresh_options_require_computed_columns(tmp_path):
table.add_columns({"doubled": "x * 2"}, auto_refresh_enabled=True)
def test_auto_refresh_resource_budget_requires_a_function(tmp_path):
table = lancedb.connect(tmp_path).create_table("invalid_budget", [{"x": 1}])
before = table.version
with pytest.raises(ValueError, match="registered Function"):
table.add_columns(
computed={"doubled": "x * 2"},
auto_refresh_resource_budget={
"cpu_millis": 1_000,
"memory_bytes": 1_024,
},
)
assert table.version == before
assert table.schema.names == ["x"]
@pytest.mark.experimental_functions
@pytest.mark.asyncio
async def test_async_computed_auto_refresh_uses_server_interval(tmp_path):
python/src/table.rs
+16 -8
View File
@@ -2217,17 +2217,21 @@ impl Table {
})
}
#[pyo3(signature = (columns, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn add_computed_columns(
self_: PyRef<'_, Self>,
#[pyo3(signature = (columns, auto_refresh_enabled=None, auto_refresh_interval_seconds=None, auto_refresh_resource_budget=None))]
pub fn add_computed_columns<'py>(
self_: PyRef<'py, Self>,
columns: Vec<(String, String)>,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
auto_refresh_resource_budget: Option<Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
let inner = self_.inner_ref()?.clone();
let options = ComputedColumnOptions {
auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true),
auto_refresh_interval_seconds,
auto_refresh_resource_budget: parse_resource_budget(
auto_refresh_resource_budget.as_ref(),
)?,
};
options.validate().infer_error()?;
future_into_py(self_.py(), async move {
@@ -2244,20 +2248,24 @@ impl Table {
})
}
#[pyo3(signature = (application_json, output_name, auto_refresh_enabled=None, auto_refresh_interval_seconds=None))]
pub fn add_function_columns(
self_: PyRef<'_, Self>,
#[pyo3(signature = (application_json, output_name, auto_refresh_enabled=None, auto_refresh_interval_seconds=None, auto_refresh_resource_budget=None))]
pub fn add_function_columns<'py>(
self_: PyRef<'py, Self>,
application_json: String,
output_name: Option<String>,
auto_refresh_enabled: Option<bool>,
auto_refresh_interval_seconds: Option<u64>,
) -> PyResult<Bound<'_, PyAny>> {
auto_refresh_resource_budget: Option<Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
let application =
lancedb::function::FunctionApplication::from_json(&application_json).infer_error()?;
let inner = self_.inner_ref()?.clone();
let options = ComputedColumnOptions {
auto_refresh_enabled: auto_refresh_enabled.unwrap_or(true),
auto_refresh_interval_seconds,
auto_refresh_resource_budget: parse_resource_budget(
auto_refresh_resource_budget.as_ref(),
)?,
};
options.validate().infer_error()?;
future_into_py(self_.py(), async move {
rust/lancedb/src/remote/table.rs
+17
View File
@@ -3951,6 +3951,9 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
if let Some(interval) = options.auto_refresh_interval_seconds {
body["auto_refresh_interval_seconds"] = serde_json::json!(interval);
}
if let Some(budget) = options.auto_refresh_resource_budget {
body["auto_refresh_resource_budget"] = serde_json::json!(budget);
}
self.apply_branch_body(&mut body);
let request = self
.client
@@ -10361,6 +10364,7 @@ mod tests {
.computed_options(crate::table::computed_columns::ComputedColumnOptions {
auto_refresh_enabled: false,
auto_refresh_interval_seconds: interval,
auto_refresh_resource_budget: None,
})
.execute()
.await
@@ -10797,6 +10801,11 @@ mod tests {
))
.unwrap();
expected["auto_refresh_enabled"] = serde_json::json!(true);
expected["auto_refresh_resource_budget"] = serde_json::json!({
"cpu_millis": 8_000,
"memory_bytes": 17_179_869_184_u64,
"gpu_count": 2,
});
assert_eq!(actual, expected);
http::Response::builder()
.status(200)
@@ -10824,6 +10833,14 @@ mod tests {
let result = table
.add_columns()
.function(application)
.computed_options(crate::table::computed_columns::ComputedColumnOptions {
auto_refresh_enabled: true,
auto_refresh_interval_seconds: None,
auto_refresh_resource_budget: Some(
crate::table::FunctionWorkerResourceBudget::new(8_000, 17_179_869_184)
.with_gpu_count(2),
),
})
.execute()
.await
.unwrap();
rust/lancedb/src/table/add_columns.rs
+42 -1
View File
@@ -171,7 +171,9 @@ impl AddColumnsBuilder {
/// Configure automatic refresh for every output declared in this call.
/// Settings are stored with the declaration. They require computed or
/// Function columns; LanceDB Embedded persists them without starting a scheduler.
/// Function columns. LanceDB Embedded persists enabled and interval settings
/// without starting a scheduler; an automatic resource budget requires an
/// Enterprise registered Function and rejects SQL-expression columns.
///
/// ```
/// # use lancedb::{Table, table::computed_columns::ComputedColumnOptions};
@@ -180,6 +182,7 @@ impl AddColumnsBuilder {
/// .computed_options(ComputedColumnOptions {
/// auto_refresh_enabled: true,
/// auto_refresh_interval_seconds: Some(300),
/// auto_refresh_resource_budget: None,
/// }).execute().await?;
/// # Ok(())
/// # }
@@ -207,6 +210,13 @@ impl AddColumnsBuilder {
.into(),
});
}
if options.auto_refresh_resource_budget.is_some() && function.is_none() {
return Err(Error::InvalidInput {
message:
"auto_refresh_resource_budget requires a registered Function declaration"
.into(),
});
}
}
let options = options.unwrap_or_default();
@@ -289,6 +299,7 @@ mod tests {
let options = ComputedColumnOptions {
auto_refresh_enabled: false,
auto_refresh_interval_seconds: Some(300),
auto_refresh_resource_budget: None,
};
let result = table
.add_columns()
@@ -373,6 +384,36 @@ mod tests {
);
}
#[tokio::test]
async fn auto_refresh_resource_budget_requires_a_function() {
let table = table_with_two_columns("budget_expression").await;
let err = table
.add_columns()
.computed("doubled", "x * 2")
.computed_options(ComputedColumnOptions {
auto_refresh_enabled: true,
auto_refresh_interval_seconds: None,
auto_refresh_resource_budget: Some(
crate::table::FunctionWorkerResourceBudget::new(1_000, 1_024),
),
})
.execute()
.await
.unwrap_err();
assert!(
matches!(&err, Error::InvalidInput { message } if message.contains("registered Function")),
"{err:?}"
);
assert!(
table
.schema()
.await
.unwrap()
.field_with_name("doubled")
.is_err()
);
}
#[tokio::test]
async fn test_requires_a_transform() {
let table = table_with_two_columns("no_transform").await;
rust/lancedb/src/table/computed_columns.rs
+156 -3
View File
@@ -42,6 +42,7 @@ use serde_json::Value;
use crate::function::{FUNCTION_BLOB_V2_TYPE, FunctionApplication, FunctionBinding};
use crate::materialized_view::without_declarations;
use crate::table::FunctionWorkerResourceBudget;
use crate::utils::resolve_arrow_field_path;
use crate::{Error, Result};
@@ -61,6 +62,9 @@ pub const INPUTS_META_KEY: &str = "computed_column.inputs";
pub const AUTO_REFRESH_ENABLED_META_KEY: &str = "computed_column.auto_refresh_enabled";
/// Minimum interval between successful automatic refreshes, in seconds.
pub const AUTO_REFRESH_INTERVAL_META_KEY: &str = "computed_column.auto_refresh_interval_seconds";
/// Aggregate worker-resource budget for automatic Function refreshes.
pub const AUTO_REFRESH_RESOURCE_BUDGET_META_KEY: &str =
"computed_column.auto_refresh_resource_budget";
/// Scheduling settings persisted on each computed output. LanceDB Embedded
/// stores these settings; automatic scheduling requires LanceDB Enterprise.
@@ -73,6 +77,10 @@ pub struct ComputedColumnOptions {
/// When omitted, the server determines when to refresh.
#[serde(skip_serializing_if = "Option::is_none")]
pub auto_refresh_interval_seconds: Option<u64>,
/// Aggregate worker-resource budget for automatic Function refreshes.
/// This option requires a registered Function output.
#[serde(skip_serializing_if = "Option::is_none")]
pub auto_refresh_resource_budget: Option<FunctionWorkerResourceBudget>,
}
impl Default for ComputedColumnOptions {
@@ -80,6 +88,7 @@ impl Default for ComputedColumnOptions {
Self {
auto_refresh_enabled: true,
auto_refresh_interval_seconds: None,
auto_refresh_resource_budget: None,
}
}
}
@@ -95,6 +104,9 @@ impl ComputedColumnOptions {
message: "auto_refresh_interval_seconds must be a positive integer no greater than u64::MAX / 1000".into(),
});
}
if let Some(budget) = &self.auto_refresh_resource_budget {
budget.validate()?;
}
Ok(())
}
@@ -112,6 +124,15 @@ impl ComputedColumnOptions {
message: "auto_refresh_interval_seconds must be a positive integer".into(),
})?);
}
if let Some(value) = metadata.get(AUTO_REFRESH_RESOURCE_BUDGET_META_KEY) {
options.auto_refresh_resource_budget = Some(serde_json::from_str(value).map_err(
|error| Error::InvalidInput {
message: format!(
"auto_refresh_resource_budget must be a valid resource budget: {error}"
),
},
)?);
}
options.validate()?;
Ok(options)
}
@@ -119,6 +140,7 @@ impl ComputedColumnOptions {
/// Attach validated settings to newly declared computed output fields.
pub fn apply_to_schema(&self, schema: &ArrowSchema) -> Result<ArrowSchema> {
self.validate()?;
let option_metadata = self.metadata()?;
let fields = schema
.fields()
.iter()
@@ -134,7 +156,7 @@ impl ComputedColumnOptions {
return field.clone();
}
let mut metadata = field.metadata().clone();
metadata.extend(self.metadata());
metadata.extend(option_metadata.clone());
Arc::new(field.as_ref().clone().with_metadata(metadata))
})
.collect::<Vec<_>>();
@@ -145,7 +167,7 @@ impl ComputedColumnOptions {
}
/// Field metadata values shared by native and server declaration paths.
pub fn metadata(&self) -> HashMap<String, String> {
pub fn metadata(&self) -> Result<HashMap<String, String>> {
let mut metadata = HashMap::from([(
AUTO_REFRESH_ENABLED_META_KEY.into(),
self.auto_refresh_enabled.to_string(),
@@ -153,7 +175,15 @@ impl ComputedColumnOptions {
if let Some(interval) = self.auto_refresh_interval_seconds {
metadata.insert(AUTO_REFRESH_INTERVAL_META_KEY.into(), interval.to_string());
}
metadata
if let Some(budget) = &self.auto_refresh_resource_budget {
metadata.insert(
AUTO_REFRESH_RESOURCE_BUDGET_META_KEY.into(),
serde_json::to_string(budget).map_err(|error| Error::InvalidInput {
message: format!("failed to encode auto_refresh_resource_budget: {error}"),
})?,
);
}
Ok(metadata)
}
}
@@ -2533,6 +2563,129 @@ pub(crate) mod test_support {
#[cfg(test)]
mod tests {
#[test]
fn computed_auto_refresh_resource_budget_metadata_round_trips() {
let budget = crate::table::FunctionWorkerResourceBudget::new(8_000, 17_179_869_184)
.with_gpu_count(2);
let options = super::ComputedColumnOptions {
auto_refresh_enabled: true,
auto_refresh_interval_seconds: Some(300),
auto_refresh_resource_budget: Some(budget),
};
assert_eq!(
super::ComputedColumnOptions::from_metadata(&HashMap::new()).unwrap(),
super::ComputedColumnOptions::default()
);
let metadata = options.metadata().unwrap();
assert_eq!(
metadata
.get(super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY)
.map(String::as_str),
Some(r#"{"cpu_millis":8000,"memory_bytes":17179869184,"gpu_count":2}"#)
);
assert_eq!(
super::ComputedColumnOptions::from_metadata(&metadata).unwrap(),
options
);
let mut omitted_gpu = metadata.clone();
omitted_gpu.insert(
super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY.into(),
r#"{"cpu_millis":8000,"memory_bytes":17179869184}"#.into(),
);
let decoded = super::ComputedColumnOptions::from_metadata(&omitted_gpu).unwrap();
assert_eq!(decoded.auto_refresh_resource_budget.unwrap().gpu_count, 0);
assert_eq!(
decoded
.metadata()
.unwrap()
.get(super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY)
.map(String::as_str),
Some(r#"{"cpu_millis":8000,"memory_bytes":17179869184,"gpu_count":0}"#)
);
for invalid in [
r#"{"memory_bytes":17179869184,"gpu_count":2}"#,
r#"{"cpu_millis":8000,"gpu_count":2}"#,
r#"{"cpu_millis":"8000","memory_bytes":17179869184,"gpu_count":2}"#,
r#"{"cpu_millis":-1,"memory_bytes":17179869184,"gpu_count":2}"#,
r#"{"cpu_millis":18446744073709551616,"memory_bytes":17179869184,"gpu_count":2}"#,
r#"{"cpu_millis":8000,"memory_bytes":17179869184,"gpu_count":-1}"#,
r#"{"cpu_millis":8000,"memory_bytes":17179869184,"gpu_count":4294967296}"#,
r#"{"cpu_millis":8000,"memory_bytes":17179869184,"gpu_count":2,"unknown":1}"#,
r#"{"cpu_millis":0,"memory_bytes":17179869184,"gpu_count":2}"#,
r#"{"cpu_millis":8000,"memory_bytes":0,"gpu_count":2}"#,
] {
let mut malformed = metadata.clone();
malformed.insert(
super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY.into(),
invalid.into(),
);
assert!(
super::ComputedColumnOptions::from_metadata(&malformed).is_err(),
"accepted invalid budget: {invalid}"
);
}
}
#[test]
fn computed_auto_refresh_resource_budget_applies_to_function_outputs() {
let function_field = |name: &str, ordinal: u32| {
ArrowField::new(name, DataType::Int32, true).with_metadata(HashMap::from([
(
super::COMPUTED_COLUMN_META_KEY.to_string(),
"true".to_string(),
),
(
super::KIND_META_KEY.to_string(),
super::FUNCTION_KIND.to_string(),
),
(
super::FUNCTION_BINDING_ID_META_KEY.to_string(),
"binding-1".to_string(),
),
(
super::FUNCTION_OUTPUT_ORDINAL_META_KEY.to_string(),
ordinal.to_string(),
),
("custom".to_string(), "preserved".to_string()),
]))
};
let schema = ArrowSchema::new(vec![
ArrowField::new("input", DataType::Int32, true),
function_field("first", 0),
function_field("second", 1),
function_field("assignment", super::FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL),
]);
let budget = crate::table::FunctionWorkerResourceBudget::new(8_000, 17_179_869_184)
.with_gpu_count(2);
let updated = super::ComputedColumnOptions {
auto_refresh_resource_budget: Some(budget),
..Default::default()
}
.apply_to_schema(&schema)
.unwrap();
for name in ["first", "second"] {
let metadata = updated.field_with_name(name).unwrap().metadata();
assert_eq!(
metadata.get(super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY),
Some(&serde_json::to_string(&budget).unwrap())
);
assert_eq!(
metadata.get("custom").map(String::as_str),
Some("preserved")
);
}
assert!(
!updated
.field_with_name("assignment")
.unwrap()
.metadata()
.contains_key(super::AUTO_REFRESH_RESOURCE_BUDGET_META_KEY)
);
}
/// The gate's reproducer: the validator applies the same schema-level
/// guard declaring does, so a staging caller is refused before it commits
/// anything else.
rust/lancedb/src/table/refresh.rs
+19 -11
View File
@@ -61,12 +61,14 @@ use crate::{Error, Result};
/// values cap the Function workers together; they are not per-worker
/// allocations. The default GPU budget is zero, which describes a CPU-only
/// Function. This option is rejected by LanceDB Embedded refreshes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
#[non_exhaustive]
pub struct FunctionWorkerResourceBudget {
pub cpu_millis: u64,
pub memory_bytes: u64,
/// Aggregate GPU budget across the refresh. Zero is CPU-only.
#[serde(default)]
pub gpu_count: u32,
}
@@ -85,6 +87,21 @@ impl FunctionWorkerResourceBudget {
self.gpu_count = gpu_count;
self
}
/// Require every per-tier resource limit to request at least one unit.
pub fn validate(&self) -> Result<()> {
if self.cpu_millis == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.cpu_millis must be greater than zero".into(),
});
}
if self.memory_bytes == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.memory_bytes must be greater than zero".into(),
});
}
Ok(())
}
}
/// Options for refreshing a computed column.
@@ -329,16 +346,7 @@ impl RefreshColumnOptions {
});
}
if let Some(budget) = &self.resource_budget {
if budget.cpu_millis == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.cpu_millis must be greater than zero".into(),
});
}
if budget.memory_bytes == 0 {
return Err(Error::InvalidInput {
message: "resource_budget.memory_bytes must be greater than zero".into(),
});
}
budget.validate()?;
}
if let Some(gpu_class) = &self.gpu_class {
validate_gpu_class_name(gpu_class)?;