From aab23eb39ee4e1a2d5898484f35a863d4c8e7036 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 3 Sep 2026 22:23:53 -0700 Subject: [PATCH] feat: support nullable named function outputs (#4123) Supports fully nullable named Function outputs while preserving the distinction between a valid all-null struct and a null/unassigned result. ## Concrete example This UDF contract is now valid: ```python @udf( input_schema=pa.schema([ pa.field("text", pa.string(), nullable=False), ]), output_schema=pa.schema([ pa.field( "embedding", pa.list_(pa.float32(), list_size=1024), nullable=True, ), pa.field("embedding_failure_reason", pa.string(), nullable=True), pa.field("embedding_failure_code", pa.int32(), nullable=True), ]), ) def embed(text): ... ``` A successful row can return: ```text embedding = [0.12, ...] embedding_failure_reason = NULL embedding_failure_code = NULL ``` If remote inference still fails after retries, it can return: ```text embedding = NULL embedding_failure_reason = "HTTP 429: rate limited" embedding_failure_code = 429 ``` An all-null but valid result struct is also assigned; it is not mistaken for unfinished work. ## Binding shapes - Mapping the result to one output column stores the `StructArray` directly, including its parent validity bitmap. - Flattening the result into top-level columns stores the parent validity in a reserved internal nullable Boolean assignment column that is not part of the UDF result mapping. - An outer null struct remains unassigned/skipped. A valid struct remains assigned regardless of which child fields are null. - Scalar Function outputs remain non-nullable. The contract is preserved through Python registration, Rust application planning, persisted `FunctionBinding` metadata, schema revalidation, and Enterprise execution. --- docs/src/python/python.md | 2 + python/python/lancedb/__init__.py | 1 + python/python/lancedb/functions.py | 28 ++- .../tests/test_first_class_function_slice2.py | 37 +++ rust/lancedb/src/function.rs | 18 +- rust/lancedb/src/table/computed_columns.rs | 217 ++++++++++++++++-- 6 files changed, 270 insertions(+), 33 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index f2bc72b3a..b0d1bb426 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -147,6 +147,8 @@ listing a storage directory. ::: lancedb.functions.OutputMapping +::: lancedb.functions.AssignmentMapping + ::: lancedb.functions.FunctionBinding ::: lancedb.functions.RefreshColumnResult diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index af325a29b..cb9b57be3 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -26,6 +26,7 @@ from .sql import AsyncQuery as AsyncSqlQuery from .sql import Query as SqlQuery from .sql import QueryDescription from .functions import ( + AssignmentMapping as AssignmentMapping, FunctionArtifactRequest as FunctionArtifactRequest, FunctionApplication as FunctionApplication, FunctionBinding as FunctionBinding, diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 2815f7f17..d3fa6da81 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -470,11 +470,7 @@ class InputBinding(_RemoteValue): class OutputMapping(_RemoteValue): - """One stable result-field mapping. - - Assignment state is outside the Slice 1 client contract. During the NULL - transition Lance exposes no public cell-flag identifier to persist here. - """ + """One stable result-field mapping.""" result_field: str output_name: str @@ -484,6 +480,13 @@ class OutputMapping(_RemoteValue): nullable: bool +class AssignmentMapping(_RemoteValue): + """Internal physical column preserving flattened struct validity.""" + + output_name: str + output_field_id: _Int32 + + class FunctionBinding(_RemoteValue): """Immutable Function binding persisted by the Enterprise table service.""" @@ -491,6 +494,7 @@ class FunctionBinding(_RemoteValue): function: FunctionVersionRef inputs: tuple[InputBinding, ...] outputs: tuple[OutputMapping, ...] + assignment: Optional[AssignmentMapping] = None input_schema: Optional[Mapping[str, Any]] = None output_schema: Optional[Mapping[str, Any]] = None @@ -911,8 +915,6 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp if not fields: raise ValueError("named-struct Function output must contain at least one field") - if any(field.nullable for field in fields): - raise ValueError("Function output fields must be non-nullable") for field in fields: _validate_exact_arrow_field(field) names = [field.name for field in fields] @@ -924,7 +926,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp FunctionResultField( name=field.name, arrow_type=_canonical_arrow_field(field), - nullable=False, + nullable=field.nullable, ) for field in fields ), @@ -1307,8 +1309,9 @@ def udf( Input and output signatures are inferred from supported annotations. For Arrow types annotations cannot express precisely, pass ``input_schema`` - and ``output_schema`` together. Nullable outputs are rejected because V1 - uses physical NULL to represent unassigned computed-column rows. + and ``output_schema`` together. Scalar outputs must be non-nullable. Every + named-struct field may be nullable; Enterprise preserves the struct's + validity when the result is expanded into sibling columns. Parameters ---------- @@ -1320,8 +1323,8 @@ def udf( Explicit input fields in the exact order of the callable parameters. Must be provided together with ``output_schema``. output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional - Explicit scalar or named-struct output. Must be non-nullable and be - provided together with ``input_schema``. + Explicit scalar or named-struct output. Scalar outputs must be + non-nullable. Must be provided together with ``input_schema``. pip : sequence of str, optional Pip requirements for the remote environment. conda : sequence of str, optional @@ -1387,6 +1390,7 @@ def udf( __all__ = [ + "AssignmentMapping", "ApplicationInput", "FunctionApplication", "FunctionArtifact", diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index d8a9aeebf..4816019d6 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -850,6 +850,43 @@ def test_named_struct_function_can_include_a_blob_result_field(): ] +def test_named_struct_function_preserves_nullable_result_fields(): + @udf( + input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + output_schema=pa.schema( + [ + pa.field("result", pa.int64(), nullable=True), + pa.field("failure_code", pa.int32(), nullable=False), + ] + ), + ) + def nullable_result(value): + return {"result": value, "failure_code": 0} + + output = nullable_result.registration_request.signature.output + assert [(field.name, field.nullable) for field in output.fields] == [ + ("result", True), + ("failure_code", False), + ] + + @udf( + input_schema=pa.schema([pa.field("value", pa.int64(), nullable=False)]), + output_schema=pa.schema( + [ + pa.field("result", pa.int64(), nullable=True), + pa.field("failure_code", pa.int32(), nullable=True), + ] + ), + ) + def all_nullable(value): + return {"result": value, "failure_code": None} + + assert all( + field.nullable + for field in all_nullable.registration_request.signature.output.fields + ) + + def test_metadata_marked_blob_field_uses_the_semantic_type(): extension = lancedb.blob("image", nullable=False).type storage = ( diff --git a/rust/lancedb/src/function.rs b/rust/lancedb/src/function.rs index 79693c031..e67d763ac 100644 --- a/rust/lancedb/src/function.rs +++ b/rust/lancedb/src/function.rs @@ -582,8 +582,8 @@ pub struct InputBinding { /// Ordered result-field to table-field mapping for a Function binding. /// -/// Assignment state is not part of the Slice 1 client contract. During the -/// NULL transition there is no public Lance cell-flag identifier to persist. +/// `nullable` describes the logical Function result. Physical computed-column +/// fields remain nullable while unassigned. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct OutputMapping { pub result_field: String, @@ -594,6 +594,14 @@ pub struct OutputMapping { pub nullable: bool, } +/// Internal physical column preserving the parent validity of a flattened +/// named-struct result. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AssignmentMapping { + pub output_name: String, + pub output_field_id: i32, +} + /// Immutable Function binding persisted by the Enterprise table service. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct FunctionBinding { @@ -601,6 +609,8 @@ pub struct FunctionBinding { function: FunctionVersionRef, inputs: Vec, outputs: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + assignment: Option, /// Exact Arrow schema presented to the Function, encoded with the Lance /// Namespace Arrow JSON representation. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -627,6 +637,10 @@ impl FunctionBinding { &self.outputs } + pub fn assignment(&self) -> Option<&AssignmentMapping> { + self.assignment.as_ref() + } + pub fn input_schema(&self) -> Option<&Value> { self.input_schema.as_ref() } diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 085876a81..0fcc5485d 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -60,6 +60,10 @@ pub const FUNCTION_BINDING_ID_META_KEY: &str = "computed_column.function.binding /// Field metadata key holding this sibling's ordered Function output ordinal. pub const FUNCTION_OUTPUT_ORDINAL_META_KEY: &str = "computed_column.function.output_ordinal"; +/// Reserved Function output ordinal for an internal flattened-result +/// assignment column. +pub const FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL: u32 = u32::MAX; + /// Schema metadata key holding all immutable Function bindings. pub const FUNCTION_BINDINGS_META_KEY: &str = "lancedb::function_bindings"; @@ -312,22 +316,29 @@ pub(crate) fn ensure_supported_function_metadata(schema: &ArrowSchema) -> Result binding_id ), })?; - let output = binding - .outputs() - .get(output_ordinal as usize) - .ok_or_else(|| Error::InvalidInput { - message: format!( - "Function output '{}' has invalid ordinal {}", - field.name(), - output_ordinal - ), - })?; - if output.output_name != field.name().as_str() { + let destination = if output_ordinal == FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL { + binding + .assignment() + .map(|assignment| assignment.output_name.as_str()) + } else { + binding + .outputs() + .get(output_ordinal as usize) + .map(|output| output.output_name.as_str()) + } + .ok_or_else(|| Error::InvalidInput { + message: format!( + "Function output '{}' has invalid ordinal {}", + field.name(), + output_ordinal + ), + })?; + if destination != field.name().as_str() { return Err(Error::InvalidInput { message: format!( "Function output '{}' does not match binding destination '{}'", field.name(), - output.output_name + destination ), }); } @@ -498,6 +509,7 @@ fn ensure_known_binding_shape(value: &Value) -> Result<()> { "function", "inputs", "outputs", + "assignment", "input_schema", "output_schema", ], @@ -546,6 +558,13 @@ fn ensure_known_binding_shape(value: &Value) -> Result<()> { "output mapping", )?; } + if let Some(assignment) = object.get("assignment") { + reject_unknown_object_fields( + assignment, + &["output_name", "output_field_id"], + "assignment mapping", + )?; + } Ok(()) } @@ -865,7 +884,7 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding output.output_name )) })?; - if field.name() != &output.output_name || !field.is_nullable() || output.nullable { + if field.name() != &output.output_name || !field.is_nullable() { return Err(invalid_function(format!( "Function output '{}' no longer matches binding '{}'", output.output_name, @@ -926,6 +945,61 @@ fn ensure_binding_matches_schema(schema: &ArrowSchema, binding: &FunctionBinding output_fields.push(json.fields.into_iter().next().unwrap()); } } + if let Some(assignment) = binding.assignment() { + if binding + .outputs() + .iter() + .any(|output| output.result_field == WHOLE_RESULT_FIELD) + { + return Err(invalid_function(format!( + "Function binding '{}' cannot attach an assignment column to a whole result", + binding.binding_id() + ))); + } + let field = schema + .field_with_name(&assignment.output_name) + .map_err(|_| { + invalid_function(format!( + "Function binding '{}' assignment column '{}' is missing", + binding.binding_id(), + assignment.output_name + )) + })?; + let metadata = field.metadata(); + if field.data_type() != &DataType::Boolean + || !field.is_nullable() + || metadata.get(COMPUTED_COLUMN_META_KEY).map(String::as_str) != Some("true") + || metadata.get(KIND_META_KEY).map(String::as_str) != Some(FUNCTION_KIND) + || metadata + .get(FUNCTION_BINDING_ID_META_KEY) + .map(String::as_str) + != Some(binding.binding_id()) + || metadata + .get(FUNCTION_OUTPUT_ORDINAL_META_KEY) + .and_then(|value| value.parse::().ok()) + != Some(FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL) + { + return Err(invalid_function(format!( + "Function binding '{}' assignment column no longer matches its declaration", + binding.binding_id() + ))); + } + let json = lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ + ArrowField::new(assignment.output_name.clone(), DataType::Boolean, true), + ])) + .map_err(|e| invalid_function(format!("invalid Function assignment schema: {e}")))?; + output_fields.push(json.fields.into_iter().next().unwrap()); + } else if binding.outputs().iter().all(|output| output.nullable) + && binding + .outputs() + .iter() + .all(|output| output.result_field != WHOLE_RESULT_FIELD) + { + return Err(invalid_function(format!( + "Function binding '{}' has no flattened-result assignment column", + binding.binding_id() + ))); + } let output_schema = JsonArrowSchema::new(output_fields); let output_schema = serde_json::to_value(output_schema).map_err(|e| { invalid_function(format!( @@ -1081,11 +1155,6 @@ pub(crate) fn plan_function_application( "named-struct Function result field names must be unique", )); } - if output.fields.iter().any(|field| field.nullable) { - return Err(invalid_function( - "Function logical outputs must be non-nullable during NULL assignment", - )); - } let unknown = application .columns() .keys() @@ -1107,7 +1176,9 @@ pub(crate) fn plan_function_application( let fields = output .fields .iter() - .map(|field| function_output_field(&field.name, false, &field.arrow_type)) + .map(|field| { + function_output_field(&field.name, field.nullable, &field.arrow_type) + }) .collect::>>()?; let mut data_type = JsonArrowDataType::new("struct".to_string()); data_type.fields = Some(fields); @@ -2858,6 +2929,17 @@ mod tests { &inputs, )); } + if let Some(assignment) = binding.assignment() { + fields.push( + ArrowField::new(&assignment.output_name, DataType::Boolean, true).with_metadata( + function_computed_column_metadata( + binding.binding_id(), + FUNCTION_ASSIGNMENT_OUTPUT_ORDINAL, + &inputs, + ), + ), + ); + } ArrowSchema::new(fields) } @@ -2875,6 +2957,74 @@ mod tests { .unwrap(); } + #[test] + fn test_binding_preserves_all_nullable_outputs_with_an_assignment_column() { + let mut raw_binding: Value = serde_json::from_str(include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_binding.json" + )) + .unwrap(); + raw_binding["outputs"][0]["nullable"] = Value::Bool(true); + raw_binding["outputs"][1]["nullable"] = Value::Bool(true); + let without_assignment: FunctionBinding = + serde_json::from_value(raw_binding.clone()).unwrap(); + let error = ensure_binding_matches_schema( + &valid_function_binding_schema(true, true, &without_assignment), + &without_assignment, + ) + .unwrap_err(); + assert!( + error + .to_string() + .contains("flattened-result assignment column") + ); + + raw_binding["assignment"] = serde_json::json!({ + "output_name": "__function_assignment_fb_01K3TEXT", + "output_field_id": -1, + }); + raw_binding["output_schema"]["fields"] + .as_array_mut() + .unwrap() + .push(serde_json::json!({ + "name": "__function_assignment_fb_01K3TEXT", + "nullable": true, + "type": {"type": "bool"}, + })); + let binding: FunctionBinding = serde_json::from_value(raw_binding).unwrap(); + ensure_binding_matches_schema( + &valid_function_binding_schema(true, true, &binding), + &binding, + ) + .unwrap(); + + let schema = ArrowSchema::new_with_metadata( + valid_function_binding_schema(true, true, &binding) + .fields() + .to_vec(), + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + function_bindings_metadata(std::slice::from_ref(&binding)).unwrap(), + )]), + ); + ensure_supported_function_metadata(&schema).unwrap(); + + let mut metadata: Value = + serde_json::from_str(schema.metadata().get(FUNCTION_BINDINGS_META_KEY).unwrap()) + .unwrap(); + metadata["bindings"][0]["assignment"]["future"] = Value::Bool(true); + let future_schema = ArrowSchema::new_with_metadata( + schema.fields().to_vec(), + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + serde_json::to_string(&metadata).unwrap(), + )]), + ); + assert!(matches!( + ensure_supported_function_metadata(&future_schema), + Err(Error::NotSupported { .. }) + )); + } + #[test] fn test_nullable_function_input_cannot_bind_to_non_nullable_parameter() { let mut raw_binding: Value = serde_json::from_str(include_str!( @@ -3094,6 +3244,35 @@ mod tests { ); } + #[test] + fn test_named_struct_plan_preserves_nullable_result_fields() { + let mut value = serde_json::to_value(named_struct_application("{}")).unwrap(); + value["output"]["fields"][0]["nullable"] = Value::Bool(true); + value["output"]["fields"][1]["nullable"] = Value::Bool(true); + let application = FunctionApplication::from_json(&value.to_string()).unwrap(); + + let expanded = + plan_function_application(&function_input_schema(), &application, None).unwrap(); + assert!( + expanded + .output_schema + .fields + .iter() + .all(|field| field.nullable) + ); + + let whole = + plan_function_application(&function_input_schema(), &application, Some("features")) + .unwrap(); + let fields = whole.output_schema.fields[0] + .r#type + .fields + .as_ref() + .unwrap(); + assert!(fields[0].nullable); + assert!(fields[1].nullable); + } + #[test] fn test_blob_function_plans_semantic_input_and_scalar_output() { let schema = ArrowSchema::new(vec![crate::blob("image", false)]);