ENH: Implement union method for Interval and IntervalArray (pandas-dev#21998)

Co-authored-by: Pedro Frigolet <[email protected]>

Author: St0rmie

doc/source/whatsnew/v3.0.0.rst

@@ -45,6 +45,7 @@ Other enhancements
 - :meth:`Series.cummin` and :meth:`Series.cummax` now supports :class:`CategoricalDtype` (:issue:`52335`)
 - :meth:`Series.plot` now correctly handle the ``ylabel`` parameter for pie charts, allowing for explicit control over the y-axis label (:issue:`58239`)
 - Added :meth:`Interval.intersection` and :meth:`IntervalArray.intersection` to calculate the intersection between interval-like objects (:issue:`21998`)
+- Added :meth:`Interval.union` and :meth:`IntervalArray.union` to calculate the union between interval-like objects (:issue:`21998`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
 
 .. ---------------------------------------------------------------------------

pandas/_libs/interval.pyx

@@ -695,6 +695,93 @@ cdef class Interval(IntervalMixin):
             closed = "neither"
         return Interval(ileft, iright, closed=closed)
 
+    def union(self, other):
+        """
+        Return the union of two intervals.
+
+        The union of two intervals are all the values in both, including
+        closed endpoints.
+
+        Parameters
+        ----------
+        other : Interval
+            Interval with which to create a union.
+
+        Returns
+        -------
+        np.array
+            numpy array with one interval if there is overlap between
+            the two intervals, with two intervals if there is no overlap.
+
+        See Also
+        --------
+        IntervalArray.union : The corresponding method for IntervalArray.
+
+        Examples
+        --------
+        >>> i0 = pd.Interval(0, 3, closed='right')
+        >>> i1 = pd.Interval(2, 4, closed='right')
+        >>> i0.union(i1)
+        [Interval(0, 4, closed='right')]
+
+        >>> i2 = pd.Interval(5, 8, closed='right')
+        >>> i0.union(i2)
+        [Interval(0, 3, closed='right') Interval(5, 8, closed='right')]
+
+        >>> i3 = pd.Interval(3, 5, closed='right')
+        >>> i0.union(i3)
+        [Interval(0, 5, closed='right')]
+        """
+        if not isinstance(other, Interval):
+            raise TypeError("`other` must be an Interval, "
+                            f"got {type(other).__name__}")
+
+        # if there is no overlap return the two intervals
+        # except if the two intervals share an endpoint were one side is closed
+        if not self.overlaps(other):
+            if(not(
+                (self.left == other.right and
+                    (self.closed_left or other.closed_right))
+                or
+                (self.right == other.left and
+                    (self.closed_right or other.closed_left)))):
+                if self.left < other.left:
+                    return np.array([self, other], dtype=object)
+                else:
+                    return np.array([other, self], dtype=object)
+
+        # Define left limit
+        if self.left < other.left:
+            uleft = self.left
+            lclosed = self.closed_left
+        elif self.left > other.left:
+            uleft = other.left
+            lclosed = other.closed_left
+        else:
+            uleft = self.left
+            lclosed = self.closed_left or other.closed_left
+
+        # Define right limit
+        if self.right > other.right:
+            uright = self.right
+            rclosed = self.closed_right
+        elif self.right < other.right:
+            uright = other.right
+            rclosed = other.closed_right
+        else:
+            uright = self.right
+            rclosed = self.closed_right or other.closed_right
+
+        if lclosed and rclosed:
+            closed = "both"
+        elif lclosed:
+            closed = "left"
+        elif rclosed:
+            closed = "right"
+        else:
+            closed = "neither"
+        return np.array([Interval(uleft, uright, closed=closed)], dtype=object)
+
 
 @cython.wraparound(False)
 @cython.boundscheck(False)

pandas/core/arrays/interval.py

@@ -1443,6 +1443,71 @@ def intersection(self, other):
             [interval.intersection(other) for interval in self], dtype=object
         )
 
+    _interval_shared_docs["union"] = textwrap.dedent(
+        """
+    Calculates union between each interval in the %(klass)s and a given
+    Interval.
+
+    The union of two intervals are all the values in both, including
+    closed endpoints.
+
+    Parameters
+    ----------
+    other : Interval
+        Interval to which to calculate the union.
+
+    Returns
+    -------
+    array
+        Array of arrays containing the unions between each interval and other.
+
+    See Also
+    --------
+    Interval.union : Calculate union between two Interval objects.
+
+    Examples
+    --------
+    %(examples)s
+
+    >>> intervals.union(pd.Interval(0, 2, 'right'))
+    [[Interval(0, 2, closed='right')]
+    [Interval(0, 5, closed='right')]
+    [Interval(0, 4, closed='right')]]
+
+    >>> intervals.union(pd.Interval(5, 8, closed='left'))
+    [array([Interval(0, 1, closed='right'), Interval(5, 8, closed='left')],
+    dtype=object)
+    array([Interval(1, 8, closed='neither')], dtype=object)
+    array([Interval(2, 4, closed='right'), Interval(5, 8, closed='left')],
+    dtype=object)]
+    """
+    )
+
+    @Appender(
+        _interval_shared_docs["union"]
+        % {
+            "klass": "IntervalArray",
+            "examples": textwrap.dedent(
+                """\
+        >>> data = [(0, 1), (1, 5), (2, 4)]
+        >>> intervals = pd.arrays.IntervalArray.from_tuples(data)
+        >>> intervals
+        <IntervalArray>
+        [(0, 1], (1, 3], (2, 4]]
+        Length: 3, dtype: interval[int64, right]
+        """
+            ),
+        }
+    )
+    def union(self, other):
+        if isinstance(other, (IntervalArray, ABCIntervalIndex)):
+            raise NotImplementedError
+        if not isinstance(other, Interval):
+            msg = f"`other` must be Interval-like, got {type(other).__name__}"
+            raise TypeError(msg)
+
+        return np.array([interval.union(other) for interval in self], dtype=object)
+
     # ---------------------------------------------------------------------
 
     @property

pandas/tests/arrays/interval/test_overlaps.py

@@ -133,3 +133,47 @@ def test_intersection_invalid_type(self, other):
         msg = f"`other` must be Interval-like, got {type(other).__name__}"
         with pytest.raises(TypeError, match=msg):
             interval_container.intersection(other)
+
+
+class TestUnion:
+    def test_union_interval_array(self):
+        interval = Interval(1, 8, "left")
+
+        tuples = [  # Intervals:
+            (1, 8),  # identical
+            (2, 4),  # nested
+            (0, 9),  # spanning
+            (4, 10),  # partial
+            (-5, 1),  # adjacent closed
+            (8, 10),  # adjacent open
+            (10, 15),  # disjoint
+        ]
+        interval_container = IntervalArray.from_tuples(tuples, "both")
+
+        expected = np.array(
+            [
+                np.array([Interval(1, 8, "both")], dtype=object),
+                np.array([Interval(1, 8, "left")], dtype=object),
+                np.array([Interval(0, 9, "both")], dtype=object),
+                np.array([Interval(1, 10, "both")], dtype=object),
+                np.array([Interval(-5, 8, "left")], dtype=object),
+                np.array([Interval(1, 10, "both")], dtype=object),
+                np.array(
+                    [Interval(1, 8, "left"), Interval(10, 15, "both")], dtype=object
+                ),
+            ],
+            dtype=object,
+        )
+        result = interval_container.union(interval)
+        tm.assert_numpy_array_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "other",
+        [10, True, "foo", Timedelta("1 day"), Timestamp("2018-01-01")],
+        ids=lambda x: type(x).__name__,
+    )
+    def test_union_invalid_type(self, other):
+        interval_container = IntervalArray.from_breaks(range(5))
+        msg = f"`other` must be Interval-like, got {type(other).__name__}"
+        with pytest.raises(TypeError, match=msg):
+            interval_container.union(other)

pandas/tests/scalar/interval/test_overlaps.py

@@ -166,3 +166,122 @@ def test_intersection_invalid_type(self, other):
         msg = f"`other` must be an Interval, got {type(other).__name__}"
         with pytest.raises(TypeError, match=msg):
             interval.intersection(other)
+
+
+class TestUnion:
+    def test_union_self(self):
+        interval = Interval(1, 8, "left")
+
+        result = interval.union(interval)
+
+        expected = np.array([interval], dtype=object)
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_union_include_limits(self):
+        other = Interval(1, 8, "left")
+
+        intervals = np.array(
+            [
+                Interval(7, 9, "left"),  # include left
+                Interval(0, 2, "right"),  # include right
+                Interval(1, 8, "right"),  # open limit
+            ]
+        )
+
+        expected = np.array(
+            [
+                np.array([Interval(1, 9, "left")], dtype=object),
+                np.array([Interval(0, 8, "neither")], dtype=object),
+                np.array([Interval(1, 8, "both")], dtype=object),
+            ],
+            dtype=object,
+        )
+
+        result = np.array([interval.union(other) for interval in intervals])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_union_overlapping(self):
+        other = Interval(1, 8, "left")
+
+        intervals = np.array(
+            [
+                Interval(2, 4, "both"),  # nested
+                Interval(0, 9, "both"),  # spanning
+                Interval(4, 10, "both"),  # partial
+            ]
+        )
+
+        expected = np.array(
+            [
+                np.array([Interval(1, 8, "left")], dtype=object),
+                np.array([Interval(0, 9, "both")], dtype=object),
+                np.array([Interval(1, 10, "both")], dtype=object),
+            ],
+            dtype=object,
+        )
+
+        result = np.array([interval.union(other) for interval in intervals])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_union_adjacent(self):
+        other = Interval(1, 8, "left")
+
+        intervals = np.array(
+            [
+                Interval(-5, 1, "both"),  # adjacent closed
+                Interval(8, 10, "both"),  # adjacent open
+                Interval(10, 15, "both"),  # disjoint
+            ]
+        )
+
+        expected = np.array(
+            [
+                np.array([Interval(-5, 8, "left")], dtype=object),
+                np.array([Interval(1, 10, "both")], dtype=object),
+                np.array([other, Interval(10, 15, "both")], dtype=object),
+            ],
+            dtype=object,
+        )
+
+        result = np.array(
+            [interval.union(other) for interval in intervals], dtype=object
+        )
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_union_timestamps(self):
+        year_2020 = Interval(
+            Timestamp("2020-01-01 00:00:00"),
+            Timestamp("2021-01-01 00:00:00"),
+            closed="left",
+        )
+
+        year_2021 = Interval(
+            Timestamp("2021-01-01 00:00:00"),
+            Timestamp("2022-01-01 00:00:00"),
+            closed="left",
+        )
+
+        expected = np.array(
+            [
+                Interval(
+                    Timestamp("2020-01-01 00:00:00"),
+                    Timestamp("2022-01-01 00:00:00"),
+                    closed="left",
+                )
+            ],
+            dtype=object,
+        )
+
+        result = year_2020.union(year_2021)
+        tm.assert_numpy_array_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "other",
+        [10, True, "foo", Timedelta("1 day"), Timestamp("2018-01-01")],
+        ids=lambda x: type(x).__name__,
+    )
+    def test_union_invalid_type(self, other):
+        interval = Interval(0, 1)
+        msg = f"`other` must be an Interval, got {type(other).__name__}"
+        with pytest.raises(TypeError, match=msg):
+            interval.union(other)