test_indexing.py

from datetime import timedelta

import numpy as np
import pytest

import pandas as pd
from pandas import (
    Categorical,
    CategoricalIndex,
    Index,
    IntervalIndex,
    MultiIndex,
    date_range,
)
from pandas.core.indexes.base import InvalidIndexError
import pandas.util.testing as tm


def test_slice_locs_partial(idx):
    sorted_idx, _ = idx.sortlevel(0)

    result = sorted_idx.slice_locs(("foo", "two"), ("qux", "one"))
    assert result == (1, 5)

    result = sorted_idx.slice_locs(None, ("qux", "one"))
    assert result == (0, 5)

    result = sorted_idx.slice_locs(("foo", "two"), None)
    assert result == (1, len(sorted_idx))

    result = sorted_idx.slice_locs("bar", "baz")
    assert result == (2, 4)


def test_slice_locs():
    df = tm.makeTimeDataFrame()
    stacked = df.stack()
    idx = stacked.index

    slob = slice(*idx.slice_locs(df.index[5], df.index[15]))
    sliced = stacked[slob]
    expected = df[5:16].stack()
    tm.assert_almost_equal(sliced.values, expected.values)

    slob = slice(
        *idx.slice_locs(
            df.index[5] + timedelta(seconds=30), df.index[15] - timedelta(seconds=30)
        )
    )
    sliced = stacked[slob]
    expected = df[6:15].stack()
    tm.assert_almost_equal(sliced.values, expected.values)


def test_slice_locs_with_type_mismatch():
    df = tm.makeTimeDataFrame()
    stacked = df.stack()
    idx = stacked.index
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs((1, 3))
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(df.index[5] + timedelta(seconds=30), (5, 2))
    df = tm.makeCustomDataframe(5, 5)
    stacked = df.stack()
    idx = stacked.index
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(timedelta(seconds=30))
    # TODO: Try creating a UnicodeDecodeError in exception message
    with pytest.raises(TypeError, match="^Level type mismatch"):
        idx.slice_locs(df.index[1], (16, "a"))


def test_slice_locs_not_sorted():
    index = MultiIndex(
        levels=[Index(np.arange(4)), Index(np.arange(4)), Index(np.arange(4))],
        codes=[
            np.array([0, 0, 1, 2, 2, 2, 3, 3]),
            np.array([0, 1, 0, 0, 0, 1, 0, 1]),
            np.array([1, 0, 1, 1, 0, 0, 1, 0]),
        ],
    )
    msg = "[Kk]ey length.*greater than MultiIndex lexsort depth"
    with pytest.raises(KeyError, match=msg):
        index.slice_locs((1, 0, 1), (2, 1, 0))

    # works
    sorted_index, _ = index.sortlevel(0)
    # should there be a test case here???
    sorted_index.slice_locs((1, 0, 1), (2, 1, 0))


def test_slice_locs_not_contained():
    # some searchsorted action

    index = MultiIndex(
        levels=[[0, 2, 4, 6], [0, 2, 4]],
        codes=[[0, 0, 0, 1, 1, 2, 3, 3, 3], [0, 1, 2, 1, 2, 2, 0, 1, 2]],
    )

    result = index.slice_locs((1, 0), (5, 2))
    assert result == (3, 6)

    result = index.slice_locs(1, 5)
    assert result == (3, 6)

    result = index.slice_locs((2, 2), (5, 2))
    assert result == (3, 6)

    result = index.slice_locs(2, 5)
    assert result == (3, 6)

    result = index.slice_locs((1, 0), (6, 3))
    assert result == (3, 8)

    result = index.slice_locs(-1, 10)
    assert result == (0, len(index))


def test_putmask_with_wrong_mask(idx):
    # GH18368

    msg = "putmask: mask and data must be the same size"
    with pytest.raises(ValueError, match=msg):
        idx.putmask(np.ones(len(idx) + 1, np.bool), 1)

    with pytest.raises(ValueError, match=msg):
        idx.putmask(np.ones(len(idx) - 1, np.bool), 1)

    with pytest.raises(ValueError, match=msg):
        idx.putmask("foo", 1)


def test_get_indexer():
    major_axis = Index(np.arange(4))
    minor_axis = Index(np.arange(2))

    major_codes = np.array([0, 0, 1, 2, 2, 3, 3], dtype=np.intp)
    minor_codes = np.array([0, 1, 0, 0, 1, 0, 1], dtype=np.intp)

    index = MultiIndex(
        levels=[major_axis, minor_axis], codes=[major_codes, minor_codes]
    )
    idx1 = index[:5]
    idx2 = index[[1, 3, 5]]

    r1 = idx1.get_indexer(idx2)
    tm.assert_almost_equal(r1, np.array([1, 3, -1], dtype=np.intp))

    r1 = idx2.get_indexer(idx1, method="pad")
    e1 = np.array([-1, 0, 0, 1, 1], dtype=np.intp)
    tm.assert_almost_equal(r1, e1)

    r2 = idx2.get_indexer(idx1[::-1], method="pad")
    tm.assert_almost_equal(r2, e1[::-1])

    rffill1 = idx2.get_indexer(idx1, method="ffill")
    tm.assert_almost_equal(r1, rffill1)

    r1 = idx2.get_indexer(idx1, method="backfill")
    e1 = np.array([0, 0, 1, 1, 2], dtype=np.intp)
    tm.assert_almost_equal(r1, e1)

    r2 = idx2.get_indexer(idx1[::-1], method="backfill")
    tm.assert_almost_equal(r2, e1[::-1])

    rbfill1 = idx2.get_indexer(idx1, method="bfill")
    tm.assert_almost_equal(r1, rbfill1)

    # pass non-MultiIndex
    r1 = idx1.get_indexer(idx2.values)
    rexp1 = idx1.get_indexer(idx2)
    tm.assert_almost_equal(r1, rexp1)

    r1 = idx1.get_indexer([1, 2, 3])
    assert (r1 == [-1, -1, -1]).all()

    # create index with duplicates
    idx1 = Index(list(range(10)) + list(range(10)))
    idx2 = Index(list(range(20)))

    msg = "Reindexing only valid with uniquely valued Index objects"
    with pytest.raises(InvalidIndexError, match=msg):
        idx1.get_indexer(idx2)


def test_get_indexer_and_fill():
    """ test getting an indexer for another index using the backfill method """
    mult_idx_1 = MultiIndex.from_product([[0], [0, 2, 3, 4]])
    mult_idx_2 = MultiIndex.from_product([[0], [1, 3, 4]])

    tm.assert_almost_equal(mult_idx_1.get_indexer(mult_idx_2), np.array([-1, 2, 3]))
    tm.assert_almost_equal(
        mult_idx_1.get_indexer(mult_idx_2, method="backfill"), np.array([1, 2, 3])
    )
    for method in ("bfill", "backfill"):
        tm.assert_almost_equal(
            mult_idx_1.get_indexer(mult_idx_2, method=method), np.array([1, 2, 3])
        )
    for method in ("ffill", "pad"):
        tm.assert_almost_equal(
            mult_idx_1.get_indexer(mult_idx_2, method=method), np.array([0, 2, 3])
        )


def test_get_indexer_three_or_more_levels():
    """ tests get_indexer() on MultiIndexes with 3+ levels

    visually, these are

    mult_idx_1:
      0: 1 2 5
      1:     7
      2:   4 5
      3:     7
      4:   6 5
      5:     7
      6: 3 2 5
      7:     7
      8:   4 5
      9:     7
     10:   6 5
     11:     7

    mult_idx_2:
      0: 1 1 8
      1: 1 5 9
      2: 1 6 7
      3: 2 1 6
      4: 2 7 6
      5: 2 7 8
      6: 3 6 8
    """
    mult_idx_1 = pd.MultiIndex.from_product([[1, 3], [2, 4, 6], [5, 7]])
    mult_idx_2 = pd.MultiIndex.from_tuples(
        [(1, 1, 8), (1, 5, 9), (1, 6, 7), (2, 1, 6), (2, 7, 7), (2, 7, 8), (3, 6, 8)]
    )
    # sanity check
    assert mult_idx_1.is_monotonic
    assert mult_idx_1.is_unique
    assert mult_idx_2.is_monotonic
    assert mult_idx_2.is_unique

    # show the relationships between the two
    assert mult_idx_2[0] < mult_idx_1[0]
    assert mult_idx_1[3] < mult_idx_2[1] < mult_idx_1[4]
    assert mult_idx_1[5] == mult_idx_2[2]
    assert mult_idx_1[5] < mult_idx_2[3] < mult_idx_1[6]
    assert mult_idx_1[5] < mult_idx_2[4] < mult_idx_1[6]
    assert mult_idx_1[5] < mult_idx_2[5] < mult_idx_1[6]
    assert mult_idx_1[-1] < mult_idx_2[6]

    indexer_no_fill = mult_idx_1.get_indexer(mult_idx_2)
    tm.assert_almost_equal(indexer_no_fill, np.array([-1, -1, 5, -1, -1, -1, -1]))

    # test with backfilling
    indexer_backfilled = mult_idx_1.get_indexer(mult_idx_2, method="backfill")
    tm.assert_almost_equal(indexer_backfilled, np.array([0, 4, 5, 6, 6, 6, -1]))

    # now, the same thing, but forward-filled (aka "padded")
    indexer_padded = mult_idx_1.get_indexer(mult_idx_2, method="pad")
    tm.assert_almost_equal(indexer_padded, np.array([-1, 3, 5, 5, 5, 5, 11]))

    # now, do the indexing in the other direction
    assert mult_idx_2[0] < mult_idx_1[0] < mult_idx_2[1]
    assert mult_idx_2[0] < mult_idx_1[1] < mult_idx_2[1]
    assert mult_idx_2[0] < mult_idx_1[2] < mult_idx_2[1]
    assert mult_idx_2[0] < mult_idx_1[3] < mult_idx_2[1]
    assert mult_idx_2[1] < mult_idx_1[4] < mult_idx_2[2]
    assert mult_idx_2[2] == mult_idx_1[5]
    assert mult_idx_2[5] < mult_idx_1[6] < mult_idx_2[6]
    assert mult_idx_2[5] < mult_idx_1[7] < mult_idx_2[6]
    assert mult_idx_2[5] < mult_idx_1[8] < mult_idx_2[6]
    assert mult_idx_2[5] < mult_idx_1[9] < mult_idx_2[6]
    assert mult_idx_2[5] < mult_idx_1[10] < mult_idx_2[6]
    assert mult_idx_2[5] < mult_idx_1[11] < mult_idx_2[6]

    tm.assert_almost_equal(
        mult_idx_2.get_indexer(mult_idx_1),
        np.array([-1, -1, -1, -1, -1, 2, -1, -1, -1, -1, -1, -1,]),
    )
    tm.assert_almost_equal(
        mult_idx_2.get_indexer(mult_idx_1, method="bfill"),
        np.array([1, 1, 1, 1, 2, 2, 6, 6, 6, 6, 6, 6,]),
    )
    tm.assert_almost_equal(
        mult_idx_2.get_indexer(mult_idx_1, method="pad"),
        np.array([0, 0, 0, 0, 1, 2, 5, 5, 5, 5, 5, 5,]),
    )


def test_get_indexer_backfill_with_carrying():
    """ tests a corner case with get_indexer() with MultiIndexes where, when we
    need to "carry" across levels, proper tuple ordering is respected

    the MultiIndexes used in this test, visually, are:

    mult_idx_1:
      0: 1 1 1 1
      1:       2
      2:     2 1
      3:       2
      4: 1 2 1 1
      5:       2
      6:     2 1
      7:       2
      8: 2 1 1 1
      9:       2
     10:     2 1
     11:       2
     12: 2 2 1 1
     13:       2
     14:     2 1
     15:       2

    mult_idx_2:
      0: 1 3 2 2
      1: 2 3 2 2
    """
    mult_idx_1 = pd.MultiIndex.from_product([[1, 2]] * 4)
    mult_idx_2 = pd.MultiIndex.from_tuples([(1, 3, 2, 2), (2, 3, 2, 2)])

    # show the tuple orderings, which get_indexer() should respect
    assert mult_idx_1[7] < mult_idx_2[0] < mult_idx_1[8]
    assert mult_idx_1[-1] < mult_idx_2[1]

    tm.assert_almost_equal(mult_idx_1.get_indexer(mult_idx_2), np.array([-1, -1]))
    tm.assert_almost_equal(
        mult_idx_1.get_indexer(mult_idx_2, method="bfill"), np.array([8, -1])
    )
    tm.assert_almost_equal(
        mult_idx_1.get_indexer(mult_idx_2, method="ffill"), np.array([7, 15])
    )


def test_get_indexer_nearest():
    midx = MultiIndex.from_tuples([("a", 1), ("b", 2)])
    msg = "method='nearest' not implemented yet for MultiIndex; see GitHub issue 9365"
    with pytest.raises(NotImplementedError, match=msg):
        midx.get_indexer(["a"], method="nearest")
    msg = "tolerance not implemented yet for MultiIndex"
    with pytest.raises(NotImplementedError, match=msg):
        midx.get_indexer(["a"], method="pad", tolerance=2)


def test_getitem(idx):
    # scalar
    assert idx[2] == ("bar", "one")

    # slice
    result = idx[2:5]
    expected = idx[[2, 3, 4]]
    assert result.equals(expected)

    # boolean
    result = idx[[True, False, True, False, True, True]]
    result2 = idx[np.array([True, False, True, False, True, True])]
    expected = idx[[0, 2, 4, 5]]
    assert result.equals(expected)
    assert result2.equals(expected)


def test_getitem_group_select(idx):
    sorted_idx, _ = idx.sortlevel(0)
    assert sorted_idx.get_loc("baz") == slice(3, 4)
    assert sorted_idx.get_loc("foo") == slice(0, 2)


def test_get_indexer_consistency(idx):
    # See GH 16819
    if isinstance(idx, IntervalIndex):
        pass

    if idx.is_unique or isinstance(idx, CategoricalIndex):
        indexer = idx.get_indexer(idx[0:2])
        assert isinstance(indexer, np.ndarray)
        assert indexer.dtype == np.intp
    else:
        e = "Reindexing only valid with uniquely valued Index objects"
        with pytest.raises(InvalidIndexError, match=e):
            idx.get_indexer(idx[0:2])

    indexer, _ = idx.get_indexer_non_unique(idx[0:2])
    assert isinstance(indexer, np.ndarray)
    assert indexer.dtype == np.intp


@pytest.mark.parametrize("ind1", [[True] * 5, pd.Index([True] * 5)])
@pytest.mark.parametrize(
    "ind2",
    [[True, False, True, False, False], pd.Index([True, False, True, False, False])],
)
def test_getitem_bool_index_all(ind1, ind2):
    # GH#22533
    idx = MultiIndex.from_tuples([(10, 1), (20, 2), (30, 3), (40, 4), (50, 5)])
    tm.assert_index_equal(idx[ind1], idx)

    expected = MultiIndex.from_tuples([(10, 1), (30, 3)])
    tm.assert_index_equal(idx[ind2], expected)


@pytest.mark.parametrize("ind1", [[True], pd.Index([True])])
@pytest.mark.parametrize("ind2", [[False], pd.Index([False])])
def test_getitem_bool_index_single(ind1, ind2):
    # GH#22533
    idx = MultiIndex.from_tuples([(10, 1)])
    tm.assert_index_equal(idx[ind1], idx)

    expected = pd.MultiIndex(
        levels=[np.array([], dtype=np.int64), np.array([], dtype=np.int64)],
        codes=[[], []],
    )
    tm.assert_index_equal(idx[ind2], expected)


def test_get_loc(idx):
    assert idx.get_loc(("foo", "two")) == 1
    assert idx.get_loc(("baz", "two")) == 3
    with pytest.raises(KeyError, match=r"^10$"):
        idx.get_loc(("bar", "two"))
    with pytest.raises(KeyError, match=r"^'quux'$"):
        idx.get_loc("quux")

    msg = "only the default get_loc method is currently supported for MultiIndex"
    with pytest.raises(NotImplementedError, match=msg):
        idx.get_loc("foo", method="nearest")

    # 3 levels
    index = MultiIndex(
        levels=[Index(np.arange(4)), Index(np.arange(4)), Index(np.arange(4))],
        codes=[
            np.array([0, 0, 1, 2, 2, 2, 3, 3]),
            np.array([0, 1, 0, 0, 0, 1, 0, 1]),
            np.array([1, 0, 1, 1, 0, 0, 1, 0]),
        ],
    )
    with pytest.raises(KeyError, match=r"^\(1, 1\)$"):
        index.get_loc((1, 1))
    assert index.get_loc((2, 0)) == slice(3, 5)


def test_get_loc_duplicates():
    index = Index([2, 2, 2, 2])
    result = index.get_loc(2)
    expected = slice(0, 4)
    assert result == expected
    # pytest.raises(Exception, index.get_loc, 2)

    index = Index(["c", "a", "a", "b", "b"])
    rs = index.get_loc("c")
    xp = 0
    assert rs == xp


def test_get_loc_level():
    index = MultiIndex(
        levels=[Index(np.arange(4)), Index(np.arange(4)), Index(np.arange(4))],
        codes=[
            np.array([0, 0, 1, 2, 2, 2, 3, 3]),
            np.array([0, 1, 0, 0, 0, 1, 0, 1]),
            np.array([1, 0, 1, 1, 0, 0, 1, 0]),
        ],
    )
    loc, new_index = index.get_loc_level((0, 1))
    expected = slice(1, 2)
    exp_index = index[expected].droplevel(0).droplevel(0)
    assert loc == expected
    assert new_index.equals(exp_index)

    loc, new_index = index.get_loc_level((0, 1, 0))
    expected = 1
    assert loc == expected
    assert new_index is None

    with pytest.raises(KeyError, match=r"^\(2, 2\)$"):
        index.get_loc_level((2, 2))
    # GH 22221: unused label
    with pytest.raises(KeyError, match=r"^2$"):
        index.drop(2).get_loc_level(2)
    # Unused label on unsorted level:
    with pytest.raises(KeyError, match=r"^2$"):
        index.drop(1, level=2).get_loc_level(2, level=2)

    index = MultiIndex(
        levels=[[2000], list(range(4))],
        codes=[np.array([0, 0, 0, 0]), np.array([0, 1, 2, 3])],
    )
    result, new_index = index.get_loc_level((2000, slice(None, None)))
    expected = slice(None, None)
    assert result == expected
    assert new_index.equals(index.droplevel(0))


@pytest.mark.parametrize("dtype1", [int, float, bool, str])
@pytest.mark.parametrize("dtype2", [int, float, bool, str])
def test_get_loc_multiple_dtypes(dtype1, dtype2):
    # GH 18520
    levels = [np.array([0, 1]).astype(dtype1), np.array([0, 1]).astype(dtype2)]
    idx = pd.MultiIndex.from_product(levels)
    assert idx.get_loc(idx[2]) == 2


@pytest.mark.parametrize("level", [0, 1])
@pytest.mark.parametrize("dtypes", [[int, float], [float, int]])
def test_get_loc_implicit_cast(level, dtypes):
    # GH 18818, GH 15994 : as flat index, cast int to float and vice-versa
    levels = [["a", "b"], ["c", "d"]]
    key = ["b", "d"]
    lev_dtype, key_dtype = dtypes
    levels[level] = np.array([0, 1], dtype=lev_dtype)
    key[level] = key_dtype(1)
    idx = MultiIndex.from_product(levels)
    assert idx.get_loc(tuple(key)) == 3


def test_get_loc_cast_bool():
    # GH 19086 : int is casted to bool, but not vice-versa
    levels = [[False, True], np.arange(2, dtype="int64")]
    idx = MultiIndex.from_product(levels)

    assert idx.get_loc((0, 1)) == 1
    assert idx.get_loc((1, 0)) == 2

    with pytest.raises(KeyError, match=r"^\(False, True\)$"):
        idx.get_loc((False, True))
    with pytest.raises(KeyError, match=r"^\(True, False\)$"):
        idx.get_loc((True, False))


@pytest.mark.parametrize("level", [0, 1])
def test_get_loc_nan(level, nulls_fixture):
    # GH 18485 : NaN in MultiIndex
    levels = [["a", "b"], ["c", "d"]]
    key = ["b", "d"]
    levels[level] = np.array([0, nulls_fixture], dtype=type(nulls_fixture))
    key[level] = nulls_fixture
    idx = MultiIndex.from_product(levels)
    assert idx.get_loc(tuple(key)) == 3


def test_get_loc_missing_nan():
    # GH 8569
    idx = MultiIndex.from_arrays([[1.0, 2.0], [3.0, 4.0]])
    assert isinstance(idx.get_loc(1), slice)
    with pytest.raises(KeyError, match=r"^3\.0$"):
        idx.get_loc(3)
    with pytest.raises(KeyError, match=r"^nan$"):
        idx.get_loc(np.nan)
    with pytest.raises(KeyError, match=r"^\[nan\]$"):
        idx.get_loc([np.nan])


def test_get_indexer_categorical_time():
    # https://github.com/pandas-dev/pandas/issues/21390
    midx = MultiIndex.from_product(
        [
            Categorical(["a", "b", "c"]),
            Categorical(date_range("2012-01-01", periods=3, freq="H")),
        ]
    )
    result = midx.get_indexer(midx)
    tm.assert_numpy_array_equal(result, np.arange(9, dtype=np.intp))


def test_timestamp_multiindex_indexer():
    # https://github.com/pandas-dev/pandas/issues/26944
    idx = pd.MultiIndex.from_product(
        [
            pd.date_range("2019-01-01T00:15:33", periods=100, freq="H", name="date"),
            ["x"],
            [3],
        ]
    )
    df = pd.DataFrame({"foo": np.arange(len(idx))}, idx)
    result = df.loc[pd.IndexSlice["2019-1-2":, "x", :], "foo"]
    qidx = pd.MultiIndex.from_product(
        [
            pd.date_range(
                start="2019-01-02T00:15:33",
                end="2019-01-05T02:15:33",
                freq="H",
                name="date",
            ),
            ["x"],
            [3],
        ]
    )
    should_be = pd.Series(data=np.arange(24, len(qidx) + 24), index=qidx, name="foo")
    tm.assert_series_equal(result, should_be)