Repository: unionai-oss/pandera

Issue:
Custom DTypes With Polars
**Describe the bug**
A clear and concise description of what the bug is.

- [x] I have checked that this issue has not already been reported.
- [x] I have confirmed this bug exists on the latest version of pandera.
- [ ] (optional) I have confirmed this bug exists on the master branch of pandera.

I'm not sure if this is a bug, intentional or just missing documentation.

#### Code Sample, a copy-pastable example

```python
from pandera import dtypes
from pandera.engines import numpy_engine, pandas_engine
from pandera.typing import Series, Bool
from pandas.api.types import infer_dtype
import pandas as pd
import pandera as pa_core


@numpy_engine.Engine.register_dtype
@dtypes.immutable
class LiteralFloat(numpy_engine.Float64):
    def coerce(self, series: Series) -> Series[Bool]:
        """If comes across a string, coerce it to a float. If it fails, return NaN."""
        if "string" in infer_dtype(series):
            series = series.apply(
                lambda x: str.replace(x, ",", "") if isinstance(x, str) else x
            )

        return pd.to_numeric(series, errors="coerce")


class Schema(pa_core.SchemaModel):
    state: str
    price: Series[LiteralFloat]

    class Config:
        strict = "filter"
        coerce = True


df = pd.DataFrame(
    {
        "state": ["FL", "FL", "FL", "CA", "CA", "CA"],
        "price": ["8,000", "12.0", "10.0", "16.0", "20.0", "18.0"],
    }
)

Schema.validate(df)
```

With the pandas API this was possible - you could write some custom dtypes that perform some basic data cleaning. For example, in our case we had a `YesNoBool` that coerces -> "yes" / "no" to booleans. This was handy since we deal with hundreds of these columns and it's a pain to write transformation logic for each one.

The documentation is pretty vague on this (not sure if this is an anti-pattern) but this was my best attempt on porting the code to polars.

```python
import pandera.polars as pa
import polars as pl
from pandera.engines import polars_engine
from pandera import dtypes
from pandera.typing import Series


@polars_engine.Engine.register_dtype
@dtypes.immutable
class LiteralFloat(pl.Float64):
    def coerce(self, series):
        """If comes across a string, remove commas and coerce it to a float. If it fails, return NaN."""
        series = series.str.replace(",", "").cast(pl.Float64, strict=False)
        return series


class Schema(pa.DataFrameModel):
    city: str
    price: Series[LiteralFloat] = pa.Field(coerce=True)

    class Config:
        strict = "filter"
        coerce = True


dl = pl.from_pandas(df)
Schema.validate(dl)
```

```
>>> SchemaInitError: Invalid annotation 'price: pandera.typing.pandas.Series[__main__.LiteralFloat]'
```

Is this intentional?

#### Desktop (please complete the following information):

 - OS: Apple Silicon
 - Version:
 ```
pandera==0.19.0b1
polars==0.20.21
pandas==2.2.2
```

#### Screenshots
None

#### Additional context
I'll be glad to open a PR to update the docs if this is just a docs issue.



Discussion:
Looks like a bug in `pandera.typing.Series`... I think you can try just the bare type and it should work:

```python
class Schema(pa.DataFrameModel):
    city: str
    price: LiteralFloat = pa.Field(coerce=True)
```

The correct implementation of the custom dtype is also:

```python
from pandera.api.polars.types import PolarsData

@polars_engine.Engine.register_dtype
@dtypes.immutable
class LiteralFloat(polars_engine.Float64):  # 👈  inherit from polars_engine.Float64, not the polars dtype
    def coerce(self, polars_data: PolarsData) -> pl.LazyFrame:  # 👈 note the input and output signature
        """If comes across a string, remove commas and coerce it to a float. If it fails, return NaN."""
        return polars_data.lazyframe.with_columns(  # 👈 must return a lazyframe
            pl.col(polars_data.key)
            .str.replace(",", "")
            .cast(pl.Float64, strict=False)
        )
```

See the polars engine DataType implementation for details on the signatures of these methods:
https://github.com/unionai-oss/pandera/blob/main/pandera/engines/polars_engine.py#L91

I'll look into fixing the `SchemaInitError: Invalid annotation 'price: pandera.typing.pandas.Series[__main__.LiteralFloat]'` issue, if you can, would be great if the polars docs can be updated with an example of a custom datatype: https://github.com/unionai-oss/pandera/blob/main/docs/source/polars.md
> I'll look into fixing the SchemaInitError: Invalid annotation 'price: pandera.typing.pandas.Series[__main__.LiteralFloat]' issue

So the whole `Series[TYPE]` syntax is only supported in the pandas DataFrameModel and will be deprecated in that API eventually... looking forward to new backends (in this case polars) the more concise bare type will be supported. I'll add a more informative error message here.
That worked, I'll open a PR shortly!

Resolving patch:
diff --git a/pandera/api/polars/model.py b/pandera/api/polars/model.py
--- a/pandera/api/polars/model.py
+++ b/pandera/api/polars/model.py
@@ -89,9 +89,19 @@ def _build_columns(  # pylint:disable=too-many-locals
                 columns[field_name] = Column(**column_kwargs)
 
             else:
+                origin_name = (
+                    f"{annotation.origin.__module__}."
+                    f"{annotation.origin.__name__}"
+                )
+                msg = (
+                    " Series[TYPE] annotations are not supported for polars. "
+                    "Use the bare TYPE directly"
+                    if origin_name == "pandera.typing.pandas.Series"
+                    else ""
+                )
                 raise SchemaInitError(
                     f"Invalid annotation '{field_name}: "
-                    f"{annotation.raw_annotation}'"
+                    f"{annotation.raw_annotation}'.{msg}"
                 )
 
         return columns


Test patch: