make finding coerce failure cases faster (#792)

* make finding coerce failure cases faster

* fix tests

* remove unneeded import

* fix tests, coverage

Author: cosmicBboy

pandera/dtypes.py

@@ -31,6 +31,10 @@ def coerce(self, data_container: Any):
         """Coerce data container to the data type."""
         raise NotImplementedError()
 
+    def coerce_value(self, value: Any):
+        """Coerce an value to a particular type."""
+        raise NotImplementedError()
+
     def try_coerce(self, data_container: Any):
         """Coerce data container to the data type,
         raises a `~pandera.errors.ParserError` if the coercion fails

pandera/engines/numpy_engine.py

@@ -51,6 +51,10 @@ def coerce(self, data_container: PandasObject) -> PandasObject:
             coerced.__str__()
         return coerced
 
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to a particular type."""
+        return self.type.type(value)
+
     def try_coerce(
         self, data_container: Union[PandasObject, np.ndarray]
     ) -> Union[PandasObject, np.ndarray]:

pandera/engines/pandas_engine.py

@@ -81,6 +81,10 @@ def coerce(self, data_container: PandasObject) -> PandasObject:
             coerced.__str__()
         return coerced
 
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to a particular type."""
+        return self.type.type(value)
+
     def try_coerce(self, data_container: PandasObject) -> PandasObject:
         try:
             return self.coerce(data_container)
@@ -197,6 +201,15 @@ class BOOL(DataType, dtypes.Bool):
     """Semantic representation of a :class:`pandas.BooleanDtype`."""
 
     type = pd.BooleanDtype()
+    _bool_like = frozenset({True, False})
+
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to specified datatime type."""
+        if value not in self._bool_like:
+            raise TypeError(
+                f"value {value} cannot be coerced to type {self.type}"
+            )
+        return super().coerce_value(value)
 
 
 ###############################################################################
@@ -416,6 +429,26 @@ def __init__(  # pylint:disable=super-init-not-called
             pd.CategoricalDtype(self.categories, self.ordered),
         )
 
+    def coerce(self, data_container: PandasObject) -> PandasObject:
+        """Pure coerce without catching exceptions."""
+        coerced = data_container.astype(self.type)
+        if (coerced.isna() & data_container.notna()).any(axis=None):
+            raise TypeError(
+                f"Data container cannot be coerced to type {self.type}"
+            )
+        if type(data_container).__module__.startswith("modin.pandas"):
+            # NOTE: this is a hack to enable catching of errors in modin
+            coerced.__str__()
+        return coerced
+
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to a particular type."""
+        if value not in self.categories:  # type: ignore
+            raise TypeError(
+                f"value {value} cannot be coerced to type {self.type}"
+            )
+        return value
+
     @classmethod
     def from_parametrized_dtype(
         cls, cat: Union[dtypes.Category, pd.CategoricalDtype]
@@ -589,6 +622,12 @@ def _to_datetime(col: pd.Series) -> pd.Series:
             return data_container.transform(_to_datetime)
         return _to_datetime(data_container)
 
+    def coerce_value(self, value: Any) -> Any:
+        """Coerce an value to specified datatime type."""
+        if value is pd.NaT:
+            return value
+        return super().coerce_value(value)
+
     @classmethod
     def from_parametrized_dtype(cls, pd_dtype: pd.DatetimeTZDtype):
         """Convert a :class:`pandas.DatetimeTZDtype` to

pandera/engines/utils.py

@@ -1,6 +1,5 @@
 """Engine module utilities."""
 
-import itertools
 from typing import Any, Union
 
 import numpy as np
@@ -23,53 +22,14 @@ def numpy_pandas_coercible(series: pd.Series, type_: Any) -> pd.Series:
 
     data_type = pandas_engine.Engine.dtype(type_)
 
-    def _bisect(series):
-        assert (
-            series.shape[0] >= 2
-        ), "cannot bisect a pandas Series of length < 2"
-        bisect_index = series.shape[0] // 2
-        return [series.iloc[:bisect_index], series.iloc[bisect_index:]]
-
-    def _coercible(series):
+    def _coercible(x):
         try:
-            data_type.coerce(series)
+            data_type.coerce_value(x)
             return True
         except Exception:  # pylint:disable=broad-except
             return False
 
-    search_list = [series] if series.size == 1 else _bisect(series)
-    failure_index = []
-    while search_list:
-        candidates = []
-        for _series in search_list:
-            if _series.shape[0] == 1 and not _coercible(_series):
-                # if series is reduced to a single value and isn't coercible,
-                # keep track of its index value.
-                failure_index.append(_series.index.item())
-            elif not _coercible(_series):
-                # if the series length > 1, add it to the candidates list
-                # to be further bisected
-                candidates.append(_series)
-
-        # the new search list is a flat list of bisected series views.
-        search_list = list(
-            itertools.chain.from_iterable([_bisect(c) for c in candidates])
-        )
-
-    # NOTE: this is a hack to support koalas. This needs to be thoroughly
-    # tested, right now koalas returns NA when a dtype value can't be coerced
-    # into the target dtype.
-    if type(series).__module__.startswith(
-        "databricks.koalas"
-    ):  # pragma: no cover
-        out = type(series)(
-            series.index.isin(failure_index).to_series().to_numpy(),  # type: ignore[union-attr]
-            index=series.index.values.to_numpy(),
-            name=series.name,
-        )
-        out.index.name = series.index.name
-        return out
-    return pd.Series(~series.index.isin(failure_index), index=series.index)
+    return series.map(_coercible)
 
 
 def numpy_pandas_coerce_failure_cases(

tests/core/test_pandas_engine.py

@@ -1,7 +1,9 @@
 """Test numpy engine."""
 
+import hypothesis.strategies as st
 import pandas as pd
 import pytest
+from hypothesis import given
 
 from pandera.engines import pandas_engine
 from pandera.errors import ParserError
@@ -42,3 +44,65 @@ def test_pandas_data_type_coerce(data_type):
         data_type().try_coerce(pd.Series(["1", "2", "a"]))
     except ParserError as exc:
         assert exc.failure_cases.shape[0] > 0
+
+
+CATEGORIES = ["A", "B", "C"]
+
+
+@given(st.lists(st.sampled_from(CATEGORIES), min_size=5))
+def test_pandas_category_dtype(data):
+    """Test pandas_engine.Category correctly coerces valid categorical data."""
+    data = pd.Series(data)
+    dtype = pandas_engine.Category(CATEGORIES)
+    coerced_data = dtype.coerce(data)
+    assert dtype.check(coerced_data.dtype)
+
+    for _, value in data.iteritems():
+        coerced_value = dtype.coerce_value(value)
+        assert coerced_value in CATEGORIES
+
+
+@given(st.lists(st.sampled_from(["X", "Y", "Z"]), min_size=5))
+def test_pandas_category_dtype_error(data):
+    """Test pandas_engine.Category raises TypeErrors on invalid data."""
+    data = pd.Series(data)
+    dtype = pandas_engine.Category(CATEGORIES)
+
+    with pytest.raises(TypeError):
+        dtype.coerce(data)
+
+    for _, value in data.iteritems():
+        with pytest.raises(TypeError):
+            dtype.coerce_value(value)
+
+
+@given(st.lists(st.sampled_from([1, 0, 1.0, 0.0, True, False]), min_size=5))
+def test_pandas_boolean_native_type(data):
+    """Test native pandas bool type correctly coerces valid bool-like data."""
+    data = pd.Series(data)
+    dtype = pandas_engine.Engine.dtype("boolean")
+
+    # the BooleanDtype can't handle Series of non-boolean, mixed dtypes
+    if data.dtype == "object":
+        with pytest.raises(TypeError):
+            dtype.coerce(data)
+    else:
+        coerced_data = dtype.coerce(data)
+        assert dtype.check(coerced_data.dtype)
+
+    for _, value in data.iteritems():
+        dtype.coerce_value(value)
+
+
+@given(st.lists(st.sampled_from(["A", "True", "False", 5, -1]), min_size=5))
+def test_pandas_boolean_native_type_error(data):
+    """Test native pandas bool type raises TypeErrors on non-bool-like data."""
+    data = pd.Series(data)
+    dtype = pandas_engine.Engine.dtype("boolean")
+
+    with pytest.raises(TypeError):
+        dtype.coerce(data)
+
+    for _, value in data.iteritems():
+        with pytest.raises(TypeError):
+            dtype.coerce_value(value)