fix: pandas pyarrow string validation (#1636)

Signed-off-by: Ajith Aravind <ajith.aravind100@gmail.com>
unionai-oss · May 14, 2024 · c815a6d · c815a6d
1 parent 18aacca
commit c815a6d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/pandera/engines/pandas_engine.py b/pandera/engines/pandas_engine.py
@@ -1660,7 +1660,9 @@ class ArrowInt8(ArrowInt16):
         type = pd.ArrowDtype(pyarrow.int8())
         bit_width: int = 8
 
-    @Engine.register_dtype(equivalents=[pyarrow.string])
+    @Engine.register_dtype(
+        equivalents=[pyarrow.string, pd.ArrowDtype(pyarrow.string())]
+    )
     @immutable
     class ArrowString(DataType, dtypes.String):
         """Semantic representation of a :class:`pyarrow.string`."""

diff --git a/tests/core/test_dtypes.py b/tests/core/test_dtypes.py
@@ -311,6 +311,12 @@ def test_invalid_pandas_extension_dtype():
 
 def test_check_equivalent(dtype: Any, pd_dtype: Any):
     """Test that a pandas-compatible dtype can be validated by check()."""
+    if (
+        pandas_engine.PYARROW_INSTALLED
+        and pandas_engine.PANDAS_2_0_0_PLUS
+        and dtype == "string[pyarrow]"
+    ):
+        pytest.skip("`string[pyarrow]` gets parsed to type `string` by pandas")
     actual_dtype = pandas_engine.Engine.dtype(pd_dtype)
     expected_dtype = pandas_engine.Engine.dtype(dtype)
     assert actual_dtype.check(expected_dtype)