fix(pandas): Annotated[pd.DataFrame, DataframeSchema(columns=["a"])] output spec

src/_bentoml_sdk/validators.py

@@ -278,7 +278,10 @@ def validate(self, obj: t.Any) -> t.Any:
 @attrs.frozen(unsafe_hash=True)
 class DataframeSchema:
     orient: str = "records"
-    columns: list[str] | None = None
+    columns: tuple[str] | None = attrs.field(
+        default=None,
+        converter=lambda x: tuple(x) if x else None,
+    )
 
     def __get_pydantic_json_schema__(
         self, schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler

tests/unit/bentoml_io/test_decorators.py

@@ -1,8 +1,11 @@
 from typing import Generator
 
+import pandas as pd
 import pytest
+from typing_extensions import Annotated
 
 import bentoml
+from bentoml.validators import DataframeSchema
 
 
 @pytest.mark.asyncio
@@ -90,3 +93,26 @@ def stream(self, name: str) -> Generator[str, None, None]:
         "Hello, world! 1",
         "Hello, world! 2",
     ]
+
+
+def test_api_decorator_pandas():
+    @bentoml.api
+    def pandas_func(
+        _,  # The decorator assumes `self` is the first arg.
+        df1: pd.DataFrame,
+        df2: Annotated[pd.DataFrame, DataframeSchema(columns=("b",))],
+    ) -> Annotated[
+        pd.DataFrame,
+        DataframeSchema(orient="columns", columns=["a", "b"]),
+    ]:
+        return pd.concat([df1, df2], axis=1)
+
+    pandas_func.input_spec.model_fields["df1"].annotation is pd.DataFrame
+    pandas_func.input_spec.model_fields["df2"].annotation is Annotated[
+        pd.DataFrame,
+        DataframeSchema(columns=("b",)),
+    ]
+    pandas_func.output_spec.model_fields["root"].annotation is Annotated[
+        pd.DataFrame,
+        DataframeSchema(orient="columns", columns=("a", "b")),
+    ]