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
from pandas.util.testing import assert_almost_equal


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]], sortorder=0)

    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)
    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)
    assert_almost_equal(r1, e1)

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

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

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

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

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

    # pass non-MultiIndex
    r1 = idx1.get_indexer(idx2.values)
    rexp1 = idx1.get_indexer(idx2)
    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_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
    dt_index = pd.date_range(start="1jan2019 00:15:33", periods=100,
                             freq="h", name="date")
    df = pd.DataFrame(index=dt_index,
                      data={'foo': range(len(dt_index)), 'bar': 'x', 'zaa': 3})
    df = df.reset_index().set_index(["date", "bar", "zaa"])
    assert df.loc[pd.IndexSlice['2019-1-2':, "x", :], 'foo'][-1] == 98