From b0966b3358032c4c5c4e2e960c5a593d4cca61af Mon Sep 17 00:00:00 2001 From: h-vetinari <33685575+h-vetinari@users.noreply.github.com> Date: Fri, 19 Oct 2018 15:13:38 +0200 Subject: [PATCH] API: better error-handling for df.set_index (#22486) --- doc/source/whatsnew/v0.24.0.txt | 2 + pandas/core/frame.py | 55 +++++++++++------ pandas/tests/frame/conftest.py | 7 ++- pandas/tests/frame/test_alter_axes.py | 87 ++++++++++++++++++--------- 4 files changed, 101 insertions(+), 50 deletions(-) diff --git a/doc/source/whatsnew/v0.24.0.txt b/doc/source/whatsnew/v0.24.0.txt index 590410363b02c4..84efe2b9a53aed 100644 --- a/doc/source/whatsnew/v0.24.0.txt +++ b/doc/source/whatsnew/v0.24.0.txt @@ -806,6 +806,8 @@ Other API Changes - :class:`pandas.io.formats.style.Styler` supports a ``number-format`` property when using :meth:`~pandas.io.formats.style.Styler.to_excel` (:issue:`22015`) - :meth:`DataFrame.corr` and :meth:`Series.corr` now raise a ``ValueError`` along with a helpful error message instead of a ``KeyError`` when supplied with an invalid method (:issue:`22298`) - :meth:`shift` will now always return a copy, instead of the previous behaviour of returning self when shifting by 0 (:issue:`22397`) +- :meth:`DataFrame.set_index` now allows all one-dimensional list-likes, raises a ``TypeError`` for incorrect types, + has an improved ``KeyError`` message, and will not fail on duplicate column names with ``drop=True``. (:issue:`22484`) - Slicing a single row of a DataFrame with multiple ExtensionArrays of the same type now preserves the dtype, rather than coercing to object (:issue:`22784`) - :class:`DateOffset` attribute `_cacheable` and method `_should_cache` have been removed (:issue:`23118`) diff --git a/pandas/core/frame.py b/pandas/core/frame.py index e8acb37ad3b586..fd18d4860997d2 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -73,6 +73,7 @@ is_sequence, is_named_tuple) from pandas.core.dtypes.concat import _get_sliced_frame_result_type +from pandas.core.dtypes.generic import ABCSeries, ABCIndexClass, ABCMultiIndex from pandas.core.dtypes.missing import isna, notna from pandas.core import algorithms @@ -3988,6 +3989,25 @@ def set_index(self, keys, drop=True, append=False, inplace=False, if not isinstance(keys, list): keys = [keys] + missing = [] + for col in keys: + if (is_scalar(col) or isinstance(col, tuple)) and col in self: + # tuples can be both column keys or list-likes + # if they are valid column keys, everything is fine + continue + elif is_scalar(col) and col not in self: + # tuples that are not column keys are considered list-like, + # not considered missing + missing.append(col) + elif (not is_list_like(col, allow_sets=False) + or getattr(col, 'ndim', 1) > 1): + raise TypeError('The parameter "keys" may only contain a ' + 'combination of valid column keys and ' + 'one-dimensional list-likes') + + if missing: + raise KeyError('{}'.format(missing)) + if inplace: frame = self else: @@ -3997,7 +4017,7 @@ def set_index(self, keys, drop=True, append=False, inplace=False, names = [] if append: names = [x for x in self.index.names] - if isinstance(self.index, MultiIndex): + if isinstance(self.index, ABCMultiIndex): for i in range(self.index.nlevels): arrays.append(self.index._get_level_values(i)) else: @@ -4005,29 +4025,29 @@ def set_index(self, keys, drop=True, append=False, inplace=False, to_remove = [] for col in keys: - if isinstance(col, MultiIndex): - # append all but the last column so we don't have to modify - # the end of this loop - for n in range(col.nlevels - 1): + if isinstance(col, ABCMultiIndex): + for n in range(col.nlevels): arrays.append(col._get_level_values(n)) - - level = col._get_level_values(col.nlevels - 1) names.extend(col.names) - elif isinstance(col, Series): - level = col._values - names.append(col.name) - elif isinstance(col, Index): - level = col + elif isinstance(col, (ABCIndexClass, ABCSeries)): + # if Index then not MultiIndex (treated above) + arrays.append(col) names.append(col.name) - elif isinstance(col, (list, np.ndarray, Index)): - level = col + elif isinstance(col, (list, np.ndarray)): + arrays.append(col) + names.append(None) + elif (is_list_like(col) + and not (isinstance(col, tuple) and col in self)): + # all other list-likes (but avoid valid column keys) + col = list(col) # ensure iterator do not get read twice etc. + arrays.append(col) names.append(None) + # from here, col can only be a column label else: - level = frame[col]._values + arrays.append(frame[col]._values) names.append(col) if drop: to_remove.append(col) - arrays.append(level) index = ensure_index_from_sequences(arrays, names) @@ -4036,7 +4056,8 @@ def set_index(self, keys, drop=True, append=False, inplace=False, raise ValueError('Index has duplicate keys: {dup}'.format( dup=duplicates)) - for c in to_remove: + # use set to handle duplicate column names gracefully in case of drop + for c in set(to_remove): del frame[c] # clear up memory usage diff --git a/pandas/tests/frame/conftest.py b/pandas/tests/frame/conftest.py index 348331fc0ccdff..ec66fb6bf55d28 100644 --- a/pandas/tests/frame/conftest.py +++ b/pandas/tests/frame/conftest.py @@ -211,12 +211,13 @@ def frame_of_index_cols(): """ Fixture for DataFrame of columns that can be used for indexing - Columns are ['A', 'B', 'C', 'D', 'E']; 'A' & 'B' contain duplicates (but - are jointly unique), the rest are unique. + Columns are ['A', 'B', 'C', 'D', 'E', ('tuple', 'as', 'label')]; + 'A' & 'B' contain duplicates (but are jointly unique), the rest are unique. """ df = DataFrame({'A': ['foo', 'foo', 'foo', 'bar', 'bar'], 'B': ['one', 'two', 'three', 'one', 'two'], 'C': ['a', 'b', 'c', 'd', 'e'], 'D': np.random.randn(5), - 'E': np.random.randn(5)}) + 'E': np.random.randn(5), + ('tuple', 'as', 'label'): np.random.randn(5)}) return df diff --git a/pandas/tests/frame/test_alter_axes.py b/pandas/tests/frame/test_alter_axes.py index 4e61c9c62266df..fdeae61f96b93d 100644 --- a/pandas/tests/frame/test_alter_axes.py +++ b/pandas/tests/frame/test_alter_axes.py @@ -49,7 +49,8 @@ def test_set_index_cast(self): tm.assert_frame_equal(df, df2) # A has duplicate values, C does not - @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B']]) + @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B'], + ('tuple', 'as', 'label')]) @pytest.mark.parametrize('inplace', [True, False]) @pytest.mark.parametrize('drop', [True, False]) def test_set_index_drop_inplace(self, frame_of_index_cols, @@ -72,7 +73,8 @@ def test_set_index_drop_inplace(self, frame_of_index_cols, tm.assert_frame_equal(result, expected) # A has duplicate values, C does not - @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B']]) + @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B'], + ('tuple', 'as', 'label')]) @pytest.mark.parametrize('drop', [True, False]) def test_set_index_append(self, frame_of_index_cols, drop, keys): df = frame_of_index_cols @@ -88,7 +90,8 @@ def test_set_index_append(self, frame_of_index_cols, drop, keys): tm.assert_frame_equal(result, expected) # A has duplicate values, C does not - @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B']]) + @pytest.mark.parametrize('keys', ['A', 'C', ['A', 'B'], + ('tuple', 'as', 'label')]) @pytest.mark.parametrize('drop', [True, False]) def test_set_index_append_to_multiindex(self, frame_of_index_cols, drop, keys): @@ -114,8 +117,10 @@ def test_set_index_after_mutation(self): tm.assert_frame_equal(result, expected) # MultiIndex constructor does not work directly on Series -> lambda + # Add list-of-list constructor because list is ambiguous -> lambda # also test index name if append=True (name is duplicate here for B) @pytest.mark.parametrize('box', [Series, Index, np.array, + list, tuple, iter, lambda x: [list(x)], lambda x: MultiIndex.from_arrays([x])]) @pytest.mark.parametrize('append, index_name', [(True, None), (True, 'B'), (True, 'test'), (False, None)]) @@ -126,21 +131,29 @@ def test_set_index_pass_single_array(self, frame_of_index_cols, df.index.name = index_name key = box(df['B']) - # np.array and list "forget" the name of B - name = [None if box in [np.array, list] else 'B'] + if box == list: + # list of strings gets interpreted as list of keys + msg = "['one', 'two', 'three', 'one', 'two']" + with tm.assert_raises_regex(KeyError, msg): + df.set_index(key, drop=drop, append=append) + else: + # np.array/tuple/iter/list-of-list "forget" the name of B + name_mi = getattr(key, 'names', None) + name = [getattr(key, 'name', None)] if name_mi is None else name_mi - result = df.set_index(key, drop=drop, append=append) + result = df.set_index(key, drop=drop, append=append) - # only valid column keys are dropped - # since B is always passed as array above, nothing is dropped - expected = df.set_index(['B'], drop=False, append=append) - expected.index.names = [index_name] + name if append else name + # only valid column keys are dropped + # since B is always passed as array above, nothing is dropped + expected = df.set_index(['B'], drop=False, append=append) + expected.index.names = [index_name] + name if append else name - tm.assert_frame_equal(result, expected) + tm.assert_frame_equal(result, expected) # MultiIndex constructor does not work directly on Series -> lambda # also test index name if append=True (name is duplicate here for A & B) - @pytest.mark.parametrize('box', [Series, Index, np.array, list, + @pytest.mark.parametrize('box', [Series, Index, np.array, + list, tuple, iter, lambda x: MultiIndex.from_arrays([x])]) @pytest.mark.parametrize('append, index_name', [(True, None), (True, 'A'), (True, 'B'), @@ -152,8 +165,8 @@ def test_set_index_pass_arrays(self, frame_of_index_cols, df.index.name = index_name keys = ['A', box(df['B'])] - # np.array and list "forget" the name of B - names = ['A', None if box in [np.array, list] else 'B'] + # np.array/list/tuple/iter "forget" the name of B + names = ['A', None if box in [np.array, list, tuple, iter] else 'B'] result = df.set_index(keys, drop=drop, append=append) @@ -168,10 +181,12 @@ def test_set_index_pass_arrays(self, frame_of_index_cols, # MultiIndex constructor does not work directly on Series -> lambda # We also emulate a "constructor" for the label -> lambda # also test index name if append=True (name is duplicate here for A) - @pytest.mark.parametrize('box2', [Series, Index, np.array, list, + @pytest.mark.parametrize('box2', [Series, Index, np.array, + list, tuple, iter, lambda x: MultiIndex.from_arrays([x]), lambda x: x.name]) - @pytest.mark.parametrize('box1', [Series, Index, np.array, list, + @pytest.mark.parametrize('box1', [Series, Index, np.array, + list, tuple, iter, lambda x: MultiIndex.from_arrays([x]), lambda x: x.name]) @pytest.mark.parametrize('append, index_name', [(True, None), @@ -183,21 +198,22 @@ def test_set_index_pass_arrays_duplicate(self, frame_of_index_cols, drop, df.index.name = index_name keys = [box1(df['A']), box2(df['A'])] + result = df.set_index(keys, drop=drop, append=append) - # == gives ambiguous Boolean for Series - if drop and keys[0] is 'A' and keys[1] is 'A': - with tm.assert_raises_regex(KeyError, '.*'): - df.set_index(keys, drop=drop, append=append) - else: - result = df.set_index(keys, drop=drop, append=append) + # if either box was iter, the content has been consumed; re-read it + keys = [box1(df['A']), box2(df['A'])] - # to test against already-tested behavior, we add sequentially, - # hence second append always True; must wrap in list, otherwise - # list-box will be illegal - expected = df.set_index([keys[0]], drop=drop, append=append) - expected = expected.set_index([keys[1]], drop=drop, append=True) + # need to adapt first drop for case that both keys are 'A' -- + # cannot drop the same column twice; + # use "is" because == would give ambiguous Boolean error for containers + first_drop = False if (keys[0] is 'A' and keys[1] is 'A') else drop - tm.assert_frame_equal(result, expected) + # to test against already-tested behaviour, we add sequentially, + # hence second append always True; must wrap keys in list, otherwise + # box = list would be illegal + expected = df.set_index([keys[0]], drop=first_drop, append=append) + expected = expected.set_index([keys[1]], drop=drop, append=True) + tm.assert_frame_equal(result, expected) @pytest.mark.parametrize('append', [True, False]) @pytest.mark.parametrize('drop', [True, False]) @@ -229,13 +245,24 @@ def test_set_index_verify_integrity(self, frame_of_index_cols): def test_set_index_raise(self, frame_of_index_cols, drop, append): df = frame_of_index_cols - with tm.assert_raises_regex(KeyError, '.*'): # column names are A-E + with tm.assert_raises_regex(KeyError, "['foo', 'bar', 'baz']"): + # column names are A-E, as well as one tuple df.set_index(['foo', 'bar', 'baz'], drop=drop, append=append) # non-existent key in list with arrays - with tm.assert_raises_regex(KeyError, '.*'): + with tm.assert_raises_regex(KeyError, 'X'): df.set_index([df['A'], df['B'], 'X'], drop=drop, append=append) + msg = 'The parameter "keys" may only contain a combination of.*' + # forbidden type, e.g. set + with tm.assert_raises_regex(TypeError, msg): + df.set_index(set(df['A']), drop=drop, append=append) + + # forbidden type in list, e.g. set + with tm.assert_raises_regex(TypeError, msg): + df.set_index(['A', df['A'], set(df['A'])], + drop=drop, append=append) + def test_construction_with_categorical_index(self): ci = tm.makeCategoricalIndex(10) ci.name = 'B'