mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(python): expose Function batch size and row timeout controls (#4807)
Python UDF registration could not set `batch_size` or `row_timeout_seconds`, and the Rust registration model discarded these fields before submitting SQL. Expose both optional controls and preserve them through registration, with invalid values rejected and unchanged canonical bytes when the controls are omitted. The timeout applies to each Python call: one row for scalar Functions and one batch for vectorized Functions. Initialization keeps its existing service deadline. Validation at `f58e55c535a7df310db14d691b329402ca4baaf5`: 6 Rust registration-contract tests, 2 outbound-SQL tests, and all 103 Python slice2 tests passed after rebuilding the native extension. Rust formatting and the repository Python CI checks (Ruff 0.9.9, python/ scope) passed. A real REST + Flight SQL server integration also passed using the rebuilt SDK: `batch_size=2` limits observed calls in a direct SELECT over 8 rows, initial materialized-view creation, and append/refresh to 12 rows. This CPU host test verifies SDK-to-runtime passthrough; it does not exercise GPU initialization. Refs ENT-3246. The original GPU initialization and video-decoding scenario still needs verification on a current deployment; this SDK change does not close that investigation. The additional repository-root style checks report existing docs/web-tests findings (nine format files and 48 lint diagnostics); all affected files are unchanged from main. These are outside the Python CI job scope. Rust compatibility: callers constructing `FunctionRegistrationRequest` must initialize the two new optional fields (`None` preserves defaults); the request no longer implements `Eq`. The final head `798f463c9f68ab73e28c9a1ff06b5e4ee3e5d01f` adds creation-guide documentation and integrates main’s namespace repr/optional OCI embedding changes. The Function registration implementation is unchanged. All 103 Python slice2 tests passed again on EC2 using the unchanged native extension; Rust and server-bridge evidence remains applicable.
This commit is contained in:
Xuanwo
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
<CodeGroup>
|
||||
```python Python icon="python"
|
||||
@udf(batch_size=32, row_timeout_seconds=90)
|
||||
def encode(value: float) -> float:
|
||||
return value
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
`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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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-"]
|
||||
)
|
||||
|
||||
@@ -663,17 +663,67 @@ pub struct FunctionArtifactRequest {
|
||||
pub adapter: PythonAdapterSpec,
|
||||
}
|
||||
|
||||
fn deserialize_batch_size<'de, D>(deserializer: D) -> std::result::Result<Option<u32>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<u32>::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<Option<f64>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
let value = Option::<f64>::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<u32>) -> 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<f64>) -> 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<u32>,
|
||||
/// 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<f64>,
|
||||
/// 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<SecretBinding>,
|
||||
}
|
||||
|
||||
impl_json!(FunctionRegistrationRequest);
|
||||
impl FunctionRegistrationRequest {
|
||||
/// Decode a registration request supplied by an SDK language binding.
|
||||
pub fn from_json(json: &str) -> Result<Self> {
|
||||
from_json(json)
|
||||
}
|
||||
|
||||
/// Encode the known client contract with bytewise-sorted JSON keys.
|
||||
pub fn to_canonical_json(&self) -> Result<String> {
|
||||
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)]
|
||||
|
||||
@@ -22,6 +22,7 @@ impl SqlClient {
|
||||
request: FunctionRegistrationRequest,
|
||||
namespace: &[String],
|
||||
) -> Result<Job<FunctionVersion>> {
|
||||
request.validate()?;
|
||||
validate_function_name(&request.name)?;
|
||||
validate_namespace(namespace)?;
|
||||
let name = quote_identifier(&request.name);
|
||||
|
||||
@@ -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<String> {
|
||||
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)] {
|
||||
|
||||
@@ -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]
|
||||
|
||||
+1
@@ -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}}}
|
||||
+45
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user