Pass axis to _ensure_listlike_indexer

Author: howsiwei

pandas/core/frame.py

@@ -2686,7 +2686,7 @@ def _setitem_array(self, key, value):
                 for k1, k2 in zip(key, value.columns):
                     self[k1] = value[k2]
             else:
-                self.loc._ensure_listlike_indexer(key)
+                self.loc._ensure_listlike_indexer(key, axis=1)
                 indexer = self.loc._get_listlike_indexer(
                     key, axis=1, raise_missing=False
                 )[1]

pandas/core/indexing.py

@@ -615,7 +615,7 @@ def _get_setitem_indexer(self, key):
                 raise
             raise IndexingError(key)
 
-    def _ensure_listlike_indexer(self, key):
+    def _ensure_listlike_indexer(self, key, axis=None):
         """
         Ensure that a list-like of column labels are all present by adding them if
         they do not already exist.
@@ -624,6 +624,7 @@ def _ensure_listlike_indexer(self, key):
         ----------
         key : _LocIndexer key or list-like of column labels
             Target labels.
+        axis : key axis if known
         """
         column_axis = 1
 
@@ -635,9 +636,11 @@ def _ensure_listlike_indexer(self, key):
             # key may be a tuple if key is a _LocIndexer key
             # in that case, set key to the column part of key
             key = key[column_axis]
+            axis = column_axis
 
         if (
-            not isinstance(self.obj._get_axis(column_axis), ABCMultiIndex)
+            axis == column_axis
+            and not isinstance(self.obj._get_axis(column_axis), ABCMultiIndex)
             and is_list_like_indexer(key)
             and not com.is_bool_indexer(key)
             and all(is_hashable(k) for k in key)