rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-09-04 12:38:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: support nullable named function outputs (#4123)

Browse Source 
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.
This commit is contained in:
Jack Ye
2026-09-03 22:23:53 -07:00
committed by GitHub
parent e639b1b650
commit aab23eb39e
6 changed files with 270 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/src/python/python.md
+2
View File
@@ -147,6 +147,8 @@ listing a storage directory.
::: lancedb.functions.OutputMapping
::: lancedb.functions.AssignmentMapping
::: lancedb.functions.FunctionBinding
::: lancedb.functions.RefreshColumnResult
python/python/lancedb/__init__.py
+1
View File
@@ -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,
python/python/lancedb/functions.py
+16 -12
View File
@@ -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",
python/python/tests/test_first_class_function_slice2.py
+37
View File
@@ -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 = (
rust/lancedb/src/function.rs
+16 -2
View File
@@ -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<InputBinding>,
outputs: Vec<OutputMapping>,
#[serde(default, skip_serializing_if = "Option::is_none")]
assignment: Option<AssignmentMapping>,
/// 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()
}
rust/lancedb/src/table/computed_columns.rs
+198 -19
View File
@@ -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::<u32>().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::<Result<Vec<_>>>()?;
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)]);