Support native PySpark.sql on Pandera

asv_bench/benchmarks/dataframe_schema.py

@@ -2,9 +2,20 @@
 import pandas as pd
 
 from pandera import (
-    Column, DataFrameSchema, Bool, Category, Check,
-    DateTime, Float, Int, Object, String, Timedelta,
-    check_input, check_output)
+    Column,
+    DataFrameSchema,
+    Bool,
+    Category,
+    Check,
+    DateTime,
+    Float,
+    Int,
+    Object,
+    String,
+    Timedelta,
+    check_input,
+    check_output,
+)
 
 
 class Validate:
@@ -14,41 +25,46 @@ class Validate:
 
     def setup(self):
         self.schema = DataFrameSchema(
-                {
-                    "a": Column(Int),
-                    "b": Column(Float),
-                    "c": Column(String),
-                    "d": Column(Bool),
-                    "e": Column(Category),
-                    "f": Column(Object),
-                    "g": Column(DateTime),
-                    "i": Column(Timedelta),
-                    },
-                )
+            {
+                "a": Column(Int),
+                "b": Column(Float),
+                "c": Column(String),
+                "d": Column(Bool),
+                "e": Column(Category),
+                "f": Column(Object),
+                "g": Column(DateTime),
+                "i": Column(Timedelta),
+            },
+        )
         self.df = pd.DataFrame(
-                {
-                    "a": [1, 2, 3],
-                    "b": [1.1, 2.5, 9.9],
-                    "c": ["z", "y", "x"],
-                    "d": [True, True, False],
-                    "e": pd.Series(["c2", "c1", "c3"], dtype="category"),
-                    "f": [(3,), (2,), (1,)],
-                    "g": [pd.Timestamp("2015-02-01"),
-                          pd.Timestamp("2015-02-02"),
-                          pd.Timestamp("2015-02-03")],
-                    "i": [pd.Timedelta(1, unit="D"),
-                          pd.Timedelta(5, unit="D"),
-                          pd.Timedelta(9, unit="D")]
-                    })
+            {
+                "a": [1, 2, 3],
+                "b": [1.1, 2.5, 9.9],
+                "c": ["z", "y", "x"],
+                "d": [True, True, False],
+                "e": pd.Series(["c2", "c1", "c3"], dtype="category"),
+                "f": [(3,), (2,), (1,)],
+                "g": [
+                    pd.Timestamp("2015-02-01"),
+                    pd.Timestamp("2015-02-02"),
+                    pd.Timestamp("2015-02-03"),
+                ],
+                "i": [
+                    pd.Timedelta(1, unit="D"),
+                    pd.Timedelta(5, unit="D"),
+                    pd.Timedelta(9, unit="D"),
+                ],
+            }
+        )
 
     def time_df_schema(self):
         self.schema.validate(self.df)
 
     def mem_df_schema(self):
-         self.schema.validate(self.df)
+        self.schema.validate(self.df)
 
     def peakmem_df_schema(self):
-         self.schema.validate(self.df)
+        self.schema.validate(self.df)
 
 
 class Decorators:

asv_bench/benchmarks/series_schema.py

@@ -2,8 +2,20 @@
 import pandas as pd
 
 from pandera import (
-    Column, DataFrameSchema, SeriesSchema, Bool, Category, Check,
-    DateTime, Float, Int, Object, String, Timedelta, String)
+    Column,
+    DataFrameSchema,
+    SeriesSchema,
+    Bool,
+    Category,
+    Check,
+    DateTime,
+    Float,
+    Int,
+    Object,
+    String,
+    Timedelta,
+    String,
+)
 
 
 class Validate:
@@ -13,23 +25,23 @@ class Validate:
 
     def setup(self):
         self.schema = SeriesSchema(
-                String,
-                checks=[
-                    Check(lambda s: s.str.startswith("foo")),
-                    Check(lambda s: s.str.endswith("bar")),
-                    Check(lambda x: len(x) > 3, element_wise=True)
-                    ],
-                nullable=False,
-                unique=False,
-                name="my_series")
-        self.series = pd.Series(["foobar", "foobar", "foobar"],
-                                name="my_series")
+            String,
+            checks=[
+                Check(lambda s: s.str.startswith("foo")),
+                Check(lambda s: s.str.endswith("bar")),
+                Check(lambda x: len(x) > 3, element_wise=True),
+            ],
+            nullable=False,
+            unique=False,
+            name="my_series",
+        )
+        self.series = pd.Series(["foobar", "foobar", "foobar"], name="my_series")
 
     def time_series_schema(self):
         self.schema.validate(self.series)
 
     def mem_series_schema(self):
-         self.schema.validate(self.series)
+        self.schema.validate(self.series)
 
     def peakmem_series_schema(self):
-         self.schema.validate(self.series)
+        self.schema.validate(self.series)

docs/source/conf.py

@@ -188,6 +188,7 @@
 )
 copybutton_prompt_is_regexp = True
 
+
 # this is a workaround to filter out forward reference issue in
 # sphinx_autodoc_typehints
 class FilterPandasTypeAnnotationWarning(pylogging.Filter):
@@ -215,6 +216,7 @@ def filter(self, record: pylogging.LogRecord) -> bool:
     FilterPandasTypeAnnotationWarning()
 )
 
+
 # based on pandas/doc/source/conf.py
 def linkcode_resolve(domain, info):
     """Determine the URL corresponding to Python object."""

docs/source/dataframe_models.rst

@@ -205,10 +205,11 @@ You can easily convert a :class:`~pandera.api.pandas.model.DataFrameModel` class
         coerce=False,
         dtype=None,
         index=None,
-        strict=False
+        strict=False,
         name=InputSchema,
         ordered=False,
-        unique_column_names=False
+        unique_column_names=False,
+        metadata=None,
     )>
 
 You can also use the :meth:`~pandera.api.pandas.model.DataFrameModel.validate` method to

docs/source/dataframe_schemas.rst

@@ -810,10 +810,11 @@ data pipeline:
         coerce=False,
         dtype=None,
         index=None,
-        strict=True
+        strict=True,
         name=None,
         ordered=False,
-        unique_column_names=False
+        unique_column_names=False,
+        metadata=None,
     )>
 
 If during the course of a data pipeline one of your columns is moved into the
@@ -858,10 +859,11 @@ the pipeline output.
             name=None,
             ordered=True
         )>,
-        strict=True
+        strict=True,
         name=None,
         ordered=False,
-        unique_column_names=False
+        unique_column_names=False,
+        metadata=None,
     )>
 
 

docs/source/schema_inference.rst

@@ -44,10 +44,11 @@ is a simple example:
         coerce=True,
         dtype=None,
         index=<Schema Index(name=None, type=DataType(int64))>,
-        strict=False
+        strict=False,
         name=None,
         ordered=False,
-        unique_column_names=False
+        unique_column_names=False,
+        metadata=None,
     )>
 
 

pandera/__init__.py

@@ -12,6 +12,7 @@
 from pandera.api.pandas.components import Column, Index, MultiIndex
 from pandera.api.pandas.model import DataFrameModel, SchemaModel
 from pandera.api.pandas.model_components import Field, check, dataframe_check
+from pandera.decorators import check_input, check_io, check_output, check_types
 from pandera.dtypes import (
     Bool,
     Category,
@@ -62,7 +63,6 @@
 import pandera.backends.pandas
 
 from pandera.schema_inference.pandas import infer_schema
-from pandera.decorators import check_input, check_io, check_output, check_types
 from pandera.version import __version__
 
 

pandera/accessors/pandas_accessor.py

@@ -43,7 +43,7 @@ class PanderaDataFrameAccessor(PanderaAccessor):
     def check_schema_type(schema):
         if not isinstance(schema, DataFrameSchema):
             raise TypeError(
-                f"schema arg must be a DataFrameSchema, found {type(schema)}"
+                f"schema arg must be a {DataFrameSchema}, found {type(schema)}"
             )
 
 
@@ -55,5 +55,5 @@ class PanderaSeriesAccessor(PanderaAccessor):
     def check_schema_type(schema):
         if not isinstance(schema, SeriesSchema):
             raise TypeError(
-                f"schema arg must be a SeriesSchema, found {type(schema)}"
+                f"schema arg must be a {SeriesSchema}, found {type(schema)}"
             )

pandera/accessors/pyspark_sql_accessor.py

@@ -1,22 +1,16 @@
-"""Custom accessor functionality for PySpark.Sql.
+"""Custom accessor functionality for PySpark.Sql. Register pyspark accessor for pandera schema metadata.
 """
 
 import warnings
-from functools import wraps
-from typing import Optional, Union
-
+from typing import Optional
 
 from pandera.api.pyspark.container import DataFrameSchema
 from pandera.api.pyspark.error_handler import ErrorHandler
 
-"""Register pyspark accessor for pandera schema metadata."""
-
-
-Schemas = Union[DataFrameSchema]
-Errors = Union[ErrorHandler]
+Schemas = DataFrameSchema  # type: ignore
+Errors = ErrorHandler  # type: ignore
 
 
-# Todo Refactor to create a seperate module for panderaAccessor
 class PanderaAccessor:
     """Pandera accessor for pyspark object."""
 
@@ -27,7 +21,7 @@ def __init__(self, pyspark_obj):
         self._errors: Optional[Errors] = None
 
     @staticmethod
-    def check_schema_type(schema: Schemas):
+    def check_schema_type(schema: Schemas):  # type: ignore
         """Abstract method for checking the schema type."""
         raise NotImplementedError
 
@@ -38,18 +32,18 @@ def add_schema(self, schema):
         return self._pyspark_obj
 
     @property
-    def schema(self) -> Optional[Schemas]:
+    def schema(self) -> Optional[Schemas]:  # type: ignore
         """Access schema metadata."""
         return self._schema
 
     @property
-    def errors(self) -> Optional[Errors]:
-        """Access errors data."""
+    def errors(self) -> Optional[Errors]:  # type: ignore
+        """Access errors details."""
         return self._errors
 
     @errors.setter
-    def errors(self, value: dict):
-        """Set errors data."""
+    def errors(self, value: Optional[Errors]):  # type: ignore
+        """Set errors details."""
         self._errors = value
 
 
@@ -133,4 +127,3 @@ def check_schema_type(schema):
 
 
 register_dataframe_accessor("pandera")(PanderaDataFrameAccessor)
-# register_series_accessor("pandera")(PanderaSeriesAccessor)