[WIP] ENH: .equals for Extension Arrays

Author: dwhu

pandas/core/arrays/base.py

@@ -78,6 +78,7 @@ class ExtensionArray:
     dropna
     factorize
     fillna
+    equals
     isna
     ravel
     repeat
@@ -350,6 +351,38 @@ def __iter__(self):
         for i in range(len(self)):
             yield self[i]
 
+    def __eq__(self, other: ABCExtensionArray) -> bool:
+        """
+        Whether the two arrays are equivalent.
+
+        Parameters
+        ----------
+        other: ExtensionArray
+            The array to compare to this array.
+
+        Returns
+        -------
+        bool
+        """
+
+        raise AbstractMethodError(self)
+
+    def __ne__(self, other: ABCExtensionArray) -> bool:
+        """
+        Whether the two arrays are not equivalent.
+
+        Parameters
+        ----------
+        other: ExtensionArray
+            The array to compare to this array.
+
+        Returns
+        -------
+        bool
+        """
+
+        raise AbstractMethodError(self)
+
     # ------------------------------------------------------------------------
     # Required attributes
     # ------------------------------------------------------------------------
@@ -657,6 +690,23 @@ def searchsorted(self, value, side="left", sorter=None):
         arr = self.astype(object)
         return arr.searchsorted(value, side=side, sorter=sorter)
 
+    def equals(self, other: ABCExtensionArray) -> bool:
+        """
+        Return if another array is equivalent to this array.
+
+        Parameters
+        ----------
+        other: ExtensionArray
+            Array to compare to this Array.
+
+        Returns
+        -------
+        boolean
+            Whether the arrays are equivalent.
+
+        """
+        return ((self == other) | (self.isna() == other.isna())).all()
+
     def _values_for_factorize(self) -> Tuple[np.ndarray, Any]:
         """
         Return an array and missing value suitable for factorization.

pandas/core/arrays/boolean.py

@@ -314,6 +314,20 @@ def __getitem__(self, item):
             return self._data[item]
         return type(self)(self._data[item], self._mask[item])
 
+    def __eq__(self, other):
+        if not isinstance(other, BooleanArray):
+            return NotImplemented
+        return (
+            hasattr(other, "_data")
+            and self._data == other._data
+            and hasattr(other, "_mask")
+            and self._mask == other._mask
+            and hasattr(other, "_dtype") & self._dtype == other._dtype
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def _coerce_to_ndarray(self, dtype=None, na_value: "Scalar" = libmissing.NA):
         """
         Coerce to an ndarray of object dtype or bool dtype (if force_bool=True).

pandas/core/arrays/categorical.py

@@ -2067,6 +2067,19 @@ def __setitem__(self, key, value):
         lindexer = self._maybe_coerce_indexer(lindexer)
         self._codes[key] = lindexer
 
+    def __eq__(self, other):
+        if not isinstance(other, Categorical):
+            return NotImplemented
+        return (
+            hasattr(other, "_codes")
+            and self._codes == other._codes
+            and hasattr(other, "_dtype")
+            and self._dtype == other._dtype
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def _reverse_indexer(self) -> Dict[Hashable, np.ndarray]:
         """
         Compute the inverse of a categorical, returning

pandas/core/arrays/integer.py

@@ -370,6 +370,19 @@ def __getitem__(self, item):
             return self._data[item]
         return type(self)(self._data[item], self._mask[item])
 
+    def __eq__(self, other):
+        if not isinstance(other, IntegerArray):
+            return NotImplemented
+        return (
+            hasattr(other, "_data")
+            and self._data == other._data
+            and hasattr(other, "_mask")
+            and self._mask == other._mask
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def _coerce_to_ndarray(self, dtype=None, na_value=lib._no_default):
         """
         coerce to an ndarary of object dtype

pandas/core/arrays/interval.py

@@ -547,6 +547,21 @@ def __setitem__(self, key, value):
         right.values[key] = value_right
         self._right = right
 
+    def __eq__(self, other):
+        if not isinstance(other, IntervalArray):
+            return NotImplementedError
+        return (
+            hasattr(other, "_left")
+            and self._left == other._left
+            and hasattr(other, "_right")
+            and self._right == other._right
+            and hasattr(other, "_closed")
+            and self._closed == other._closed
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def fillna(self, value=None, method=None, limit=None):
         """
         Fill NA/NaN values using the specified method.

pandas/tests/extension/json/array.py

@@ -107,6 +107,16 @@ def __setitem__(self, key, value):
                     assert isinstance(v, self.dtype.type)
                     self.data[k] = v
 
+    def __eq__(self, other):
+        return (
+            isinstance(other, JSONArray)
+            and hasattr(other, "data")
+            and self.data == other.data
+        )
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def __len__(self) -> int:
         return len(self.data)