Repository: unionai-oss/pandera

Issue:
Add a polars `Series` type
**Is your feature request related to a problem? Please describe.**
This would provide a similar `pandera.typing.Series` type for the polars API that exists for pandas and other backends. It is not strictly necessary since it does not get checked, but it would be good to have to match the existing API and to help type checkers understand accessing string column names from the class attribute names.

**Describe the solution you'd like**
Allow for annotating polars data frame fields as:
```python
import pandera.polars as pa
from pandera.typing.polars import Series

class MySchema(pa.DataFrameModel):
    a: Series[int]
```



Discussion:
Working on this now

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
@@ -21,6 +21,7 @@
 from pandera.engines import polars_engine as pe
 from pandera.errors import SchemaInitError
 from pandera.typing import AnnotationInfo
+from pandera.typing.polars import Series
 
 
 class DataFrameModel(_DataFrameModel[pl.LazyFrame, DataFrameSchema]):
@@ -74,6 +75,7 @@ def _build_columns(  # pylint:disable=too-many-locals
             if (
                 annotation.origin is None
                 or isinstance(annotation.origin, pl.datatypes.DataTypeClass)
+                or annotation.origin is Series
                 or engine_dtype
             ):
                 if check_name is False:
@@ -94,19 +96,9 @@ 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}'.{msg}"
+                    f"{annotation.raw_annotation}'."
                 )
 
         return columns
diff --git a/pandera/typing/polars.py b/pandera/typing/polars.py
--- a/pandera/typing/polars.py
+++ b/pandera/typing/polars.py
@@ -4,7 +4,11 @@
 
 from packaging import version
 
-from pandera.typing.common import DataFrameBase, DataFrameModel
+from pandera.typing.common import (
+    DataFrameBase,
+    DataFrameModel,
+    SeriesBase,
+)
 
 
 try:
@@ -35,3 +39,11 @@ class LazyFrame(DataFrameBase, pl.LazyFrame, Generic[T]):
 
         *new in 0.19.0*
         """
+
+    # pylint: disable=too-few-public-methods
+    class Series(SeriesBase, pl.Series, Generic[T]):
+        """
+        Pandera generic for pl.Series, only used for type annotation.
+
+        *new in 0.19.0*
+        """


Test patch:
diff --git a/tests/polars/test_polars_dataframe_generic.py b/tests/polars/test_polars_dataframe_generic.py
--- a/tests/polars/test_polars_dataframe_generic.py
+++ b/tests/polars/test_polars_dataframe_generic.py
@@ -4,7 +4,24 @@
 import pytest
 
 import pandera.polars as pa
-from pandera.typing.polars import LazyFrame
+from pandera.typing.polars import LazyFrame, Series
+
+
+def test_series_annotation():
+    class Model(pa.DataFrameModel):
+        col1: Series[pl.Int64]
+
+    data = pl.LazyFrame(
+        {
+            "col1": [1, 2, 3],
+        }
+    )
+
+    assert data.collect().equals(Model.validate(data).collect())
+
+    invalid_data = data.cast({"col1": pl.Float64})
+    with pytest.raises(pa.errors.SchemaError):
+        Model.validate(invalid_data).collect()
 
 
 def test_lazyframe_generic_simple():