diff --git a/docs/web/curation/create-function.mdx b/docs/web/curation/create-function.mdx index 1a35e31be..941d37b2c 100644 --- a/docs/web/curation/create-function.mdx +++ b/docs/web/curation/create-function.mdx @@ -51,7 +51,11 @@ import pyarrow as pa import pyarrow.compute as pc from lancedb import udf -@udf(pip=["numpy==2.2.0"]) +@udf( + pip=["numpy==2.2.0"], + batch_size=32, + row_timeout_seconds=90, +) def score_batch(value: pa.Array) -> pa.Array: return pc.multiply(value, 2) ``` @@ -61,6 +65,13 @@ Batching amortizes per-call overhead across many rows, which matters most once a function also does something with real fixed cost per call -- loading a model, opening a connection, making a network round trip. +`batch_size` caps the number of input rows in each execution batch; actual +batches may be smaller. `row_timeout_seconds` bounds each Python call: one +row for scalar Functions or the entire batch for vectorized Functions. It +does not include the instance's `__init__`. Both settings are fixed into the +registered version. Omit them to use adaptive batching and the service's +default per-call deadline. + ## Stateful Functions: load once, call many times {#stateful} A `@udf`-decorated **class** is the same callable with state. Each remote diff --git a/docs/web/management/function.mdx b/docs/web/management/function.mdx index aeda61ec3..7ea939a71 100644 --- a/docs/web/management/function.mdx +++ b/docs/web/management/function.mdx @@ -65,6 +65,24 @@ fields: whether a row has a result at all is the struct's own validity, which Enterprise keeps in a separate internal column when it expands the result into sibling columns, so a null field is a value rather than a missing row. +Registration can also cap each execution batch and set its row deadline: + + +```python Python icon="python" +@udf(batch_size=32, row_timeout_seconds=90) +def encode(value: float) -> float: + return value +``` + + +`batch_size` is the maximum number of input rows in a batch; execution may use +fewer. `row_timeout_seconds` bounds each Python call: one row for scalar +Functions and one batch for vectorized Functions. It does not include a Function +instance's `__init__`. Both values are fixed into the registered version and +apply to direct calls, computed columns, and materialized-view refreshes. When +omitted, the service chooses batches adaptively and applies its default per-call +budget. + ## Connect and create a table {#table} The remaining steps run against your Enterprise deployment. Replace the database diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 80d1e08eb..ae6de3ab4 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -86,6 +86,27 @@ def _validate_gpu_class_name(value: Any) -> str: return value +def _validate_batch_size(value: Any) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError("batch_size must be an integer") + if not 1 <= value <= 2**32 - 1: + raise ValueError("batch_size must be between 1 and 4294967295") + return value + + +def _validate_row_timeout(value: Any) -> Optional[float]: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("row_timeout_seconds must be a number") + value = float(value) + if not math.isfinite(value) or value <= 0: + raise ValueError("row_timeout_seconds must be positive and finite") + return value + + class _FrozenDict(dict): def _immutable(self, *args, **kwargs): raise TypeError("remote canonical values are immutable") @@ -547,8 +568,20 @@ class FunctionRegistrationRequest(_RemoteValue): artifact: FunctionArtifactRequest signature: FunctionSignature runtime: PythonRuntimeSpec + batch_size: Optional[int] = None + row_timeout_seconds: Optional[float] = None secret_bindings: tuple[SecretBinding, ...] = () + @field_validator("batch_size", mode="before") + @classmethod + def _validate_batch_size(cls, value): + return _validate_batch_size(value) + + @field_validator("row_timeout_seconds", mode="before") + @classmethod + def _validate_row_timeout(cls, value): + return _validate_row_timeout(value) + class FunctionVersionRef(_OpenRemoteValue): name: str @@ -1983,6 +2016,8 @@ class UdfDefinition: pip: tuple[str, ...], env: Mapping[str, str], python_version: Optional[str], + batch_size: Optional[int] = None, + row_timeout_seconds: Optional[float] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] = (), @@ -2105,6 +2140,8 @@ class UdfDefinition: ), signature=signature, runtime=runtime, + batch_size=batch_size, + row_timeout_seconds=row_timeout_seconds, ) # A class's own __dict__ holds its methods; copying it would shadow # this wrapper's attributes. @@ -2199,6 +2236,8 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + batch_size: Optional[int] = None, + row_timeout_seconds: Optional[float] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] | list[str] = (), @@ -2218,6 +2257,8 @@ def udf( pip: tuple[str, ...] | list[str] = (), env: Optional[Mapping[str, str]] = None, python_version: Optional[str] = None, + batch_size: Optional[int] = None, + row_timeout_seconds: Optional[float] = None, gpu: bool = False, gpu_class: Optional[str] = None, conda: tuple[str, ...] | list[str] = (), @@ -2339,6 +2380,15 @@ def udf( values that differ between bindings. python_version : str, optional Remote Python major/minor version. Defaults to the client version. + batch_size : int, optional + Maximum number of input rows supplied in one execution batch. Must be + between 1 and 4294967295. Remote execution may use smaller batches. + When omitted, the service chooses the batch size adaptively. + row_timeout_seconds : float, optional + Positive, finite deadline in seconds for each Python call: one row for + scalar Functions and one batch for vectorized Functions. This does not + govern class initialization. When omitted, the service applies its + default per-call budget. gpu : bool, default False Whether every remote execution requires a GPU. The execution platform selects one compatible GPU for each worker. The requirement is part of @@ -2457,6 +2507,8 @@ def udf( pip=tuple(pip), env={} if env is None else env, python_version=python_version, + batch_size=batch_size, + row_timeout_seconds=row_timeout_seconds, gpu=gpu, gpu_class=gpu_class, conda=tuple(conda), diff --git a/python/python/tests/test_e2e_remote_db.py b/python/python/tests/test_e2e_remote_db.py index 74d5cd7de..5a2f1e290 100644 --- a/python/python/tests/test_e2e_remote_db.py +++ b/python/python/tests/test_e2e_remote_db.py @@ -10,7 +10,7 @@ import lancedb import numpy as np import pyarrow as pa import pytest -from lancedb import LanceDBConnection +from lancedb import LanceDBConnection, udf from lancedb.index import FTS # TODO: setup integ test mark and script @@ -74,6 +74,106 @@ def test_json_fts_against_server(): conn.drop_table(name) +@pytest.mark.skipif( + not os.environ.get("LANCEDB_URI"), + reason="Set LANCEDB_URI to a server with Python Functions enabled", +) +def test_function_execution_controls_through_sql_and_materialized_view(): + """Exercise registered batch controls through the real remote runtime. + + The Function returns the input batch length for every row. A cap of two must + therefore produce only one- or two-element observations in both a direct + SELECT and a materialized view refreshed after an append. + """ + suffix = uuid4().hex[:12] + conn = lancedb.connect( + uri=os.environ["LANCEDB_URI"], + api_key=os.environ.get("LANCEDB_API_KEY"), + region=os.environ.get("LANCEDB_REGION", "us-east-1"), + host_override=os.environ.get("LANCEDB_HOST_OVERRIDE"), + sql_host_override=os.environ.get("LANCEDB_SQL_HOST_OVERRIDE"), + ) + function_name = f"batch_probe_{suffix}" + source_name = f"batch_probe_source_{suffix}" + view_name = f"batch_probe_view_{suffix}" + registered_version = None + + @udf( + name=function_name, + input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + output_schema=pa.int64(), + batch_size=2, + row_timeout_seconds=15, + ) + def observed_batch(value: pa.Array) -> pa.Array: + size = len(value) + return pa.array([size] * size, type=pa.int64()) + + def row_batch_sizes(table_name): + rows = ( + conn.execute_query( + f"SELECT {function_name}(value) AS observed FROM {table_name}" + ) + .read_all() + .to_pylist() + ) + return [row["observed"] for row in rows] + + def drop_table_if_present(table_name): + try: + conn.drop_table(table_name) + except Exception: + pass + + try: + table = conn.create_table( + source_name, + pa.table( + {"value": pa.array(range(8), type=pa.int64())}, + schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + ), + ) + registered_version = conn.create_function(observed_batch, namespace_path=None) + + direct = row_batch_sizes(source_name) + assert len(direct) == 8 + assert all(1 <= size <= 2 for size in direct) + + view = conn.create_materialized_view( + view_name, + source=source_name, + select=[("observed", f"{function_name}(value)")], + ) + initial_view = [ + row["observed"] + for row in view.table.search().select(["observed"]).to_list() + ] + assert len(initial_view) == 8 + assert all(1 <= size <= 2 for size in initial_view) + + table.add( + pa.table( + {"value": pa.array(range(8, 12), type=pa.int64())}, + schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + ) + ) + view.refresh() + refreshed = [ + row["observed"] + for row in view.table.search().select(["observed"]).to_list() + ] + assert len(refreshed) == 12 + assert all(1 <= size <= 2 for size in refreshed) + finally: + drop_table_if_present(view_name) + drop_table_if_present(source_name) + if registered_version is not None: + try: + conn.drop_function(function_name, version=registered_version.version) + except Exception: + pass + + @pytest.mark.skip(reason="Need to set up a local server") def test_against_local_server(): conn = LanceDBConnection("lancedb+http://localhost:10024") diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 4478fd3e5..535df89b4 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -92,6 +92,57 @@ def test_scalar_udf_matches_shared_registration_golden_and_remains_callable(): } +def test_udf_execution_controls_match_shared_registration_golden(): + controlled = udf( + pip=["numpy>=2"], + env={"MODE": "test"}, + python_version="3.12", + batch_size=2, + row_timeout_seconds=0.75, + )(normalize_score.__wrapped__) + assert ( + controlled.registration_request.to_canonical_json() + == (FIXTURES / "remote_function_controlled_registration_request.canonical.json") + .read_text() + .strip() + ) + + +def test_udf_execution_controls_apply_to_vectorized_functions(): + @udf( + input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + output_schema=pa.int64(), + batch_size=3, + row_timeout_seconds=1.25, + ) + def identity(value: pa.Array) -> pa.Array: + return value + + request = json.loads(identity.registration_request.to_canonical_json()) + assert request["batch_size"] == 3 + assert request["row_timeout_seconds"] == 1.25 + + +@pytest.mark.parametrize( + ("options", "field"), + [ + ({"batch_size": 0}, "batch_size"), + ({"batch_size": 2**32}, "batch_size"), + ({"batch_size": True}, "batch_size"), + ({"row_timeout_seconds": 0}, "row_timeout_seconds"), + ({"row_timeout_seconds": -0.25}, "row_timeout_seconds"), + ({"row_timeout_seconds": float("inf")}, "row_timeout_seconds"), + ({"row_timeout_seconds": True}, "row_timeout_seconds"), + ], +) +def test_udf_execution_controls_require_their_positive_ranges(options, field): + def identity(value: int) -> int: + return value + + with pytest.raises(ValueError, match=field): + udf(**options)(identity) + + @pytest.mark.parametrize( "name", ["examples.image_info", "image-info", "image_info.", "image_info-"] ) diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 988fd40af..1f6c4faaf 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -663,17 +663,67 @@ pub struct FunctionArtifactRequest { pub adapter: PythonAdapterSpec, } +fn deserialize_batch_size<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer) + .map_err(|error| de::Error::custom(format!("batch_size: {error}")))?; + validate_batch_size_value(value).map_err(de::Error::custom)?; + Ok(value) +} + +fn deserialize_row_timeout<'de, D>(deserializer: D) -> std::result::Result, D::Error> +where + D: Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer) + .map_err(|error| de::Error::custom(format!("row_timeout_seconds: {error}")))?; + validate_row_timeout_value(value).map_err(de::Error::custom)?; + Ok(value) +} + +fn validate_batch_size_value(value: Option) -> std::result::Result<(), &'static str> { + if value == Some(0) { + return Err("batch_size must be between 1 and u32::MAX"); + } + Ok(()) +} + +fn validate_row_timeout_value(value: Option) -> std::result::Result<(), &'static str> { + if value.is_some_and(|seconds| !seconds.is_finite() || seconds <= 0.0) { + return Err("row_timeout_seconds must be positive and finite"); + } + Ok(()) +} + /// Stable request envelope for remote immutable Function registration. /// /// Credential values deliberately have no field here. The only secret-shaped /// thing a client sends is `secret_bindings`: the name of a Secret the /// database already holds, which Sophon resolves inside the remote runtime. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct FunctionRegistrationRequest { pub name: String, pub artifact: FunctionArtifactRequest, pub signature: FunctionSignature, pub runtime: PythonRuntimeSpec, + /// Maximum rows in one execution batch; the executor may use fewer. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_batch_size" + )] + pub batch_size: Option, + /// Deadline for each Python call, in seconds: one row for scalar Functions + /// and one batch for vectorized Functions. It does not govern Function + /// initialization. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_row_timeout" + )] + pub row_timeout_seconds: Option, /// Declared environment variable name to the Secret it binds. A binding is /// a reference: whether the Secret exists is answered when a column is /// declared against this version, not here. @@ -681,7 +731,28 @@ pub struct FunctionRegistrationRequest { pub secret_bindings: Vec, } -impl_json!(FunctionRegistrationRequest); +impl FunctionRegistrationRequest { + /// Decode a registration request supplied by an SDK language binding. + pub fn from_json(json: &str) -> Result { + from_json(json) + } + + /// Encode the known client contract with bytewise-sorted JSON keys. + pub fn to_canonical_json(&self) -> Result { + self.validate()?; + canonical_json(self) + } + + /// Reject explicit execution controls that cannot be submitted safely. + pub fn validate(&self) -> Result<()> { + let invalid_input = |message: &'static str| Error::InvalidInput { + message: message.to_string(), + }; + validate_batch_size_value(self.batch_size).map_err(invalid_input)?; + validate_row_timeout_value(self.row_timeout_seconds).map_err(invalid_input)?; + Ok(()) + } +} /// Exact FunctionVersion reference embedded in applications and bindings. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/rust/lancedb/src/remote/sql/function.rs b/rust/lancedb/src/remote/sql/function.rs index 671f3e3a4..7682b11c1 100644 --- a/rust/lancedb/src/remote/sql/function.rs +++ b/rust/lancedb/src/remote/sql/function.rs @@ -22,6 +22,7 @@ impl SqlClient { request: FunctionRegistrationRequest, namespace: &[String], ) -> Result> { + request.validate()?; validate_function_name(&request.name)?; validate_namespace(namespace)?; let name = quote_identifier(&request.name); diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 02efc96ab..b2973c679 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -2300,6 +2300,8 @@ async fn function_creation_and_description_use_sql_query_handles() { ); for defer in [false, true] { let mut request_json: serde_json::Value = serde_json::from_str(REQUEST).unwrap(); + request_json["batch_size"] = serde_json::json!(2); + request_json["row_timeout_seconds"] = serde_json::json!(0.75); request_json["secret_bindings"] = serde_json::json!([{ "kind": "env", "variable": "HF_TOKEN", @@ -2358,6 +2360,37 @@ async fn function_creation_and_description_use_sql_query_handles() { shutdown.send(()).unwrap(); } } + +#[tokio::test] +async fn function_creation_rejects_invalid_execution_controls_before_sql() { + const REQUEST: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.json" + ); + let service = TestSqlService::default(); + let statements = service.statements.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()); + }) + .with_sql_client(client.clone()); + + for timeout in [f64::NAN, f64::INFINITY, -1.0] { + let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap(); + request.row_timeout_seconds = Some(timeout); + let error = db.create_function_async(request, &[]).await.unwrap_err(); + assert!( + matches!(error, Error::InvalidInput { message } if message.contains("row_timeout_seconds")) + ); + } + + let mut request = crate::function::FunctionRegistrationRequest::from_json(REQUEST).unwrap(); + request.batch_size = Some(0); + let error = db.create_function_async(request, &[]).await.unwrap_err(); + assert!(matches!(error, Error::InvalidInput { message } if message.contains("batch_size"))); + assert!(statements.lock().unwrap().is_empty()); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn function_drop_keeps_the_query_id_and_reports_catalog_removal() { for (status, expected) in [("dropped", true), ("not_found", false)] { diff --git a/rust/lancedb/tests/first_class_function_slice2.rs b/rust/lancedb/tests/first_class_function_slice2.rs index 27d16323d..815ddcf08 100644 --- a/rust/lancedb/tests/first_class_function_slice2.rs +++ b/rust/lancedb/tests/first_class_function_slice2.rs @@ -58,6 +58,57 @@ fn registration_request_matches_shared_canonical_golden() { assert_no_secret_values(&value); } +#[test] +fn execution_controls_match_shared_canonical_golden() { + let request = FunctionRegistrationRequest::from_json(&fixture( + "remote_function_controlled_registration_request.json", + )) + .expect("controlled registration request"); + assert_eq!( + request.to_canonical_json().expect("canonical request"), + fixture("remote_function_controlled_registration_request.canonical.json").trim() + ); + + let value: Value = serde_json::from_str(&request.to_canonical_json().unwrap()).unwrap(); + assert_eq!(value["batch_size"], 2); + assert_eq!(value["row_timeout_seconds"], 0.75); +} + +#[test] +fn execution_controls_require_their_positive_ranges() { + let source = fixture("remote_function_registration_request.json"); + for (field, value) in [ + ("batch_size", Value::from(0)), + ("batch_size", Value::from(4_294_967_296_u64)), + ("row_timeout_seconds", Value::from(0)), + ("row_timeout_seconds", Value::from(-0.25)), + ] { + let mut request: Value = serde_json::from_str(&source).unwrap(); + request[field] = value; + let error = FunctionRegistrationRequest::from_json(&request.to_string()).unwrap_err(); + assert!( + error.to_string().contains(field), + "{field} validation error was {error}" + ); + } +} + +#[test] +fn execution_controls_are_validated_before_canonical_serialization() { + let source = fixture("remote_function_controlled_registration_request.json"); + for timeout in [f64::NAN, f64::INFINITY, -0.25] { + let mut request = FunctionRegistrationRequest::from_json(&source).unwrap(); + request.row_timeout_seconds = Some(timeout); + let error = request.to_canonical_json().unwrap_err(); + assert!(error.to_string().contains("row_timeout_seconds"), "{error}"); + } + + let mut request = FunctionRegistrationRequest::from_json(&source).unwrap(); + request.batch_size = Some(0); + let error = request.to_canonical_json().unwrap_err(); + assert!(error.to_string().contains("batch_size"), "{error}"); +} + /// The same shared golden as the Python suite builds from `@udf(secrets=...)` /// plus `bind_secrets`, so both clients agree byte for byte on a bound request. #[test] diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.canonical.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.canonical.json new file mode 100644 index 000000000..9ce9a289f --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.canonical.json @@ -0,0 +1 @@ +{"artifact":{"adapter":{"kind":"scalar_to_arrow_batch","version":1},"content":{"data":"ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK","encoding":"base64"},"digest":"sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f","entrypoint":"normalize_score","kind":"python_callable"},"batch_size":2,"name":"normalize_score","row_timeout_seconds":0.75,"runtime":{"env":{"MODE":"test"},"environment":{"kind":"pip","packages":["numpy>=2"]},"kind":"python","python_version":"3.12"},"signature":{"inputs":[{"arrow_type":"float64","name":"value","nullable":false}],"output":{"arrow_type":"float64","kind":"scalar","nullable":false}}} diff --git a/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.json b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.json new file mode 100644 index 000000000..cb0961eba --- /dev/null +++ b/rust/lancedb/tests/fixtures/first_class_functions/v1/remote_function_controlled_registration_request.json @@ -0,0 +1,45 @@ +{ + "name": "normalize_score", + "batch_size": 2, + "row_timeout_seconds": 0.75, + "artifact": { + "kind": "python_callable", + "digest": "sha256:760784bdcef57b802f389b97804cc0b618aae39e86044733451bcd0b13089a7f", + "entrypoint": "normalize_score", + "content": { + "encoding": "base64", + "data": "ZnJvbSBfX2Z1dHVyZV9fIGltcG9ydCBhbm5vdGF0aW9ucwoKZGVmIG5vcm1hbGl6ZV9zY29yZSh2YWx1ZTogZmxvYXQpIC0+IGZsb2F0OgogICAgcmV0dXJuIHZhbHVlIC8gMTAwLjAK" + }, + "adapter": { + "kind": "scalar_to_arrow_batch", + "version": 1 + } + }, + "signature": { + "inputs": [ + { + "name": "value", + "arrow_type": "float64", + "nullable": false + } + ], + "output": { + "kind": "scalar", + "arrow_type": "float64", + "nullable": false + } + }, + "runtime": { + "kind": "python", + "python_version": "3.12", + "environment": { + "kind": "pip", + "packages": [ + "numpy>=2" + ] + }, + "env": { + "MODE": "test" + } + } +}