feat: support nested blob function signatures

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}")
 
 

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):

rust/lancedb/src/table/computed_columns.rs

@@ -589,16 +589,13 @@ fn canonical_input_arrow_type(field: &JsonArrowField) -> Result<String> {
         .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>` -> (`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<JsonArrowDataType> {
 }
 
 fn function_output_field(name: &str, nullable: bool, raw: &str) -> Result<JsonArrowField> {
-    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::<Vec<_>>(),
+            ["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);