From e6867f7d0433054b140fc7a3b2c67787029bf3d0 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 1 Sep 2026 23:51:08 +0800 Subject: [PATCH] feat: support nested blob function signatures (#4109) Function signatures currently reject Blob v2 fields nested inside structs, preventing UDFs from accepting or returning structured values that contain blobs. Accept canonical Blob v2 fields as direct or recursive struct children while preserving exact field metadata and nullability. Blob fields under list, large-list, fixed-size-list, or map ancestors remain rejected because collection runtime adaptation is outside the supported Function ABI. A whole named struct result can bind directly to one destination column without introducing an extra wrapper level. --- python/python/lancedb/functions.py | 107 ++++++++- .../tests/test_first_class_function_slice2.py | 177 +++++++++++++-- rust/lancedb/src/table/computed_columns.rs | 206 ++++++++++++++++-- 3 files changed, 446 insertions(+), 44 deletions(-) diff --git a/python/python/lancedb/functions.py b/python/python/lancedb/functions.py index 3bdd117cb..2815f7f17 100644 --- a/python/python/lancedb/functions.py +++ b/python/python/lancedb/functions.py @@ -521,6 +521,12 @@ class RefreshColumnResult(_RemoteValue): _FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$") _FUNCTION_BLOB_V2_TYPE = "blob_v2" +_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name" +_BLOB_V2_EXTENSION_NAME = "lance.blob.v2" +_NESTED_BLOB_COLLECTION_ERROR = ( + "unsupported Arrow type for Function signature: Blob v2 fields nested under " + "collection types are not supported" +) _GRAMMAR_PRIMITIVES = ( @@ -591,6 +597,19 @@ def _validate_exact_arrow_field(field: pa.Field) -> None: "unsupported Arrow type for Function signature: lance.blob.v2 " f"requires a supported Blob storage layout, got {field}" ) + metadata = { + (key.decode() if isinstance(key, bytes) else key): ( + value.decode() if isinstance(value, bytes) else value + ) + for key, value in (field.metadata or {}).items() + } + if metadata and metadata != { + _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME + }: + raise TypeError( + "unsupported Arrow type for Function signature: lance.blob.v2 " + "field metadata must contain only its canonical extension marker" + ) elif field.metadata: raise TypeError( "unsupported Arrow type for Function signature: field metadata " @@ -655,23 +674,84 @@ def _canonical_arrow_field(field: pa.Field) -> str: return _canonical_arrow_type(field.type) -def _exact_arrow_field(field: pa.Field) -> dict[str, Any]: +def _blob_storage_type(field: pa.Field) -> pa.DataType: + data_type = field.type + if isinstance(data_type, pa.ExtensionType): + return data_type.storage_type + return data_type + + +def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]: + storage = _blob_storage_type(field) + if not pa.types.is_struct(storage): + raise TypeError( + "unsupported Arrow type for Function signature: lance.blob.v2 " + "requires struct storage" + ) + return { + "type": "struct", + "fields": [ + { + "name": child.name, + "nullable": child.nullable, + "type": ( + {"type": "large_binary"} + if pa.types.is_large_binary(child.type) + else _exact_arrow_type(child.type) + ), + } + for child in storage + ], + } + + +def _data_type_has_blob_v2(data_type: pa.DataType) -> bool: + if pa.types.is_struct(data_type): + return any( + _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + for field in data_type + ) + if ( + pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type) + ): + field = data_type.value_field + return _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + if pa.types.is_map(data_type): + return any( + _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type) + for field in (data_type.key_field, data_type.item_field) + ) + return False + + +def _exact_arrow_field( + field: pa.Field, *, inside_collection: bool = False +) -> dict[str, Any]: _validate_exact_arrow_field(field) if _is_blob_v2_field(field): - raise TypeError( - "unsupported Arrow type for Function signature: nested Blob v2 " - "fields are not supported; declare Blob parameters or named result " - "fields directly" - ) + if inside_collection: + raise TypeError(_NESTED_BLOB_COLLECTION_ERROR) + return { + "name": field.name, + "nullable": field.nullable, + "type": _exact_blob_storage_type(field), + "metadata": { + _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME, + }, + } value = { "name": field.name, "nullable": field.nullable, - "type": _exact_arrow_type(field.type), + "type": _exact_arrow_type(field.type, inside_collection=inside_collection), } return value -def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: +def _exact_arrow_type( + data_type: pa.DataType, *, inside_collection: bool = False +) -> dict[str, Any]: for candidate, name in _GRAMMAR_PRIMITIVES: if data_type == candidate: return {"type": name} @@ -685,7 +765,10 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: ) return { "type": "struct", - "fields": [_exact_arrow_field(field) for field in fields], + "fields": [ + _exact_arrow_field(field, inside_collection=inside_collection) + for field in fields + ], } if ( pa.types.is_list(data_type) @@ -710,11 +793,15 @@ def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]: if pa.types.is_large_list(data_type) else "fixed_size_list" ), - "fields": [_exact_arrow_field(data_type.value_field)], + "fields": [ + _exact_arrow_field(data_type.value_field, inside_collection=True) + ], } if pa.types.is_fixed_size_list(data_type): value["length"] = data_type.list_size return value + if pa.types.is_map(data_type) and _data_type_has_blob_v2(data_type): + raise TypeError(_NESTED_BLOB_COLLECTION_ERROR) raise TypeError(f"unsupported Arrow type for Function signature: {data_type}") diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index b1308a2bc..d8a9aeebf 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -668,6 +668,167 @@ def test_blob_fields_use_the_scalar_function_semantic_type(): assert signature.output.arrow_type == "blob_v2" +def test_whole_named_struct_function_can_include_a_blob_result_field(): + @udf( + input_schema=pa.schema([lancedb.blob("image", nullable=False)]), + output_schema=pa.field( + "payload", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + lancedb.blob("image", nullable=False), + ] + ), + nullable=False, + ), + ) + def inspect_blob(image): + return {"mime_type": "image/png", "image": image} + + output = inspect_blob.registration_request.signature.output + assert output.kind == "named_struct" + assert [(field.name, field.arrow_type) for field in output.fields] == [ + ("mime_type", "utf8"), + ("image", "blob_v2"), + ] + + +def test_struct_blob_signature_fields_preserve_exact_metadata_and_nullability(): + nested_input = pa.field( + "payload", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + pa.field( + "nested", + pa.struct([lancedb.blob("image", nullable=True)]), + nullable=True, + ), + ] + ), + nullable=True, + ) + nested_output = pa.field( + "result", + pa.struct( + [ + pa.field("mime_type", pa.string(), nullable=False), + pa.field( + "nested", + pa.struct([lancedb.blob("image", nullable=True)]), + nullable=False, + ), + ] + ), + nullable=False, + ) + + @udf(input_schema=pa.schema([nested_input]), output_schema=nested_output) + def copy_payload(payload): + return payload + + signature = copy_payload.registration_request.signature + input_type = json.loads(signature.inputs[0].arrow_type) + assert input_type["fields"][1]["nullable"] is True + input_blob = input_type["fields"][1]["type"]["fields"][0] + assert input_blob["nullable"] is True + assert input_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"} + + assert signature.output.kind == "named_struct" + nested_result = next( + field for field in signature.output.fields if field.name == "nested" + ) + output_type = json.loads(nested_result.arrow_type) + output_blob = output_type["fields"][0] + assert output_blob["nullable"] is True + assert output_blob["metadata"] == {"ARROW:extension:name": "lance.blob.v2"} + + +def test_struct_blob_signature_supports_multiple_struct_levels(): + recursive = pa.field( + "value", + pa.struct( + [ + pa.field( + "level_1", + pa.struct( + [ + pa.field( + "level_2", + pa.struct([lancedb.blob("image", nullable=False)]), + nullable=False, + ) + ] + ), + nullable=False, + ) + ] + ), + nullable=False, + ) + + @udf( + input_schema=pa.schema([recursive]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value["level_1"]["level_2"]["image"]) + + encoded = json.loads(blob_size.registration_request.signature.inputs[0].arrow_type) + blob = encoded["fields"][0]["type"]["fields"][0]["type"]["fields"][0] + assert blob["metadata"]["ARROW:extension:name"] == "lance.blob.v2" + + +@pytest.mark.parametrize( + "data_type", + [ + pa.list_(lancedb.blob("item", nullable=False)), + pa.large_list(lancedb.blob("item", nullable=False)), + pa.list_(lancedb.blob("item", nullable=False), 2), + pa.map_(pa.string(), lancedb.blob("value", nullable=False).type), + ], +) +def test_blob_signature_rejects_collection_ancestors(data_type): + with pytest.raises( + TypeError, + match="Blob v2 fields nested under collection types are not supported", + ): + + @udf( + input_schema=pa.schema([pa.field("value", data_type, nullable=False)]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value) + + +def test_blob_signature_rejects_collection_below_a_struct(): + nested = pa.field( + "value", + pa.struct( + [ + pa.field( + "images", + pa.list_(lancedb.blob("item", nullable=False)), + nullable=False, + ) + ] + ), + nullable=False, + ) + with pytest.raises( + TypeError, + match="Blob v2 fields nested under collection types are not supported", + ): + + @udf( + input_schema=pa.schema([nested]), + output_schema=pa.field("size", pa.int64(), nullable=False), + ) + def blob_size(value): + return len(value["images"]) + + def test_named_struct_function_can_include_a_blob_result_field(): @udf( input_schema=pa.schema([lancedb.blob("image", nullable=False)]), @@ -729,22 +890,6 @@ def test_blob_marker_rejects_invalid_storage_layout(): return len(image) -def test_nested_blob_signature_field_has_a_clear_error(): - nested = pa.field( - "value", - pa.struct([lancedb.blob("image", nullable=False)]), - nullable=False, - ) - with pytest.raises(TypeError, match="nested Blob v2 fields are not supported"): - - @udf( - input_schema=pa.schema([nested]), - output_schema=pa.field("size", pa.int64(), nullable=False), - ) - def blob_size(value): - return len(value["image"]) - - def test_nested_non_blob_extension_is_not_silently_unwrapped(): class TestExtension(pa.ExtensionType): def __init__(self): diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 77e3d0a4d..085876a81 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -589,16 +589,13 @@ fn canonical_input_arrow_type(field: &JsonArrowField) -> Result { .and_then(|metadata| metadata.get(ARROW_EXT_NAME_KEY)) .map(String::as_str) == Some(BLOB_V2_EXT_NAME); - if is_blob_v2 { + if is_blob_v2 || field.r#type.fields.is_some() { let arrow_field = lance_namespace::schema::convert_json_arrow_field(field) .map_err(|e| invalid_function(format!("invalid Function input field: {e}")))?; - if !has_supported_blob_v2_layout(&arrow_field) { - return Err(invalid_function(format!( - "Function input '{}' has an invalid Blob v2 storage layout", - arrow_field.name() - ))); + validate_function_blob_nesting(&arrow_field, false)?; + if is_blob_v2 { + return Ok(FUNCTION_BLOB_V2_TYPE.to_string()); } - return Ok(FUNCTION_BLOB_V2_TYPE.to_string()); } if field.r#type.fields.is_none() && field.r#type.length.is_none() { Ok(field.r#type.r#type.clone()) @@ -617,6 +614,34 @@ fn has_supported_blob_v2_layout(field: &ArrowField) -> bool { ) } +fn validate_function_blob_nesting(field: &ArrowField, inside_collection: bool) -> Result<()> { + if field.is_blob_v2() { + if inside_collection { + return Err(invalid_function(format!( + "Function field '{}' nests Blob v2 under a collection, which Function signatures do not support", + field.name() + ))); + } + if !has_supported_blob_v2_layout(field) { + return Err(invalid_function(format!( + "Function field '{}' has an invalid Blob v2 storage layout", + field.name() + ))); + } + return Ok(()); + } + match field.data_type() { + DataType::Struct(fields) => fields + .iter() + .try_for_each(|field| validate_function_blob_nesting(field, inside_collection)), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => validate_function_blob_nesting(field, true), + _ => Ok(()), + } +} + /// `fixed_size_list` -> (`item`, `size`); the comma must sit outside /// any nested `<...>`. fn split_fixed_size_list(raw: &str) -> Option<(&str, i32)> { @@ -697,21 +722,22 @@ fn parse_output_arrow_type(raw: &str) -> Result { } fn function_output_field(name: &str, nullable: bool, raw: &str) -> Result { - if raw == FUNCTION_BLOB_V2_TYPE { - return lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![ - crate::blob(name, nullable), - ])) + let field = if raw == FUNCTION_BLOB_V2_TYPE { + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![crate::blob( + name, nullable, + )])) .map_err(|e| invalid_function(format!("could not encode Blob v2 output field: {e}")))? .fields .into_iter() .next() - .ok_or_else(|| invalid_function("Blob v2 output field is missing")); - } - Ok(JsonArrowField::new( - name.to_string(), - nullable, - parse_output_arrow_type(raw)?, - )) + .ok_or_else(|| invalid_function("Blob v2 output field is missing"))? + } else { + JsonArrowField::new(name.to_string(), nullable, parse_output_arrow_type(raw)?) + }; + let arrow_field = lance_namespace::schema::convert_json_arrow_field(&field) + .map_err(|e| invalid_function(format!("invalid Function output field: {e}")))?; + validate_function_blob_nesting(&arrow_field, false)?; + Ok(field) } fn function_output_field_matches(expected: &ArrowField, actual: &ArrowField) -> bool { @@ -2719,6 +2745,28 @@ mod tests { .unwrap() } + fn exact_arrow_type(field: ArrowField) -> String { + let json = + lance_namespace::schema::arrow_schema_to_json(&ArrowSchema::new(vec![field])).unwrap(); + serde_json::to_string(json.fields[0].r#type.as_ref()).unwrap() + } + + fn single_input_application(path: &str) -> FunctionApplication { + FunctionApplication::from_json( + &serde_json::json!({ + "function": {"name": "inspect", "version": "fv_nested_blob"}, + "inputs": [{ + "parameter": "value", + "kind": "column", + "value": {"path": path} + }], + "output": {"kind": "scalar", "arrow_type": "int64", "nullable": false} + }) + .to_string(), + ) + .unwrap() + } + fn binding_from_plan(plan: &FunctionDeclarationPlan) -> FunctionBinding { let inputs = plan .input_bindings @@ -3158,6 +3206,128 @@ mod tests { assert_eq!(fields[1].data_type(), &DataType::Int32); } + #[test] + fn test_struct_blob_input_preserves_exact_schema_and_nullability() { + let payload = ArrowField::new( + "payload", + DataType::Struct(Fields::from(vec![ + ArrowField::new("mime_type", DataType::Utf8, false), + ArrowField::new( + "nested", + DataType::Struct(Fields::from(vec![crate::blob("image", true)])), + true, + ), + ])), + true, + ); + let plan = plan_function_application( + &ArrowSchema::new(vec![payload]), + &single_input_application("payload"), + Some("size"), + ) + .unwrap(); + + let declared: JsonArrowDataType = + serde_json::from_str(&plan.input_bindings[0].arrow_type).unwrap(); + let DataType::Struct(fields) = + lance_namespace::schema::convert_json_arrow_type(&declared).unwrap() + else { + panic!("expected a struct Function input") + }; + assert!(fields[1].is_nullable()); + let DataType::Struct(nested) = fields[1].data_type() else { + panic!("expected a recursive struct Function input") + }; + assert!(nested[0].is_blob_v2()); + assert!(nested[0].is_nullable()); + + let exact = lance_namespace::schema::convert_json_arrow_schema(&plan.input_schema).unwrap(); + let DataType::Struct(fields) = exact.field(0).data_type() else { + panic!("expected exact input schema to retain the struct") + }; + let DataType::Struct(nested) = fields[1].data_type() else { + panic!("expected exact input schema to retain the nested struct") + }; + assert!(nested[0].is_blob_v2()); + } + + #[test] + fn test_recursive_blob_result_plans_one_whole_named_struct_column() { + let details_type = exact_arrow_type(ArrowField::new( + "details", + DataType::Struct(Fields::from(vec![crate::blob("image", true)])), + false, + )); + let application = FunctionApplication::from_json( + &serde_json::json!({ + "function": {"name": "inspect", "version": "fv_nested_blob"}, + "inputs": [], + "output": { + "kind": "named_struct", + "fields": [ + {"name": "mime_type", "arrow_type": "utf8", "nullable": false}, + {"name": "details", "arrow_type": details_type, "nullable": false} + ] + } + }) + .to_string(), + ) + .unwrap(); + let plan = plan_function_application(&ArrowSchema::empty(), &application, Some("payload")) + .unwrap(); + + assert_eq!(plan.outputs.len(), 1); + assert_eq!(plan.outputs[0].result_field, WHOLE_RESULT_FIELD); + let schema = + lance_namespace::schema::convert_json_arrow_schema(&plan.output_schema).unwrap(); + assert_eq!(schema.field(0).name(), "payload"); + let DataType::Struct(fields) = schema.field(0).data_type() else { + panic!("whole named result must be one struct column") + }; + assert_eq!( + fields.iter().map(|field| field.name()).collect::>(), + ["mime_type", "details"] + ); + let DataType::Struct(details) = fields[1].data_type() else { + panic!("expected recursive result struct") + }; + assert!(details[0].is_blob_v2()); + assert!(!fields.iter().any(|field| field.name() == "payload")); + } + + #[test] + fn test_blob_children_under_collections_are_rejected() { + let collections = vec![ + DataType::List(Arc::new(crate::blob("item", false))), + DataType::LargeList(Arc::new(crate::blob("item", false))), + DataType::FixedSizeList(Arc::new(crate::blob("item", false)), 2), + DataType::Map( + Arc::new(ArrowField::new( + "entries", + DataType::Struct(Fields::from(vec![ + ArrowField::new("key", DataType::Utf8, false), + crate::blob("value", false), + ])), + false, + )), + false, + ), + ]; + for data_type in collections { + let schema = ArrowSchema::new(vec![ArrowField::new("value", data_type, false)]); + let error = plan_function_application( + &schema, + &single_input_application("value"), + Some("size"), + ) + .unwrap_err(); + assert!( + error.to_string().contains("under a collection"), + "got: {error}" + ); + } + } + #[test] fn test_blob_whole_struct_binding_accepts_full_logical_layout() { let input = crate::blob("image", false);