From 8029428c7152e6b018dcab31f5240e3144caab68 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 11 Nov 2026 11:23:17 +0800 Subject: [PATCH] feat: persist automatic Function refresh resource budgets (#4953) ## 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. --- .../add-computed-column.mdx | 46 ++++- .../feature-engineering/refresh-column.mdx | 8 +- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/remote/table.py | 2 + python/python/lancedb/table.py | 19 +++ .../tests/test_first_class_function_slice1.py | 6 +- python/python/tests/test_remote_db.py | 87 ++++++++++ python/python/tests/test_table.py | 15 ++ python/src/table.rs | 24 ++- rust/lancedb/src/remote/table.rs | 17 ++ rust/lancedb/src/table/add_columns.rs | 43 ++++- rust/lancedb/src/table/computed_columns.rs | 159 +++++++++++++++++- rust/lancedb/src/table/refresh.rs | 30 ++-- 13 files changed, 427 insertions(+), 31 deletions(-) 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)?;