Fix code to handle multi-dimensional index properly

Author: Si Wei How

doc/source/whatsnew/v0.25.0.rst

@@ -597,7 +597,7 @@ Indexing
 - Bug in which :meth:`DataFrame.to_csv` caused a segfault for a reindexed data frame, when the indices were single-level :class:`MultiIndex` (:issue:`26303`).
 - Fixed bug where assigning a :class:`arrays.PandasArray` to a :class:`pandas.core.frame.DataFrame` would raise error (:issue:`26390`)
 - Allow keyword arguments for callable local reference used in the :method:`DataFrame.query` string (:issue:`26426`)
-- Bug in assignment to multiple columns of a `DataFrame` when some of the columns do not exist (:issue:`13658`)
+- Bug in assignment to multiple columns of a :class:`DataFrame` when some of the columns do not exist (:issue:`13658`)
 
 
 Missing

pandas/core/indexing.py

@@ -1299,17 +1299,21 @@ def _convert_to_indexer(self, obj, axis=None, is_setter=False,
                 inds, = obj.nonzero()
                 return inds
             else:
-                # allow missing columns when setting
-                if is_setter and axis == 1:
-                    return [
-                        self._convert_to_indexer(k, axis=axis,
-                                                 is_setter=is_setter,
-                                                 raise_missing=raise_missing)
-                        for k in obj]
-                # When setting, missing keys are not allowed, even with .loc:
-                kwargs = {'raise_missing': True if is_setter else
-                          raise_missing}
-                return self._get_listlike_indexer(obj, axis, **kwargs)[1]
+                try:
+                    # When setting, missing keys are not allowed, even with
+                    # .loc:
+                    kwargs = {'raise_missing': True if is_setter else
+                              raise_missing}
+                    return self._get_listlike_indexer(obj, axis, **kwargs)[1]
+                except KeyError:
+                    # allow missing columns when setting
+                    if is_setter and axis == 1:
+                        return [
+                            self._convert_to_indexer(
+                                k, axis=axis, is_setter=is_setter,
+                                raise_missing=raise_missing)
+                            for k in obj]
+                    raise
         else:
             try:
                 return labels.get_loc(obj)

pandas/tests/frame/test_indexing.py

@@ -18,8 +18,7 @@
 from pandas.tests.frame.common import TestData
 import pandas.util.testing as tm
 from pandas.util.testing import (
-    assert_almost_equal, assert_frame_equal, assert_index_equal,
-    assert_numpy_array_equal, assert_series_equal)
+    assert_almost_equal, assert_frame_equal, assert_series_equal)
 
 from pandas.tseries.offsets import BDay
 
@@ -187,37 +186,41 @@ def test_setitem_list_of_tuples(self):
         assert_series_equal(result, expected)
 
     def test_setitem_list_all_missing_columns_scalar(self):
-        df = self.frame.copy()
-        df[['E', 'F']] = 1
-        df2 = self.frame.copy()
-        df2['E'] = 1
-        df2['F'] = 1
-        assert_frame_equal(df, df2)
+        # GH 26534
+        result = self.frame.copy()
+        result[['E', 'F']] = 1
+        expected = self.frame.copy()
+        expected['E'] = 1
+        expected['F'] = 1
+        assert_frame_equal(result, expected)
 
     def test_setitem_list_some_missing_columns_list(self):
-        df = self.frame.copy()
-        df[['A', 'E']] = [1, 2]
-        df2 = self.frame.copy()
-        df2['A'] = 1
-        df2['E'] = 2
-        assert_frame_equal(df, df2)
+        # GH 26534
+        result = self.frame.copy()
+        result[['A', 'E']] = [1, 2]
+        expected = self.frame.copy()
+        expected['A'] = 1
+        expected['E'] = 2
+        assert_frame_equal(result, expected)
 
     def test_setitem_list_some_missing_columns_dataframe(self):
-        df = self.frame.copy()
-        df[['A', 'E']] = self.frame[['B', 'C']]
-        df2 = self.frame.copy()
-        df2['A'] = self.frame['B']
-        df2['E'] = self.frame['C']
-        assert_frame_equal(df, df2)
+        # GH 26534
+        result = self.frame.copy()
+        result[['A', 'E']] = self.frame[['B', 'C']]
+        expected = self.frame.copy()
+        expected['A'] = self.frame['B']
+        expected['E'] = self.frame['C']
+        assert_frame_equal(result, expected)
 
     def test_setitem_list_some_missing_columns_2dlist(self):
-        df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
-        df[['B', 'C', 'D']] = [[7, 8, 9], [10, 11, 12], [13, 14, 15]]
+        # GH 26534
+        result = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
+        result[['B', 'C', 'D']] = [[7, 8, 9], [10, 11, 12], [13, 14, 15]]
         expected = pd.DataFrame([[1, 7, 8, 9],
                                  [3, 10, 11, 12],
                                  [5, 13, 14, 15]],
                                 columns=['A', 'B', 'C', 'D'])
-        tm.assert_frame_equal(df, expected)
+        tm.assert_frame_equal(result, expected)
 
     def test_setitem_mulit_index(self):
         # GH7655, test that assigning to a sub-frame of a frame

pandas/tests/indexing/test_loc.py

@@ -539,6 +539,7 @@ def test_loc_setitem_with_scalar_index(self, indexer, value):
         assert is_scalar(result) and result == 'Z'
 
     def test_loc_setitem_missing_columns_scalar_index_list_value(self):
+        # GH 26534
         df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
         df.loc[1, ['C', 'D']] = [7, 8]
         expected = pd.DataFrame([[1, 2, np.nan, np.nan],
@@ -548,6 +549,7 @@ def test_loc_setitem_missing_columns_scalar_index_list_value(self):
         tm.assert_frame_equal(df, expected)
 
     def test_loc_setitem_missing_columns_full_index_dataframe_value(self):
+        # GH 26534
         df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
         df2 = pd.DataFrame([[7, 8], [9, 10], [11, 12]], columns=['A', 'C'])
         df.loc[:, ['A', 'C']] = df2
@@ -556,6 +558,7 @@ def test_loc_setitem_missing_columns_full_index_dataframe_value(self):
         tm.assert_frame_equal(df, expected)
 
     def test_loc_setitem_missing_columns_list_index_scalar_value(self):
+        # GH 26534
         df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
         df.loc[[0, 2], ['B', 'C', 'D']] = 9
         expected = pd.DataFrame([[1, 9, 9, 9],
@@ -565,6 +568,7 @@ def test_loc_setitem_missing_columns_list_index_scalar_value(self):
         tm.assert_frame_equal(df, expected)
 
     def test_loc_setitem_missing_columns_range_index_2dlist_value(self):
+        # GH 26534
         df = pd.DataFrame([[1, 2], [3, 4], [5, 6]], columns=['A', 'B'])
         df.loc[1:3, ['B', 'C', 'D']] = [[7, 8, 9], [10, 11, 12]]
         expected = pd.DataFrame([[1, 2, np.nan, np.nan],