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

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

Author: St0rmie

doc/source/whatsnew/v3.0.0.rst

@@ -46,6 +46,7 @@ Other enhancements
 - :meth:`DataFrame.pivot_table` and :func:`pivot_table` now allow the passing of keyword arguments to ``aggfunc`` through ``**kwargs`` (:issue:`57884`)
 - :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`)
 - Restore support for reading Stata 104-format and enable reading 103-format dta files (:issue:`58554`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
 

pandas/_libs/interval.pyx

@@ -620,6 +620,81 @@ cdef class Interval(IntervalMixin):
         # (simplifying the negation allows this to be done in less operations)
         return op1(self.left, other.right) and op2(other.left, self.right)
 
+    def intersection(self, other):
+        """
+        Return the intersection of two intervals.
+
+        The intersection of two intervals is the common points shared between both,
+        including closed endpoints. Open endpoints are not included.
+
+        Parameters
+        ----------
+        other : Interval
+            Interval to which to calculate the intersection.
+
+        Returns
+        -------
+        Interval or None
+            Interval containing the shared points and its closedness or None in
+            case there's no intersection.
+
+        See Also
+        --------
+        IntervalArray.intersection : The corresponding method for IntervalArray.
+
+        Examples
+        --------
+        >>> i0 = pd.Interval(0, 3, closed='right')
+        >>> i1 = pd.Interval(2, 4, closed='right')
+        >>> i0.intersection(i1)
+        Interval(2, 3, closed='right')
+
+        Intervals that have no intersection:
+
+        >>> i2 = pd.Interval(5, 8, closed='right')
+        >>> i0.intersection(i2)
+        None
+        """
+        if not isinstance(other, Interval):
+            raise TypeError("`other` must be an Interval, "
+                            f"got {type(other).__name__}")
+
+        # Define left limit
+        if self.left < other.left:
+            ileft = other.left
+            lclosed = other.closed_left
+        elif self.left > other.left:
+            ileft = self.left
+            lclosed = other.closed_left
+        else:
+            ileft = self.left
+            lclosed = self.closed_left and other.closed_left
+
+        # Define right limit
+        if self.right < other.right:
+            iright = self.right
+            rclosed = self.closed_right
+        elif self.right > other.right:
+            iright = other.right
+            rclosed = other.closed_right
+        else:
+            iright = self.right
+            rclosed = self.closed_right and other.closed_right
+
+        # No intersection if there is no overlap
+        if iright < ileft or (iright == ileft and not (lclosed and rclosed)):
+            return None
+
+        if lclosed and rclosed:
+            closed = "both"
+        elif lclosed:
+            closed = "left"
+        elif rclosed:
+            closed = "right"
+        else:
+            closed = "neither"
+        return Interval(ileft, iright, closed=closed)
+
 
 @cython.wraparound(False)
 @cython.boundscheck(False)

pandas/core/arrays/interval.py

@@ -1380,6 +1380,70 @@ def overlaps(self, other):
         # (simplifying the negation allows this to be done in less operations)
         return op1(self.left, other.right) & op2(other.left, self.right)
 
+    _interval_shared_docs["intersection"] = textwrap.dedent(
+        """
+    Calculates intersection between all intervals in the %(klass)s and a given
+    Interval.
+
+    The intersection of two intervals is the common points shared between both,
+    including closed endpoints. Open endpoints are not included.
+
+    Parameters
+    ----------
+    other : Interval
+        Interval to which to calculate the intersection.
+
+    Returns
+    -------
+    array
+        Array containing the Intersections between each interval and other.
+
+    See Also
+    --------
+    Interval.intersection : Calculate intersection between two Interval objects.
+
+    Examples
+    --------
+    %(examples)s
+
+    >>> intervals.intersection(pd.Interval(0, 2, 'right'))
+    array([Interval(0, 1, closed='right'), Interval(1, 2, closed='right'),
+           None], dtype=object)
+
+    Intersection with a single value:
+
+    >>> intervals.intersection(pd.Interval(5, 8, closed='left'))
+    array([None, Interval(5, 5, closed='both'), None], dtype=object)
+    """
+    )
+
+    @Appender(
+        _interval_shared_docs["intersection"]
+        % {
+            "klass": "IntervalArray",
+            "examples": textwrap.dedent(
+                """\
+        >>> data = [(0, 1), (1, 5), (2, 4)]
+        >>> intervals = pd.arrays.IntervalArray.from_tuples(data)
+        >>> intervals
+        <IntervalArray>
+        [(0, 1], (1, 5], (2, 4]]
+        Length: 3, dtype: interval[int64, right]
+        """
+            ),
+        }
+    )
+    def intersection(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.intersection(other) for interval in self], dtype=object
+        )
+
     # ---------------------------------------------------------------------
 
     @property

pandas/tests/arrays/interval/test_overlaps.py

@@ -92,3 +92,44 @@ def test_overlaps_invalid_type(self, constructor, other):
         msg = f"`other` must be Interval-like, got {type(other).__name__}"
         with pytest.raises(TypeError, match=msg):
             interval_container.overlaps(other)
+
+
+class TestIntersection:
+    def test_intersection_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(
+            [
+                Interval(1, 8, "left"),
+                Interval(2, 4, "both"),
+                Interval(1, 8, "left"),
+                Interval(4, 8, "left"),
+                Interval(1, 1, "both"),
+                None,
+                None,
+            ]
+        )
+        result = interval_container.intersection(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_intersection_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.intersection(other)

pandas/tests/scalar/interval/test_overlaps.py

@@ -1,10 +1,12 @@
+import numpy as np
 import pytest
 
 from pandas import (
     Interval,
     Timedelta,
     Timestamp,
 )
+import pandas._testing as tm
 
 
 @pytest.fixture(
@@ -65,3 +67,102 @@ def test_overlaps_invalid_type(self, other):
         msg = f"`other` must be an Interval, got {type(other).__name__}"
         with pytest.raises(TypeError, match=msg):
             interval.overlaps(other)
+
+
+class TestIntersection:
+    def test_intersection_self(self):
+        interval = Interval(1, 8, "left")
+        assert interval.intersection(interval) == interval
+
+    def test_intersection_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(
+            [
+                Interval(7, 8, "left"),
+                Interval(1, 2, "both"),
+                Interval(1, 8, "neither"),
+            ]
+        )
+
+        result = np.array([interval.intersection(other) for interval in intervals])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_intersection_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(
+            [
+                Interval(2, 4, "both"),
+                Interval(1, 8, "left"),
+                Interval(4, 8, "left"),
+            ]
+        )
+
+        result = np.array([interval.intersection(other) for interval in intervals])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_intersection_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(
+            [
+                Interval(1, 1, "both"),
+                None,
+                None,
+            ]
+        )
+
+        result = np.array([interval.intersection(other) for interval in intervals])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_intersection_timestamps(self):
+        year_2020 = Interval(
+            Timestamp("2020-01-01 00:00:00"),
+            Timestamp("2021-01-01 00:00:00"),
+            closed="left",
+        )
+
+        march_2020 = Interval(
+            Timestamp("2020-03-01 00:00:00"),
+            Timestamp("2020-04-01 00:00:00"),
+            closed="left",
+        )
+
+        result = year_2020.intersection(march_2020)
+        assert result == march_2020
+
+    @pytest.mark.parametrize(
+        "other",
+        [10, True, "foo", Timedelta("1 day"), Timestamp("2018-01-01")],
+        ids=lambda x: type(x).__name__,
+    )
+    def test_intersection_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.intersection(other)