test_interval.py

import numpy as np
import pytest

import pandas as pd
from pandas import (
    DataFrame,
    IntervalIndex,
    Series,
)
import pandas._testing as tm


class TestIntervalIndex:
    @pytest.fixture
    def series_with_interval_index(self):
        return Series(np.arange(5), IntervalIndex.from_breaks(np.arange(6)))

    def test_getitem_with_scalar(self, series_with_interval_index, indexer_sl):
        ser = series_with_interval_index.copy()

        expected = ser.iloc[:3]
        tm.assert_series_equal(expected, indexer_sl(ser)[:3])
        tm.assert_series_equal(expected, indexer_sl(ser)[:2.5])
        tm.assert_series_equal(expected, indexer_sl(ser)[0.1:2.5])
        if indexer_sl is tm.loc:
            tm.assert_series_equal(expected, ser.loc[-1:3])

        expected = ser.iloc[1:4]
        tm.assert_series_equal(expected, indexer_sl(ser)[[1.5, 2.5, 3.5]])
        tm.assert_series_equal(expected, indexer_sl(ser)[[2, 3, 4]])
        tm.assert_series_equal(expected, indexer_sl(ser)[[1.5, 3, 4]])

        expected = ser.iloc[2:5]
        tm.assert_series_equal(expected, indexer_sl(ser)[ser >= 2])

    @pytest.mark.parametrize("direction", ["increasing", "decreasing"])
    def test_getitem_nonoverlapping_monotonic(self, direction, closed, indexer_sl):
        tpls = [(0, 1), (2, 3), (4, 5)]
        if direction == "decreasing":
            tpls = tpls[::-1]

        idx = IntervalIndex.from_tuples(tpls, closed=closed)
        ser = Series(list("abc"), idx)

        for key, expected in zip(idx.left, ser):
            if idx.closed_left:
                assert indexer_sl(ser)[key] == expected
            else:
                with pytest.raises(KeyError, match=str(key)):
                    indexer_sl(ser)[key]

        for key, expected in zip(idx.right, ser):
            if idx.closed_right:
                assert indexer_sl(ser)[key] == expected
            else:
                with pytest.raises(KeyError, match=str(key)):
                    indexer_sl(ser)[key]

        for key, expected in zip(idx.mid, ser):
            assert indexer_sl(ser)[key] == expected

    def test_getitem_non_matching(self, series_with_interval_index, indexer_sl):
        ser = series_with_interval_index.copy()

        # this is a departure from our current
        # indexing scheme, but simpler
        with pytest.raises(KeyError, match=r"\[-1\] not in index"):
            indexer_sl(ser)[[-1, 3, 4, 5]]

        with pytest.raises(KeyError, match=r"\[-1\] not in index"):
            indexer_sl(ser)[[-1, 3]]

    @pytest.mark.slow
    def test_loc_getitem_large_series(self):
        ser = Series(
            np.arange(1000000), index=IntervalIndex.from_breaks(np.arange(1000001))
        )

        result1 = ser.loc[:80000]
        result2 = ser.loc[0:80000]
        result3 = ser.loc[0:80000:1]
        tm.assert_series_equal(result1, result2)
        tm.assert_series_equal(result1, result3)

    def test_loc_getitem_frame(self):
        # CategoricalIndex with IntervalIndex categories
        df = DataFrame({"A": range(10)})
        ser = pd.cut(df.A, 5)
        df["B"] = ser
        df = df.set_index("B")

        result = df.loc[4]
        expected = df.iloc[4:6]
        tm.assert_frame_equal(result, expected)

        with pytest.raises(KeyError, match="10"):
            df.loc[10]

        # single list-like
        result = df.loc[[4]]
        expected = df.iloc[4:6]
        tm.assert_frame_equal(result, expected)

        # non-unique
        result = df.loc[[4, 5]]
        expected = df.take([4, 5, 4, 5])
        tm.assert_frame_equal(result, expected)

        with pytest.raises(KeyError, match=r"None of \[\[10\]\] are"):
            df.loc[[10]]

        # partial missing
        with pytest.raises(KeyError, match=r"\[10\] not in index"):
            df.loc[[10, 4]]

    def test_getitem_interval_with_nans(self, frame_or_series, indexer_sl):
        # GH#41831

        index = IntervalIndex([np.nan, np.nan])
        key = index[:-1]

        obj = frame_or_series(range(2), index=index)
        if frame_or_series is DataFrame and indexer_sl is tm.setitem:
            obj = obj.T

        result = indexer_sl(obj)[key]
        expected = obj

        tm.assert_equal(result, expected)


class TestIntervalIndexInsideMultiIndex:
    def test_mi_intervalindex_slicing_with_scalar(self):
        # GH#27456
        ii = IntervalIndex.from_arrays(
            [0, 1, 10, 11, 0, 1, 10, 11], [1, 2, 11, 12, 1, 2, 11, 12], name="MP"
        )
        idx = pd.MultiIndex.from_arrays(
            [
                pd.Index(["FC", "FC", "FC", "FC", "OWNER", "OWNER", "OWNER", "OWNER"]),
                pd.Index(
                    ["RID1", "RID1", "RID2", "RID2", "RID1", "RID1", "RID2", "RID2"]
                ),
                ii,
            ]
        )

        idx.names = ["Item", "RID", "MP"]
        df = DataFrame({"value": [1, 2, 3, 4, 5, 6, 7, 8]})
        df.index = idx

        query_df = DataFrame(
            {
                "Item": ["FC", "OWNER", "FC", "OWNER", "OWNER"],
                "RID": ["RID1", "RID1", "RID1", "RID2", "RID2"],
                "MP": [0.2, 1.5, 1.6, 11.1, 10.9],
            }
        )

        query_df = query_df.sort_index()

        idx = pd.MultiIndex.from_arrays([query_df.Item, query_df.RID, query_df.MP])
        query_df.index = idx
        result = df.value.loc[query_df.index]

        # the IntervalIndex level is indexed with floats, which map to
        #  the intervals containing them.  Matching the behavior we would get
        #  with _only_ an IntervalIndex, we get an IntervalIndex level back.
        sliced_level = ii.take([0, 1, 1, 3, 2])
        expected_index = pd.MultiIndex.from_arrays(
            [idx.get_level_values(0), idx.get_level_values(1), sliced_level]
        )
        expected = Series([1, 6, 2, 8, 7], index=expected_index, name="value")
        tm.assert_series_equal(result, expected)

    @pytest.mark.parametrize(
        "base, expected_result",
        [
            (10, Series([np.nan, 0], index=[np.nan, 1.0], dtype=np.float64)),
            (100, Series([np.nan, 0], index=[np.nan, 1.0], dtype=np.float64)),
            (101, Series([np.nan, 0], index=[np.nan, 1.0], dtype=np.float64)),
            (1010, Series([np.nan, 0], index=[np.nan, 1.0], dtype=np.float64)),
        ],
    )
    def test_reindex_behavior_with_interval_index(self, base, expected_result):
        # GH 51826
        left = np.arange(base)
        right = np.arange(1, base + 1)
        d = Series(range(base), index=IntervalIndex.from_arrays(left, right))
        result = d.reindex(index=[np.nan, 1.0])
        tm.assert_series_equal(result, expected_result)