import numpy as np
|
import pytest
|
|
import pandas as pd
|
from pandas import (
|
DataFrame,
|
Series,
|
)
|
import pandas._testing as tm
|
|
|
@pytest.fixture
|
def m():
|
return 5
|
|
|
@pytest.fixture
|
def n():
|
return 100
|
|
|
@pytest.fixture
|
def cols():
|
return ["jim", "joe", "jolie", "joline", "jolia"]
|
|
|
@pytest.fixture
|
def vals(n):
|
vals = [
|
np.random.default_rng(2).integers(0, 10, n),
|
np.random.default_rng(2).choice(list("abcdefghij"), n),
|
np.random.default_rng(2).choice(
|
pd.date_range("20141009", periods=10).tolist(), n
|
),
|
np.random.default_rng(2).choice(list("ZYXWVUTSRQ"), n),
|
np.random.default_rng(2).standard_normal(n),
|
]
|
vals = list(map(tuple, zip(*vals)))
|
return vals
|
|
|
@pytest.fixture
|
def keys(n, m, vals):
|
# bunch of keys for testing
|
keys = [
|
np.random.default_rng(2).integers(0, 11, m),
|
np.random.default_rng(2).choice(list("abcdefghijk"), m),
|
np.random.default_rng(2).choice(
|
pd.date_range("20141009", periods=11).tolist(), m
|
),
|
np.random.default_rng(2).choice(list("ZYXWVUTSRQP"), m),
|
]
|
keys = list(map(tuple, zip(*keys)))
|
keys += [t[:-1] for t in vals[:: n // m]]
|
return keys
|
|
|
# covers both unique index and non-unique index
|
@pytest.fixture
|
def df(vals, cols):
|
return DataFrame(vals, columns=cols)
|
|
|
@pytest.fixture
|
def a(df):
|
return pd.concat([df, df])
|
|
|
@pytest.fixture
|
def b(df, cols):
|
return df.drop_duplicates(subset=cols[:-1])
|
|
|
@pytest.mark.filterwarnings("ignore::pandas.errors.PerformanceWarning")
|
@pytest.mark.parametrize("lexsort_depth", list(range(5)))
|
@pytest.mark.parametrize("frame_fixture", ["a", "b"])
|
def test_multiindex_get_loc(request, lexsort_depth, keys, frame_fixture, cols):
|
# GH7724, GH2646
|
|
frame = request.getfixturevalue(frame_fixture)
|
if lexsort_depth == 0:
|
df = frame.copy(deep=False)
|
else:
|
df = frame.sort_values(by=cols[:lexsort_depth])
|
|
mi = df.set_index(cols[:-1])
|
assert not mi.index._lexsort_depth < lexsort_depth
|
for key in keys:
|
mask = np.ones(len(df), dtype=bool)
|
|
# test for all partials of this key
|
for i, k in enumerate(key):
|
mask &= df.iloc[:, i] == k
|
|
if not mask.any():
|
assert key[: i + 1] not in mi.index
|
continue
|
|
assert key[: i + 1] in mi.index
|
right = df[mask].copy(deep=False)
|
|
if i + 1 != len(key): # partial key
|
return_value = right.drop(cols[: i + 1], axis=1, inplace=True)
|
assert return_value is None
|
return_value = right.set_index(cols[i + 1 : -1], inplace=True)
|
assert return_value is None
|
tm.assert_frame_equal(mi.loc[key[: i + 1]], right)
|
|
else: # full key
|
return_value = right.set_index(cols[:-1], inplace=True)
|
assert return_value is None
|
if len(right) == 1: # single hit
|
right = Series(
|
right["jolia"].values, name=right.index[0], index=["jolia"]
|
)
|
tm.assert_series_equal(mi.loc[key[: i + 1]], right)
|
else: # multi hit
|
tm.assert_frame_equal(mi.loc[key[: i + 1]], right)
|
