ENH: IntervalArray comparisons

Author: jbrockmendel

pandas/core/arrays/_mixins.py

@@ -3,7 +3,6 @@
 import numpy as np
 
 from pandas._libs import lib
-from pandas.compat.numpy import function as nv
 from pandas.errors import AbstractMethodError
 from pandas.util._decorators import cache_readonly, doc
 from pandas.util._validators import validate_fillna_kwargs
@@ -139,7 +138,6 @@ def repeat(self: _T, repeats, axis=None) -> _T:
         --------
         numpy.ndarray.repeat
         """
-        nv.validate_repeat(tuple(), dict(axis=axis))
         new_data = self._ndarray.repeat(repeats, axis=axis)
         return self._from_backing_data(new_data)
 

pandas/core/arrays/interval.py

@@ -1,3 +1,4 @@
+import operator
 from operator import le, lt
 import textwrap
 from typing import TYPE_CHECKING, Optional, Tuple, Union, cast
@@ -48,7 +49,7 @@
 from pandas.core.construction import array, extract_array
 from pandas.core.indexers import check_array_indexer
 from pandas.core.indexes.base import ensure_index
-from pandas.core.ops import unpack_zerodim_and_defer
+from pandas.core.ops import invalid_comparison, unpack_zerodim_and_defer
 
 if TYPE_CHECKING:
     from pandas import Index
@@ -520,16 +521,15 @@ def __setitem__(self, key, value):
         self._left[key] = value_left
         self._right[key] = value_right
 
-    @unpack_zerodim_and_defer("__eq__")
-    def __eq__(self, other):
+    def _cmp_method(self, other, op):
         # ensure pandas array for list-like and eliminate non-interval scalars
         if is_list_like(other):
             if len(self) != len(other):
                 raise ValueError("Lengths must match to compare")
             other = array(other)
         elif not isinstance(other, Interval):
             # non-interval scalar -> no matches
-            return np.zeros(len(self), dtype=bool)
+            return invalid_comparison(self, other, op)
 
         # determine the dtype of the elements we want to compare
         if isinstance(other, Interval):
@@ -543,33 +543,87 @@ def __eq__(self, other):
             # extract intervals if we have interval categories with matching closed
             if is_interval_dtype(other_dtype):
                 if self.closed != other.categories.closed:
-                    return np.zeros(len(self), dtype=bool)
+                    return invalid_comparison(self, other, op)
                 other = other.categories.take(other.codes)
 
         # interval-like -> need same closed and matching endpoints
         if is_interval_dtype(other_dtype):
             if self.closed != other.closed:
-                return np.zeros(len(self), dtype=bool)
-            return (self._left == other.left) & (self._right == other.right)
+                return invalid_comparison(self, other, op)
+            if isinstance(other, Interval):
+                other = type(self)._from_sequence([other])
+                if self._combined.dtype.kind in ["m", "M"]:
+                    # Need to repeat bc we do not broadcast length-1
+                    # TODO: would be helpful to have a tile method to do
+                    #  this without copies
+                    other = other.repeat(len(self))
+            else:
+                other = type(self)(other)
+
+            if op is operator.eq:
+                return (self._combined[:, 0] == other._left) & (
+                    self._combined[:, 1] == other._right
+                )
+            elif op is operator.ne:
+                return (self._combined[:, 0] != other._left) | (
+                    self._combined[:, 1] != other._right
+                )
+            elif op is operator.gt:
+                return (self._combined[:, 0] > other._combined[:, 0]) | (
+                    (self._combined[:, 0] == other._left)
+                    & (self._combined[:, 1] > other._right)
+                )
+            elif op is operator.ge:
+                return (self == other) | (self > other)
+            elif op is operator.lt:
+                return (self._combined[:, 0] < other._combined[:, 0]) | (
+                    (self._combined[:, 0] == other._left)
+                    & (self._combined[:, 1] < other._right)
+                )
+            else:
+                # operator.lt
+                return (self == other) | (self < other)
 
         # non-interval/non-object dtype -> no matches
         if not is_object_dtype(other_dtype):
-            return np.zeros(len(self), dtype=bool)
+            return invalid_comparison(self, other, op)
 
         # object dtype -> iteratively check for intervals
-        result = np.zeros(len(self), dtype=bool)
-        for i, obj in enumerate(other):
-            # need object to be an Interval with same closed and endpoints
-            if (
-                isinstance(obj, Interval)
-                and self.closed == obj.closed
-                and self._left[i] == obj.left
-                and self._right[i] == obj.right
-            ):
-                result[i] = True
-
+        try:
+            result = np.zeros(len(self), dtype=bool)
+            for i, obj in enumerate(other):
+                result[i] = op(self[i], obj)
+        except TypeError:
+            # pd.NA
+            result = np.zeros(len(self), dtype=object)
+            for i, obj in enumerate(other):
+                result[i] = op(self[i], obj)
         return result
 
+    @unpack_zerodim_and_defer("__eq__")
+    def __eq__(self, other):
+        return self._cmp_method(other, operator.eq)
+
+    @unpack_zerodim_and_defer("__ne__")
+    def __ne__(self, other):
+        return self._cmp_method(other, operator.ne)
+
+    @unpack_zerodim_and_defer("__gt__")
+    def __gt__(self, other):
+        return self._cmp_method(other, operator.gt)
+
+    @unpack_zerodim_and_defer("__ge__")
+    def __ge__(self, other):
+        return self._cmp_method(other, operator.ge)
+
+    @unpack_zerodim_and_defer("__lt__")
+    def __lt__(self, other):
+        return self._cmp_method(other, operator.lt)
+
+    @unpack_zerodim_and_defer("__le__")
+    def __le__(self, other):
+        return self._cmp_method(other, operator.le)
+
     def fillna(self, value=None, method=None, limit=None):
         """
         Fill NA/NaN values using the specified method.

pandas/tests/arithmetic/test_interval.py

@@ -216,11 +216,6 @@ def test_compare_list_like_nan(self, op, array, nulls_fixture, request):
         result = op(array, other)
         expected = self.elementwise_comparison(op, array, other)
 
-        if nulls_fixture is pd.NA and array.dtype.subtype != "i8":
-            reason = "broken for non-integer IntervalArray; see GH 31882"
-            mark = pytest.mark.xfail(reason=reason)
-            request.node.add_marker(mark)
-
         tm.assert_numpy_array_equal(result, expected)
 
     @pytest.mark.parametrize(

pandas/tests/extension/base/methods.py

@@ -443,7 +443,7 @@ def test_repeat(self, data, repeats, as_series, use_numpy):
     @pytest.mark.parametrize(
         "repeats, kwargs, error, msg",
         [
-            (2, dict(axis=1), ValueError, "'axis"),
+            (2, dict(axis=1), ValueError, "axis"),
             (-1, dict(), ValueError, "negative"),
             ([1, 2], dict(), ValueError, "shape"),
             (2, dict(foo="bar"), TypeError, "'foo'"),

pandas/tests/indexes/interval/test_interval.py

@@ -579,9 +579,11 @@ def test_comparison(self):
         actual = self.index == self.index.left
         tm.assert_numpy_array_equal(actual, np.array([False, False]))
 
-        msg = (
-            "not supported between instances of 'int' and "
-            "'pandas._libs.interval.Interval'"
+        msg = "|".join(
+            [
+                "not supported between instances of 'int' and '.*.Interval'",
+                r"Invalid comparison between dtype=interval\[int64\] and ",
+            ]
         )
         with pytest.raises(TypeError, match=msg):
             self.index > 0