feat(python): support nested Function Arrow types (#4088)

Teach Python Function authoring to retain the compact V1 grammar for
existing types and emit canonical exact JSON for nested struct
signatures. Adds coverage for recursive struct/list schemas and exact
field properties.

python/python/lancedb/functions.py

@@ -502,31 +502,107 @@ _GRAMMAR_PRIMITIVES = (
 
 
 def _canonical_arrow_type(data_type: pa.DataType) -> str:
-    """The server's V1 Function type grammar. Anything outside it is rejected
-    here rather than at registration."""
+    """The compact Function grammar, or canonical exact JSON for nested types."""
+    grammar = _grammar_arrow_type(data_type)
+    if grammar is not None:
+        return grammar
+    exact = _exact_arrow_type(data_type)
+    return json.dumps(exact, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+
+
+def _grammar_arrow_type(data_type: pa.DataType) -> Optional[str]:
     for candidate, name in _GRAMMAR_PRIMITIVES:
         if data_type == candidate:
             return name
     if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
+        item = _grammar_list_item(data_type)
+        if item is None:
+            return None
         prefix = "list" if pa.types.is_list(data_type) else "large_list"
-        return f"{prefix}<{_canonical_list_item(data_type)}>"
+        return f"{prefix}<{item}>"
     if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
-        return (
-            f"fixed_size_list<{_canonical_list_item(data_type)}, {data_type.list_size}>"
-        )
-    raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
+        item = _grammar_list_item(data_type)
+        if item is not None:
+            return f"fixed_size_list<{item}, {data_type.list_size}>"
+    return None
 
 
-def _canonical_list_item(data_type: pa.DataType) -> str:
+def _grammar_list_item(data_type: pa.DataType) -> Optional[str]:
     """The grammar names only the item type; it always means a non-nullable
-    child called `item`, so any other child metadata cannot be represented."""
+    child called `item`, so other child properties require exact JSON."""
     child = data_type.value_field
     if child.name != "item" or child.nullable or child.metadata:
+        return None
+    return _grammar_arrow_type(child.type)
+
+
+def _validate_exact_arrow_field(field: pa.Field) -> None:
+    if not field.name:
         raise TypeError(
-            "unsupported Arrow type for Function signature: list items must be a "
-            f"non-nullable field named 'item', got {child}"
+            "unsupported Arrow type for Function signature: field names "
+            "must not be empty"
         )
-    return _canonical_arrow_type(child.type)
+    if field.metadata:
+        raise TypeError(
+            "unsupported Arrow type for Function signature: field metadata "
+            f"is not supported, got {field}"
+        )
+
+
+def _exact_arrow_field(field: pa.Field) -> dict[str, Any]:
+    _validate_exact_arrow_field(field)
+    return {
+        "name": field.name,
+        "nullable": field.nullable,
+        "type": _exact_arrow_type(field.type),
+    }
+
+
+def _exact_arrow_type(data_type: pa.DataType) -> dict[str, Any]:
+    for candidate, name in _GRAMMAR_PRIMITIVES:
+        if data_type == candidate:
+            return {"type": name}
+    if pa.types.is_struct(data_type):
+        fields = list(data_type)
+        names = [field.name for field in fields]
+        if not fields or len(set(names)) != len(names):
+            raise TypeError(
+                "unsupported Arrow type for Function signature: structs must have "
+                "non-empty, uniquely named fields"
+            )
+        return {
+            "type": "struct",
+            "fields": [_exact_arrow_field(field) for field in fields],
+        }
+    if (
+        pa.types.is_list(data_type)
+        or pa.types.is_large_list(data_type)
+        or pa.types.is_fixed_size_list(data_type)
+    ):
+        if pa.types.is_fixed_size_list(data_type):
+            if data_type.value_field.name != "item":
+                raise TypeError(
+                    "unsupported Arrow type for Function signature: fixed-size list "
+                    "items must be named 'item'"
+                )
+            if data_type.list_size <= 0:
+                raise TypeError(
+                    f"unsupported Arrow type for Function signature: {data_type}"
+                )
+        value: dict[str, Any] = {
+            "type": (
+                "list"
+                if pa.types.is_list(data_type)
+                else "large_list"
+                if pa.types.is_large_list(data_type)
+                else "fixed_size_list"
+            ),
+            "fields": [_exact_arrow_field(data_type.value_field)],
+        }
+        if pa.types.is_fixed_size_list(data_type):
+            value["length"] = data_type.list_size
+        return value
+    raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")
 
 
 def _list_of(item: pa.DataType) -> pa.DataType:
@@ -600,8 +676,11 @@ def _callable_parameters(function: Callable[..., Any]) -> tuple[inspect.Paramete
 
 def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutput:
     if isinstance(output, pa.Schema):
+        if output.metadata:
+            raise TypeError("Function output schema metadata is not supported")
         fields = tuple(output)
     elif isinstance(output, pa.Field) and pa.types.is_struct(output.type):
+        _validate_exact_arrow_field(output)
         if output.nullable:
             raise ValueError("Function output must be non-nullable")
         fields = tuple(output.type)
@@ -617,6 +696,7 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
             raise TypeError(
                 "output_schema must be a PyArrow DataType, Field, or Schema"
             )
+        _validate_exact_arrow_field(field)
         if field.nullable:
             raise ValueError("Function output must be non-nullable")
         return FunctionOutput(
@@ -629,6 +709,8 @@ def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutp
         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]
     if len(set(names)) != len(names):
         raise ValueError("Function output field names must be unique")
@@ -657,6 +739,10 @@ def _infer_signature(
     if input_schema is not None:
         if not isinstance(input_schema, pa.Schema):
             raise TypeError("input_schema must be a PyArrow Schema")
+        if input_schema.metadata:
+            raise TypeError("Function input schema metadata is not supported")
+        for field in input_schema:
+            _validate_exact_arrow_field(field)
         expected = tuple(parameter.name for parameter in parameters)
         actual = tuple(input_schema.names)
         if actual != expected:

python/python/tests/test_first_class_function_slice2.py

@@ -168,7 +168,7 @@ def test_udf_resolves_module_globals_before_builtins(tmp_path):
         udf(module.uses_callable_shadow)
 
 
-def test_canonical_arrow_type_is_exactly_the_grammar():
+def test_canonical_arrow_type_prefers_the_compact_grammar():
     from lancedb.functions import _GRAMMAR_PRIMITIVES, _canonical_arrow_type
 
     golden = json.loads(
@@ -181,6 +181,13 @@ def test_canonical_arrow_type_is_exactly_the_grammar():
         case["arrow_type"] for case in golden["valid"] if "<" not in case["arrow_type"]
     ]
     assert [name for _, name in _GRAMMAR_PRIMITIVES] == primitives
+    assert _canonical_arrow_type(pa.list_(pa.field("item", pa.float32(), False))) == (
+        "list<float32>"
+    )
+    assert (
+        _canonical_arrow_type(pa.large_list(pa.field("item", pa.float32(), False)))
+        == "large_list<float32>"
+    )
     for outside in [
         pa.timestamp("us"),
         pa.decimal128(10, 2),
@@ -188,7 +195,6 @@ def test_canonical_arrow_type_is_exactly_the_grammar():
         pa.large_binary(),
         pa.binary(4),
         pa.duration("s"),
-        pa.struct([pa.field("a", pa.int32())]),
         pa.list_(pa.float32(), 0),
         pa.list_(pa.timestamp("us")),
     ]:
@@ -378,14 +384,29 @@ def test_udf_recursion_versus_a_rebound_module_name(tmp_path):
         udf(raw_fact)
 
 
-def test_canonical_arrow_type_rejects_unrepresentable_list_children():
+def test_canonical_arrow_type_uses_exact_json_for_list_child_properties():
     from lancedb.functions import _canonical_arrow_type
 
+    nullable = pa.list_(pa.float32())
+    assert json.loads(_canonical_arrow_type(nullable)) == {
+        "type": "list",
+        "fields": [
+            {
+                "name": "item",
+                "nullable": True,
+                "type": {"type": "float32"},
+            }
+        ],
+    }
+    named = pa.list_(pa.field("custom", pa.float32(), nullable=False))
+    assert json.loads(_canonical_arrow_type(named))["fields"][0]["name"] == "custom"
     for outside in [
-        pa.list_(pa.float32()),  # pyarrow default: nullable child
-        pa.list_(pa.field("custom", pa.float32(), nullable=False)),
         pa.list_(pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"})),
         pa.list_(pa.field("item", pa.float32(), nullable=False), 0),
+        pa.list_(
+            pa.field("item", pa.float32(), nullable=False, metadata={"k": "v"}), 3
+        ),
+        pa.list_(pa.field("custom", pa.float32(), nullable=False), 3),
     ]:
         with pytest.raises(TypeError, match="unsupported Arrow type"):
             _canonical_arrow_type(outside)
@@ -395,6 +416,29 @@ def test_canonical_arrow_type_rejects_unrepresentable_list_children():
         )
         == "fixed_size_list<float32, 3>"
     )
+    fixed = json.loads(_canonical_arrow_type(pa.list_(pa.float32(), 3)))
+    assert fixed == {
+        "type": "fixed_size_list",
+        "fields": [
+            {
+                "name": "item",
+                "nullable": True,
+                "type": {"type": "float32"},
+            }
+        ],
+        "length": 3,
+    }
+    large = json.loads(_canonical_arrow_type(pa.large_list(pa.float32())))
+    assert large["type"] == "large_list"
+    assert large["fields"][0]["nullable"] is True
+
+    for invalid_struct in [
+        pa.struct([]),
+        pa.struct([pa.field("a", pa.int32()), pa.field("a", pa.int64())]),
+        pa.struct([pa.field("", pa.int32())]),
+    ]:
+        with pytest.raises(TypeError, match="unsupported Arrow type"):
+            _canonical_arrow_type(invalid_struct)
 
 
 def _calls_missing(value: int) -> int:
@@ -482,6 +526,105 @@ def test_explicit_arrow_schema_is_deterministic():
     assert signature.output.nullable is False
 
 
+def test_nested_struct_output_uses_canonical_exact_json():
+    token = pa.struct(
+        [
+            pa.field("position", pa.int32(), nullable=False),
+            pa.field("value", pa.string(), nullable=False),
+            pa.field("length", pa.int32(), nullable=False),
+        ]
+    )
+    analysis = pa.struct(
+        [
+            pa.field("normalized_text", pa.string(), nullable=False),
+            pa.field("has_content", pa.bool_(), nullable=False),
+            pa.field(
+                "metrics",
+                pa.struct(
+                    [
+                        pa.field("character_count", pa.int64(), nullable=False),
+                        pa.field("word_count", pa.int32(), nullable=False),
+                        pa.field("average_word_length", pa.float64(), nullable=False),
+                    ]
+                ),
+                nullable=False,
+            ),
+            pa.field(
+                "diagnostics",
+                pa.struct(
+                    [
+                        pa.field("status", pa.string(), nullable=False),
+                        pa.field(
+                            "normalization",
+                            pa.struct(
+                                [
+                                    pa.field("changed", pa.bool_(), nullable=False),
+                                    pa.field(
+                                        "original_length", pa.int64(), nullable=False
+                                    ),
+                                ]
+                            ),
+                            nullable=False,
+                        ),
+                    ]
+                ),
+                nullable=False,
+            ),
+            pa.field(
+                "token_preview",
+                pa.list_(pa.field("item", token, nullable=False)),
+                nullable=False,
+            ),
+        ]
+    )
+
+    @udf(
+        input_schema=pa.schema([pa.field("text", pa.string(), nullable=False)]),
+        output_schema=pa.field("analysis", analysis, nullable=False),
+    )
+    def analyze(text):
+        return {"normalized_text": text}
+
+    output = analyze.registration_request.signature.output
+    assert output.kind == "named_struct"
+    assert [field.name for field in output.fields] == [
+        "normalized_text",
+        "has_content",
+        "metrics",
+        "diagnostics",
+        "token_preview",
+    ]
+    metrics = json.loads(output.fields[2].arrow_type)
+    assert metrics == {
+        "type": "struct",
+        "fields": [
+            {
+                "name": "character_count",
+                "nullable": False,
+                "type": {"type": "int64"},
+            },
+            {
+                "name": "word_count",
+                "nullable": False,
+                "type": {"type": "int32"},
+            },
+            {
+                "name": "average_word_length",
+                "nullable": False,
+                "type": {"type": "float64"},
+            },
+        ],
+    }
+    preview = json.loads(output.fields[4].arrow_type)
+    assert preview["type"] == "list"
+    assert preview["fields"][0]["type"]["type"] == "struct"
+    assert [field["name"] for field in preview["fields"][0]["type"]["fields"]] == [
+        "position",
+        "value",
+        "length",
+    ]
+
+
 def test_annotation_and_explicit_schema_validation_fail_closed():
     with pytest.raises(TypeError, match="missing Function annotations"):
 
@@ -525,6 +668,72 @@ def test_annotation_and_explicit_schema_validation_fail_closed():
         def nullable_explicit(value):
             return value
 
+    for invalid_field in [
+        pa.field("", pa.int32(), nullable=False),
+        pa.field("result", pa.int32(), nullable=False, metadata={"k": "v"}),
+    ]:
+        with pytest.raises(TypeError, match="unsupported Arrow type"):
+
+            @udf(
+                input_schema=pa.schema([pa.field("value", pa.int64())]),
+                output_schema=pa.schema([invalid_field]),
+            )
+            def invalid_explicit_field(value):
+                return value
+
+    with pytest.raises(TypeError, match="unsupported Arrow type"):
+
+        @udf(
+            input_schema=pa.schema(
+                [pa.field("value", pa.int64(), metadata={"k": "v"})]
+            ),
+            output_schema=pa.int64(),
+        )
+        def input_field_metadata(value):
+            return value
+
+    with pytest.raises(TypeError, match="unsupported Arrow type"):
+
+        @udf(
+            input_schema=pa.schema([pa.field("value", pa.int64())]),
+            output_schema=pa.field(
+                "result", pa.int64(), nullable=False, metadata={"k": "v"}
+            ),
+        )
+        def scalar_output_field_metadata(value):
+            return value
+
+    struct_type = pa.struct([pa.field("value", pa.int64(), nullable=False)])
+    with pytest.raises(TypeError, match="unsupported Arrow type"):
+
+        @udf(
+            input_schema=pa.schema([pa.field("value", pa.int64())]),
+            output_schema=pa.field(
+                "result", struct_type, nullable=False, metadata={"k": "v"}
+            ),
+        )
+        def struct_output_field_metadata(value):
+            return {"value": value}
+
+    for input_schema, output_schema in [
+        (
+            pa.schema([pa.field("value", pa.int64())], metadata={"k": "v"}),
+            pa.int64(),
+        ),
+        (
+            pa.schema([pa.field("value", pa.int64())]),
+            pa.schema(
+                [pa.field("result", pa.int64(), nullable=False)],
+                metadata={"k": "v"},
+            ),
+        ),
+    ]:
+        with pytest.raises(TypeError, match="schema metadata"):
+
+            @udf(input_schema=input_schema, output_schema=output_schema)
+            def schema_metadata(value):
+                return value
+
 
 def test_local_function_catalog_operations_are_not_supported(tmp_path):
     db = lancedb.connect(tmp_path)