Custom DTypes With Polars

[Filimoa]

Custom DTypes With Polars

[cosmicBboy]

Looks like a bug in pandera.typing.Series... I think you can try just the bare type and it should work:

class Schema(pa.DataFrameModel):
    city: str
    price: LiteralFloat = pa.Field(coerce=True)

The correct implementation of the custom dtype is also:

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

[cosmicBboy]

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.

[Filimoa]

That worked, I'll open a PR shortly!