Register default backends when restoring state (#1550)

Signed-off-by: Alex Kment <[email protected]>

Author: alkment

pandera/api/base/schema.py

@@ -132,6 +132,10 @@ def _register_default_backends(self):
         overridden in those subclasses.
         """
 
+    def __setstate__(self, state):
+        self.__dict__ = state
+        self._register_default_backends()
+
 
 def inferred_schema_guard(method):
     """

tests/core/test_pandas_parallel.py

@@ -4,10 +4,11 @@
 from joblib import Parallel, delayed
 from pandera import Column, DataFrameSchema
 
+schema = DataFrameSchema({"a": Column("int64")}, coerce=True)
+
 
 def test_polars_parallel():
     def fn():
-        schema = DataFrameSchema({"a": Column("int64")}, coerce=True)
         return schema.validate(pd.DataFrame({"a": [1]}))
 
     results = Parallel(2)([delayed(fn)() for _ in range(10)])

tests/polars/test_polars_parallel.py

@@ -4,10 +4,11 @@
 from joblib import Parallel, delayed
 from pandera.polars import Column, DataFrameSchema
 
+schema = DataFrameSchema({"a": Column(pl.Int32)}, coerce=True)
+
 
 def test_polars_parallel():
     def fn():
-        schema = DataFrameSchema({"a": Column(pl.Int32)}, coerce=True)
         return schema.validate(pl.DataFrame({"a": [1]}))
 
     results = Parallel(2)([delayed(fn)() for _ in range(10)])