diff --git a/docs/web/feature-engineering/add-computed-column.mdx b/docs/web/feature-engineering/add-computed-column.mdx
index 2c9e9be4c..851b37d9b 100644
--- a/docs/web/feature-engineering/add-computed-column.mdx
+++ b/docs/web/feature-engineering/add-computed-column.mdx
@@ -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:
+
+
+```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?;
+```
+
+
+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}
diff --git a/docs/web/feature-engineering/refresh-column.mdx b/docs/web/feature-engineering/refresh-column.mdx
index 043c712cb..62328d790 100644
--- a/docs/web/feature-engineering/refresh-column.mdx
+++ b/docs/web/feature-engineering/refresh-column.mdx
@@ -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
diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi
index 4f676de8c..fafa21cbf 100644
--- a/python/python/lancedb/_lancedb.pyi
+++ b/python/python/lancedb/_lancedb.pyi
@@ -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,
diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py
index f389f313f..4aca6bcc5 100644
--- a/python/python/lancedb/remote/table.py
+++ b/python/python/lancedb/remote/table.py
@@ -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,
)
)
diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py
index f113baa05..a4d40c901 100644
--- a/python/python/lancedb/table.py
+++ b/python/python/lancedb/table.py
@@ -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(
diff --git a/python/python/tests/test_first_class_function_slice1.py b/python/python/tests/test_first_class_function_slice1.py
index 27a9a9376..c392fdc7f 100644
--- a/python/python/tests/test_first_class_function_slice1.py
+++ b/python/python/tests/test_first_class_function_slice1.py
@@ -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
diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py
index 109dbfbdb..2547cf807 100644
--- a/python/python/tests/test_remote_db.py
+++ b/python/python/tests/test_remote_db.py
@@ -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 = []
diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py
index 2422437c2..6c9485ced 100644
--- a/python/python/tests/test_table.py
+++ b/python/python/tests/test_table.py
@@ -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):
diff --git a/python/src/table.rs b/python/src/table.rs
index 56fc3bf5c..07dc49979 100644
--- a/python/src/table.rs
+++ b/python/src/table.rs
@@ -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,
auto_refresh_interval_seconds: Option,
- ) -> PyResult> {
+ auto_refresh_resource_budget: Option>,
+ ) -> PyResult> {
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,
auto_refresh_enabled: Option,
auto_refresh_interval_seconds: Option,
- ) -> PyResult> {
+ auto_refresh_resource_budget: Option>,
+ ) -> PyResult> {
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 {
diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs
index 4f428809c..ee9b3827b 100644
--- a/rust/lancedb/src/remote/table.rs
+++ b/rust/lancedb/src/remote/table.rs
@@ -3951,6 +3951,9 @@ impl BaseTable for RemoteTable {
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();
diff --git a/rust/lancedb/src/table/add_columns.rs b/rust/lancedb/src/table/add_columns.rs
index 1b827bb37..a76c0d043 100644
--- a/rust/lancedb/src/table/add_columns.rs
+++ b/rust/lancedb/src/table/add_columns.rs
@@ -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;
diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs
index e6d7d432c..8fce7a260 100644
--- a/rust/lancedb/src/table/computed_columns.rs
+++ b/rust/lancedb/src/table/computed_columns.rs
@@ -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,
+ /// 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,
}
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 {
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::>();
@@ -145,7 +167,7 @@ impl ComputedColumnOptions {
}
/// Field metadata values shared by native and server declaration paths.
- pub fn metadata(&self) -> HashMap {
+ pub fn metadata(&self) -> Result> {
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.
diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs
index 742d60f18..31c165be6 100644
--- a/rust/lancedb/src/table/refresh.rs
+++ b/rust/lancedb/src/table/refresh.rs
@@ -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)?;