diff --git a/pandas/tests/indexing/multiindex/test_xs.py b/pandas/tests/indexing/multiindex/test_xs.py index b8380e3a95f2a..4a8e95b240e39 100644 --- a/pandas/tests/indexing/multiindex/test_xs.py +++ b/pandas/tests/indexing/multiindex/test_xs.py @@ -1,164 +1,213 @@ import numpy as np import pytest -from pandas.compat import StringIO, lrange, product as cart_product +from pandas.compat import lrange, product as cart_product -from pandas import DataFrame, Index, MultiIndex, concat, read_csv +from pandas import DataFrame, Index, MultiIndex, concat import pandas.core.common as com from pandas.util import testing as tm -class TestMultiIndexXs(object): - - def test_xs_multiindex(self): - - # GH2903 - columns = MultiIndex.from_tuples( - [('a', 'foo'), ('a', 'bar'), ('b', 'hello'), - ('b', 'world')], names=['lvl0', 'lvl1']) - df = DataFrame(np.random.randn(4, 4), columns=columns) - df.sort_index(axis=1, inplace=True) - result = df.xs('a', level='lvl0', axis=1) - expected = df.iloc[:, 0:2].loc[:, 'a'] - tm.assert_frame_equal(result, expected) - - result = df.xs('foo', level='lvl1', axis=1) - expected = df.iloc[:, 1:2].copy() - expected.columns = expected.columns.droplevel('lvl1') - tm.assert_frame_equal(result, expected) - - def test_xs(self, multiindex_dataframe_random_data): - frame = multiindex_dataframe_random_data - xs = frame.xs(('bar', 'two')) - xs2 = frame.loc[('bar', 'two')] - - tm.assert_series_equal(xs, xs2) - tm.assert_almost_equal(xs.values, frame.values[4]) - - # GH 6574 - # missing values in returned index should be preserrved - acc = [ - ('a', 'abcde', 1), - ('b', 'bbcde', 2), - ('y', 'yzcde', 25), - ('z', 'xbcde', 24), - ('z', None, 26), - ('z', 'zbcde', 25), - ('z', 'ybcde', 26), - ] - df = DataFrame(acc, - columns=['a1', 'a2', 'cnt']).set_index(['a1', 'a2']) - expected = DataFrame({'cnt': [24, 26, 25, 26]}, index=Index( - ['xbcde', np.nan, 'zbcde', 'ybcde'], name='a2')) - - result = df.xs('z', level='a1') - tm.assert_frame_equal(result, expected) - - def test_xs_with_duplicates(self, multiindex_dataframe_random_data): - # Issue #13719 - frame = multiindex_dataframe_random_data - df_dup = concat([frame] * 2) - assert df_dup.index.is_unique is False - expected = concat([frame.xs('one', level='second')] * 2) - tm.assert_frame_equal(df_dup.xs('one', level='second'), expected) - tm.assert_frame_equal(df_dup.xs(['one'], level=['second']), expected) - - def test_xs_level(self, multiindex_dataframe_random_data): - frame = multiindex_dataframe_random_data - result = frame.xs('two', level='second') - expected = frame[frame.index.get_level_values(1) == 'two'] - expected.index = expected.index.droplevel(1) - - tm.assert_frame_equal(result, expected) - - index = MultiIndex.from_tuples([('x', 'y', 'z'), ('a', 'b', 'c'), ( - 'p', 'q', 'r')]) - df = DataFrame(np.random.randn(3, 5), index=index) - result = df.xs('c', level=2) - expected = df[1:2] - expected.index = expected.index.droplevel(2) - tm.assert_frame_equal(result, expected) - - # this is a copy in 0.14 - result = frame.xs('two', level='second') - - # setting this will give a SettingWithCopyError - # as we are trying to write a view - def f(x): - x[:] = 10 - - pytest.raises(com.SettingWithCopyError, f, result) - - def test_xs_level_multiple(self): - text = """ A B C D E -one two three four -a b 10.0032 5 -0.5109 -2.3358 -0.4645 0.05076 0.3640 -a q 20 4 0.4473 1.4152 0.2834 1.00661 0.1744 -x q 30 3 -0.6662 -0.5243 -0.3580 0.89145 2.5838""" - - df = read_csv(StringIO(text), sep=r'\s+', engine='python') - - result = df.xs(('a', 4), level=['one', 'four']) - expected = df.xs('a').xs(4, level='four') - tm.assert_frame_equal(result, expected) - - # this is a copy in 0.14 - result = df.xs(('a', 4), level=['one', 'four']) - - # setting this will give a SettingWithCopyError - # as we are trying to write a view - def f(x): - x[:] = 10 - - pytest.raises(com.SettingWithCopyError, f, result) - - # GH2107 - dates = lrange(20111201, 20111205) - ids = 'abcde' - idx = MultiIndex.from_tuples([x for x in cart_product(dates, ids)]) - idx.names = ['date', 'secid'] - df = DataFrame(np.random.randn(len(idx), 3), idx, ['X', 'Y', 'Z']) - - rs = df.xs(20111201, level='date') - xp = df.loc[20111201, :] - tm.assert_frame_equal(rs, xp) - - def test_xs_level0(self): - text = """ A B C D E -one two three four -a b 10.0032 5 -0.5109 -2.3358 -0.4645 0.05076 0.3640 -a q 20 4 0.4473 1.4152 0.2834 1.00661 0.1744 -x q 30 3 -0.6662 -0.5243 -0.3580 0.89145 2.5838""" - - df = read_csv(StringIO(text), sep=r'\s+', engine='python') - - result = df.xs('a', level=0) - expected = df.xs('a') - assert len(result) == 2 - tm.assert_frame_equal(result, expected) - - def test_xs_level_series(self, multiindex_dataframe_random_data, - multiindex_year_month_day_dataframe_random_data): - frame = multiindex_dataframe_random_data - ymd = multiindex_year_month_day_dataframe_random_data - s = frame['A'] - result = s[:, 'two'] - expected = frame.xs('two', level=1)['A'] - tm.assert_series_equal(result, expected) - - s = ymd['A'] - result = s[2000, 5] - expected = ymd.loc[2000, 5]['A'] - tm.assert_series_equal(result, expected) - - # not implementing this for now - - pytest.raises(TypeError, s.__getitem__, (2000, slice(3, 4))) - - # result = s[2000, 3:4] - # lv =s.index.get_level_values(1) - # expected = s[(lv == 3) | (lv == 4)] - # expected.index = expected.index.droplevel(0) - # tm.assert_series_equal(result, expected) - - # can do this though +@pytest.fixture +def four_level_index_dataframe(): + arr = np.array([[-0.5109, -2.3358, -0.4645, 0.05076, 0.364], + [0.4473, 1.4152, 0.2834, 1.00661, 0.1744], + [-0.6662, -0.5243, -0.358, 0.89145, 2.5838]]) + index = MultiIndex( + levels=[['a', 'x'], ['b', 'q'], [10.0032, 20.0, 30.0], [3, 4, 5]], + codes=[[0, 0, 1], [0, 1, 1], [0, 1, 2], [2, 1, 0]], + names=['one', 'two', 'three', 'four']) + return DataFrame(arr, index=index, columns=list('ABCDE')) + + +@pytest.mark.parametrize('key, level, exp_arr, exp_index', [ + ('a', 'lvl0', lambda x: x[:, 0:2], Index(['bar', 'foo'], name='lvl1')), + ('foo', 'lvl1', lambda x: x[:, 1:2], Index(['a'], name='lvl0')) +]) +def test_xs_named_levels_axis_eq_1(key, level, exp_arr, exp_index): + # see gh-2903 + arr = np.random.randn(4, 4) + index = MultiIndex(levels=[['a', 'b'], ['bar', 'foo', 'hello', 'world']], + codes=[[0, 0, 1, 1], [0, 1, 2, 3]], + names=['lvl0', 'lvl1']) + df = DataFrame(arr, columns=index) + result = df.xs(key, level=level, axis=1) + expected = DataFrame(exp_arr(arr), columns=exp_index) + tm.assert_frame_equal(result, expected) + + +def test_xs_values(multiindex_dataframe_random_data): + df = multiindex_dataframe_random_data + result = df.xs(('bar', 'two')).values + expected = df.values[4] + tm.assert_almost_equal(result, expected) + + +def test_xs_loc_equality(multiindex_dataframe_random_data): + df = multiindex_dataframe_random_data + result = df.xs(('bar', 'two')) + expected = df.loc[('bar', 'two')] + tm.assert_series_equal(result, expected) + + +def test_xs_missing_values_in_index(): + # see gh-6574 + # missing values in returned index should be preserrved + acc = [ + ('a', 'abcde', 1), + ('b', 'bbcde', 2), + ('y', 'yzcde', 25), + ('z', 'xbcde', 24), + ('z', None, 26), + ('z', 'zbcde', 25), + ('z', 'ybcde', 26), + ] + df = DataFrame(acc, + columns=['a1', 'a2', 'cnt']).set_index(['a1', 'a2']) + expected = DataFrame({'cnt': [24, 26, 25, 26]}, index=Index( + ['xbcde', np.nan, 'zbcde', 'ybcde'], name='a2')) + + result = df.xs('z', level='a1') + tm.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize('key, level', [ + ('one', 'second'), + (['one'], ['second']) +]) +def test_xs_with_duplicates(key, level, multiindex_dataframe_random_data): + # see gh-13719 + frame = multiindex_dataframe_random_data + df = concat([frame] * 2) + assert df.index.is_unique is False + expected = concat([frame.xs('one', level='second')] * 2) + + result = df.xs(key, level=level) + tm.assert_frame_equal(result, expected) + + +def test_xs_level(multiindex_dataframe_random_data): + df = multiindex_dataframe_random_data + result = df.xs('two', level='second') + expected = df[df.index.get_level_values(1) == 'two'] + expected.index = Index(['foo', 'bar', 'baz', 'qux'], name='first') + tm.assert_frame_equal(result, expected) + + +def test_xs_level_eq_2(): + arr = np.random.randn(3, 5) + index = MultiIndex( + levels=[['a', 'p', 'x'], ['b', 'q', 'y'], ['c', 'r', 'z']], + codes=[[2, 0, 1], [2, 0, 1], [2, 0, 1]]) + df = DataFrame(arr, index=index) + expected = DataFrame(arr[1:2], index=[['a'], ['b']]) + result = df.xs('c', level=2) + tm.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize('indexer', [ + lambda df: df.xs(('a', 4), level=['one', 'four']), + lambda df: df.xs('a').xs(4, level='four') +]) +def test_xs_level_multiple(indexer, four_level_index_dataframe): + df = four_level_index_dataframe + expected_values = [[0.4473, 1.4152, 0.2834, 1.00661, 0.1744]] + expected_index = MultiIndex( + levels=[['q'], [20.0]], + codes=[[0], [0]], + names=['two', 'three']) + expected = DataFrame( + expected_values, index=expected_index, columns=list('ABCDE')) + result = indexer(df) + tm.assert_frame_equal(result, expected) + + +def test_xs_setting_with_copy_error(multiindex_dataframe_random_data): + # this is a copy in 0.14 + df = multiindex_dataframe_random_data + result = df.xs('two', level='second') + + # setting this will give a SettingWithCopyError + # as we are trying to write a view + msg = 'A value is trying to be set on a copy of a slice from a DataFrame' + with pytest.raises(com.SettingWithCopyError, match=msg): + result[:] = 10 + + +def test_xs_setting_with_copy_error_multiple(four_level_index_dataframe): + # this is a copy in 0.14 + df = four_level_index_dataframe + result = df.xs(('a', 4), level=['one', 'four']) + + # setting this will give a SettingWithCopyError + # as we are trying to write a view + msg = 'A value is trying to be set on a copy of a slice from a DataFrame' + with pytest.raises(com.SettingWithCopyError, match=msg): + result[:] = 10 + + +def test_xs_integer_key(): + # see gh-2107 + dates = lrange(20111201, 20111205) + ids = 'abcde' + index = MultiIndex.from_tuples( + [x for x in cart_product(dates, ids)], + names=['date', 'secid']) + df = DataFrame( + np.random.randn(len(index), 3), index, ['X', 'Y', 'Z']) + + result = df.xs(20111201, level='date') + expected = df.loc[20111201, :] + tm.assert_frame_equal(result, expected) + + +@pytest.mark.parametrize('indexer', [ + lambda df: df.xs('a', level=0), + lambda df: df.xs('a') +]) +def test_xs_level0(indexer, four_level_index_dataframe): + df = four_level_index_dataframe + expected_values = [[-0.5109, -2.3358, -0.4645, 0.05076, 0.364], + [0.4473, 1.4152, 0.2834, 1.00661, 0.1744]] + expected_index = MultiIndex( + levels=[['b', 'q'], [10.0032, 20.0], [4, 5]], + codes=[[0, 1], [0, 1], [1, 0]], + names=['two', 'three', 'four']) + expected = DataFrame( + expected_values, index=expected_index, columns=list('ABCDE')) + + result = indexer(df) + tm.assert_frame_equal(result, expected) + + +def test_xs_level_series(multiindex_dataframe_random_data): + # this test is not explicitly testing .xs functionality + # TODO: move to another module or refactor + df = multiindex_dataframe_random_data + s = df['A'] + result = s[:, 'two'] + expected = df.xs('two', level=1)['A'] + tm.assert_series_equal(result, expected) + + +def test_xs_level_series_ymd(multiindex_year_month_day_dataframe_random_data): + # this test is not explicitly testing .xs functionality + # TODO: move to another module or refactor + df = multiindex_year_month_day_dataframe_random_data + s = df['A'] + result = s[2000, 5] + expected = df.loc[2000, 5]['A'] + tm.assert_series_equal(result, expected) + + +def test_xs_level_series_slice_not_implemented( + multiindex_year_month_day_dataframe_random_data): + # this test is not explicitly testing .xs functionality + # TODO: move to another module or refactor + # not implementing this for now + df = multiindex_year_month_day_dataframe_random_data + s = df['A'] + + msg = r'\(2000, slice\(3, 4, None\)\)' + with pytest.raises(TypeError, match=msg): + s[2000, 3:4]