Repository: unionai-oss/pandera

Issue:
Unexpected behavior with `coerce=True` in DataFrameSchema
## Bug description
* Using the `@check_output` or `@check_input` decorators do not actually coerce all elements in the dataframe. Instead, it updates the column dtype label, but each individual element is not necessarily coerced into the desired datatype.
---
**For context, here is the issue I am dealing with**
* I have saved my dataframe schema in a yaml file, example entry:
     * Notably, a have set `coerce=True`
```yaml
  numerical_column:
    title: null
    description: This is a numerical column that could possible contain null values
    dtype: float64
    nullable: true
    checks: null
    unique: false
    coerce: true
    required: true
    regex: false
```
* I load this `DataFrameSchema` object from the yaml file: `schema: pandera.DataFrameSchema = DataFrameSchema.from_yaml('schema.yaml')`
* I then use it to check and enforce data types from my functions using the decorator syntax:
```python
@check_input(input_schema)
@check_output(output_schema)
def add_columns_to_df(df: typing.DataFrame[schema]) -> pd.DataFrame:
     do_something_here(df)
     return df
```
* When I check the dtype of the `numerical_column` using `df.numerical_column.dtype` column it is correctly displayed as `float64`, but when I check individual values in the column Series, the nan values are python `None`, and the values are `int64`. 

## Expected Behavior
* When I set `coerce = True` I expect a similar output/behavior as calling `df.numerical_column.astype('float64')` which looks at each element and converts it to the desired datatype

## Temporary solution
I have implemented a temporary solution for anybody dealing with this issue, which also uses the decorator syntax. Simply add the following decorator definition to your pipeline and stack it on top of your other decorators:

```python
def enforce_dtype_coercion(schema: pandera.DataFrameSchema):
    def decorator(func):
        def wrapper(*args, **kwargs):

            df = func(*args, **kwargs)
            schema_dtypes = schema.get_dtypes(df)

            for column_name in df.columns:
                column_type = schema_dtypes.get(column_name)
                if column_type:
                    df[column_name] = df[column_name].astype(column_type.type)

            return df
        return wrapper
    return decorator
```

Then call it like this:
```python
@check_input(input_schema)
@check_output(output_schema)
@enforce_dtype_coercion(output_schema)
def add_columns_to_df(df: typing.DataFrame[schema]) -> pd.DataFrame:
     do_something_here(df)
     return df
```

* This produced the behavior I desired, but I assume that the core developers will likely need a more sophisticated solution.


Discussion:
Thanks for reporting this @ecthompson99. Can you please provide a minimally reproducible example of what you're seeing? Sometimes things get lost in translation, details are lost/data is quirky, but I'll try to repro based on the code you have above.
I was able reproduce the bug for `check_output` but not for `check_input` (which appears to correctly coerce the dtype)

```python
import pandas as pd
import pandera as pa
import pandera.io

yaml_schema = """
schema_type: dataframe
columns:
  numerical_column:
    title: null
    description: This is a numerical column that could possible contain null values
    dtype: float64
    nullable: true
    checks: null
    unique: false
    coerce: true
    required: true
    regex: false
"""


input_schema = pandera.io.from_yaml(yaml_schema)

output_schema = input_schema.add_columns({
    "other_column": pa.Column("float64", coerce=True, nullable=True)
})

@pa.check_input(input_schema)
@pa.check_output(output_schema)
def add_columns_to_df(df: pd.DataFrame) -> pd.DataFrame:
    return df.assign(other_column=1)


data = pd.DataFrame({
    "numerical_column": pd.Series([1, 2, 3, None, None], dtype="object")
})
print(data.dtypes)
output_data = add_columns_to_df(data)
print(output_data)
print(output_data.dtypes)

```

Output:

```
numerical_column    object
dtype: object
   numerical_column  other_column
0               1.0             1
1               2.0             1
2               3.0             1
3               NaN             1
4               NaN             1
numerical_column    float64
other_column          int64
dtype: object
```

Can you confirm that this matches the behavior you're seeing?

Resolving patch:
diff --git a/pandera/decorators.py b/pandera/decorators.py
--- a/pandera/decorators.py
+++ b/pandera/decorators.py
@@ -314,23 +314,49 @@ def check_output(
     See :ref:`here<decorators>` for more usage details.
     """
 
+    # make sure that callable obj_getter doesn't work when the schema has
+    # any component that requires coercion, since there's no way to re-assign
+    # the output to the coerced data.
+    # pylint: disable=too-many-boolean-expressions
+    if callable(obj_getter) and (
+        schema.coerce
+        or (schema.index is not None and schema.index.coerce)
+        or (
+            isinstance(schema, schemas.DataFrameSchema)
+            and any(col.coerce for col in schema.columns.values())
+        )
+    ):
+        raise ValueError(
+            "Cannot use callable obj_getter when the schema uses coercion."
+        )
+
     def validate(out: Any, fn: Callable) -> None:
+        def _try_validate(obj: Any):
+            try:
+                return schema.validate(
+                    obj, head, tail, sample, random_state, lazy, inplace
+                )
+            except errors.SchemaError as e:
+                _handle_schema_error("check_output", fn, schema, obj, e)
+
         if obj_getter is None:
-            obj = out
+            return _try_validate(out)
         elif isinstance(obj_getter, (int, str)):
             obj = out[obj_getter]
+            validated = _try_validate(obj)
+            if isinstance(out, tuple):
+                out = list(out)
+                out[obj_getter] = validated
+                out = tuple(out)
+            else:
+                out[obj_getter] = validated
+            return out
         elif callable(obj_getter):
             obj = obj_getter(out)
-        else:
-            raise TypeError(
-                f"obj_getter is unrecognized type: {type(obj_getter)}"
-            )
-        try:
-            schema.validate(
-                obj, head, tail, sample, random_state, lazy, inplace
-            )
-        except errors.SchemaError as e:
-            _handle_schema_error("check_output", fn, schema, obj, e)
+            _try_validate(obj)
+            return out
+
+        raise TypeError(f"obj_getter is unrecognized type: {type(obj_getter)}")
 
     @wrapt.decorator
     def _wrapper(
@@ -360,8 +386,7 @@ async def aio_wrapper():
             return aio_wrapper()
         else:
             out = fn(*args, **kwargs)
-            validate(out, fn)
-            return out
+            return validate(out, fn)
 
     return _wrapper
 


Test patch:
diff --git a/tests/core/test_decorators.py b/tests/core/test_decorators.py
--- a/tests/core/test_decorators.py
+++ b/tests/core/test_decorators.py
@@ -193,6 +193,64 @@ def test_incorrect_check_input_index(df):
         test_incorrect_check_input_index(pd.DataFrame({"column1": [1, 2, 3]}))
 
 
+def test_check_decorator_coercion() -> None:
+    """Test that check decorators correctly coerce input/output data."""
+
+    in_schema = DataFrameSchema({"column1": Column(int, coerce=True)})
+    out_schema = DataFrameSchema({"column2": Column(float, coerce=True)})
+
+    @check_input(in_schema)
+    @check_output(out_schema)
+    def test_func_io(df):
+        return df.assign(column2=10)
+
+    @check_input(in_schema)
+    @check_output(out_schema, obj_getter=1)
+    def test_func_out_tuple_obj_getter(df):
+        return None, df.assign(column2=10)
+
+    @check_input(in_schema)
+    @check_output(out_schema, obj_getter=1)
+    def test_func_out_list_obj_getter(df):
+        return None, df.assign(column2=10)
+
+    @check_input(in_schema)
+    @check_output(out_schema, obj_getter="key")
+    def test_func_out_dict_obj_getter(df):
+        return {"key": df.assign(column2=10)}
+
+    cases: typing.Iterable[
+        typing.Tuple[typing.Callable, typing.Union[int, str, None]]
+    ] = [
+        (test_func_io, None),
+        (test_func_out_tuple_obj_getter, 1),
+        (test_func_out_list_obj_getter, 1),
+        (test_func_out_dict_obj_getter, "key"),
+    ]
+    for fn, key in cases:
+        out = fn(pd.DataFrame({"column1": ["1", "2", "3"]}))
+        if key is not None:
+            out = out[key]
+        assert out.dtypes["column1"] == "int64"
+        assert out.dtypes["column2"] == "float64"
+
+
+def test_check_output_coercion_error() -> None:
+    """Test that check_output raises ValueError when obj_getter is callable."""
+
+    with pytest.raises(
+        ValueError,
+        match="Cannot use callable obj_getter when the schema uses coercion",
+    ):
+
+        @check_output(
+            DataFrameSchema({"column2": Column(float, coerce=True)}),
+            obj_getter=lambda x: x[0]["key"],
+        )
+        def test_func(df):  # pylint: disable=unused-argument
+            ...
+
+
 def test_check_instance_method_decorator_error() -> None:
     """Test error message on methods."""
     # pylint: disable-next=missing-class-docstring
@@ -309,6 +367,7 @@ def multiple_outputs_tuple(df):
         return df, df
 
     @check_io(
+        # pylint: disable=undefined-loop-variable
         out=[(0, schema), ("foo", schema), (lambda x: x[2]["bar"], schema)]
     )
     def multiple_outputs_dict(df):
@@ -339,6 +398,14 @@ def validate_inplace(df):
     invalid_df = pd.DataFrame({"col": [-1, -1, -1]})
     expected = pd.DataFrame({"col": [3, 3, 3]})
 
+    def _assert_equals(actual, expect):
+        if isinstance(actual, pd.Series):
+            assert (actual == expect).all()
+        elif isinstance(actual, pd.DataFrame):
+            assert (actual == expect).all(axis=None)
+        else:
+            assert actual == expect
+
     for fn, valid, invalid, out in [
         (simple_func, [df1, df2], [invalid_df, invalid_df], expected),
         (simple_func_no_out, [df1, df2], [invalid_df, invalid_df], expected),
@@ -358,11 +425,12 @@ def validate_inplace(df):
     ]:
         result = fn(*valid)  # type: ignore[operator]
         if isinstance(result, pd.Series):
-            assert (result == out).all()
+            _assert_equals(result, out)
         if isinstance(result, pd.DataFrame):
-            assert (result == out).all(axis=None)
+            _assert_equals(result, out)
         else:
-            assert result == out
+            for _result, _out in zip(result, out):
+                _assert_equals(_result, _out)
 
         expected_error = (
             errors.SchemaErrors if fn is validate_lazy else errors.SchemaError