API (string dtype): implement hierarchy (NA > NaN, pyarrow > python) for consistent comparisons between different string dtypes

Author: jorisvandenbossche

pandas/core/arrays/arrow/array.py

@@ -33,7 +33,6 @@
     infer_dtype_from_scalar,
 )
 from pandas.core.dtypes.common import (
-    CategoricalDtype,
     is_array_like,
     is_bool_dtype,
     is_float_dtype,
@@ -730,9 +729,7 @@ def __setstate__(self, state) -> None:
 
     def _cmp_method(self, other, op) -> ArrowExtensionArray:
         pc_func = ARROW_CMP_FUNCS[op.__name__]
-        if isinstance(
-            other, (ArrowExtensionArray, np.ndarray, list, BaseMaskedArray)
-        ) or isinstance(getattr(other, "dtype", None), CategoricalDtype):
+        if isinstance(other, (ExtensionArray, np.ndarray, list)):
             try:
                 result = pc_func(self._pa_array, self._box_pa(other))
             except pa.ArrowNotImplementedError:

pandas/core/arrays/string_.py

@@ -1018,7 +1018,30 @@ def searchsorted(
         return super().searchsorted(value=value, side=side, sorter=sorter)
 
     def _cmp_method(self, other, op):
-        from pandas.arrays import BooleanArray
+        from pandas.arrays import (
+            ArrowExtensionArray,
+            BooleanArray,
+        )
+
+        if (
+            isinstance(other, BaseStringArray)
+            and self.dtype.na_value is not libmissing.NA
+            and other.dtype.na_value is libmissing.NA
+        ):
+            # NA has priority of NaN semantics
+            return NotImplemented
+
+        if isinstance(other, ArrowExtensionArray):
+            if isinstance(other, BaseStringArray):
+                # pyarrow storage has priority over python storage
+                # (except if we have NA semantics and other not)
+                if not (
+                    self.dtype.na_value is libmissing.NA
+                    and other.dtype.na_value is not libmissing.NA
+                ):
+                    return NotImplemented
+            else:
+                return NotImplemented
 
         if isinstance(other, StringArray):
             other = other._ndarray

pandas/core/arrays/string_arrow.py

@@ -473,6 +473,14 @@ def value_counts(self, dropna: bool = True) -> Series:
         return result
 
     def _cmp_method(self, other, op):
+        if (
+            isinstance(other, BaseStringArray)
+            and self.dtype.na_value is not libmissing.NA
+            and other.dtype.na_value is libmissing.NA
+        ):
+            # NA has priority of NaN semantics
+            return NotImplemented
+
         result = super()._cmp_method(other, op)
         if self.dtype.na_value is np.nan:
             if op == operator.ne:

pandas/tests/arrays/string_/test_string.py

@@ -45,6 +45,14 @@ def cls(dtype):
     return dtype.construct_array_type()
 
 
+DTYPE_HIERARCHY = [
+    pd.StringDtype("python", na_value=np.nan),
+    pd.StringDtype("pyarrow", na_value=np.nan),
+    pd.StringDtype("python", na_value=pd.NA),
+    pd.StringDtype("pyarrow", na_value=pd.NA),
+]
+
+
 def test_dtype_constructor():
     pytest.importorskip("pyarrow")
 
@@ -319,37 +327,43 @@ def test_comparison_methods_scalar_not_string(comparison_op, dtype):
         tm.assert_extension_array_equal(result, expected)
 
 
-def test_comparison_methods_array(comparison_op, dtype):
+def test_comparison_methods_array(comparison_op, dtype, dtype2):
     op_name = f"__{comparison_op.__name__}__"
 
     a = pd.array(["a", None, "c"], dtype=dtype)
-    other = [None, None, "c"]
-    result = getattr(a, op_name)(other)
-    if dtype.na_value is np.nan:
+    other = pd.array([None, None, "c"], dtype=dtype2)
+    result = comparison_op(a, other)
+
+    # ensure operation is commutative
+    result2 = comparison_op(other, a)
+    tm.assert_equal(result, result2)
+
+    if dtype.na_value is np.nan and dtype2.na_value is np.nan:
         if operator.ne == comparison_op:
             expected = np.array([True, True, False])
         else:
             expected = np.array([False, False, False])
             expected[-1] = getattr(other[-1], op_name)(a[-1])
         tm.assert_numpy_array_equal(result, expected)
 
-        result = getattr(a, op_name)(pd.NA)
-        if operator.ne == comparison_op:
-            expected = np.array([True, True, True])
+    else:
+        h1 = DTYPE_HIERARCHY.index(dtype)
+        h2 = DTYPE_HIERARCHY.index(dtype2)
+        max_dtype = DTYPE_HIERARCHY[max(h1, h2)]
+        if max_dtype.storage == "python":
+            expected_dtype = "boolean"
         else:
-            expected = np.array([False, False, False])
-        tm.assert_numpy_array_equal(result, expected)
+            expected_dtype = "bool[pyarrow]"
 
-    else:
-        expected_dtype = "boolean[pyarrow]" if dtype.storage == "pyarrow" else "boolean"
         expected = np.full(len(a), fill_value=None, dtype="object")
         expected[-1] = getattr(other[-1], op_name)(a[-1])
         expected = pd.array(expected, dtype=expected_dtype)
         tm.assert_extension_array_equal(result, expected)
 
-        result = getattr(a, op_name)(pd.NA)
-        expected = pd.array([None, None, None], dtype=expected_dtype)
-        tm.assert_extension_array_equal(result, expected)
+    # # with list
+    # other = [None, None, "c"]
+    # result3 = getattr(a, op_name)(other)
+    # tm.assert_equal(result, result3)
 
 
 def test_constructor_raises(cls):