allow for top and mid-level assignment to DataFrames with MultIndex columns

Author: iamlemec

pandas/core/frame.py

@@ -3057,10 +3057,15 @@ def __setitem__(self, key, value):
             #  to a slice for partial-string date indexing
             return self._setitem_slice(indexer, value)
 
+        # mimic getitem behavior
+        is_single_key = isinstance(key, tuple) or not is_list_like(key)
+
         if isinstance(key, DataFrame) or getattr(key, "ndim", None) == 2:
             self._setitem_frame(key, value)
         elif isinstance(key, (Series, np.ndarray, list, Index)):
             self._setitem_array(key, value)
+        elif is_single_key and self.columns.nlevels > 1:
+            return self._setitem_multilevel(key, value)
         else:
             # set column
             self._set_item(key, value)
@@ -3085,10 +3090,18 @@ def _setitem_array(self, key, value):
             self.iloc._setitem_with_indexer(indexer, value)
         else:
             if isinstance(value, DataFrame):
-                if len(value.columns) != len(key):
-                    raise ValueError("Columns must be same length as key")
-                for k1, k2 in zip(key, value.columns):
-                    self[k1] = value[k2]
+                columns = value.columns
+                if len(columns) == len(key):
+                    for k1, k2 in zip(key, columns):
+                        self[k1] = value[k2]
+                elif columns.nlevels > 1 and len(columns.levels[0]) == len(key):
+                    for k1, k2 in zip(key, columns.levels[0]):
+                        self[k1] = value[k2]
+                else:
+                    raise ValueError(
+                        "Key must be same length as columns or top level of "
+                        "MultiIndex"
+                    )
             else:
                 self.loc._ensure_listlike_indexer(key, axis=1)
                 indexer = self.loc._get_listlike_indexer(
@@ -3114,6 +3127,30 @@ def _setitem_frame(self, key, value):
         self._check_setitem_copy()
         self._where(-key, value, inplace=True)
 
+    def _setitem_multilevel(self, key, value):
+        # self.columns is a MultiIndex
+        if key in self.columns:
+            self._set_item(key, value)
+        else:
+            if not isinstance(key, tuple):
+                key = (key,)
+            if isinstance(value, DataFrame):
+                if len(key) + value.columns.nlevels != self.columns.nlevels:
+                    raise TypeError(
+                        "Must pass key/value pair that conforms with number of column "
+                        "levels"
+                    )
+                if value.columns.nlevels > 1:
+                    items = MultiIndex.from_tuples([key + i for i in value.columns])
+                else:
+                    items = MultiIndex.from_tuples([key + (i,) for i in value.columns])
+            else:
+                if len(key) < self.columns.nlevels:
+                    key = key + ("",) * (self.columns.nlevels - len(key))
+                items = MultiIndex.from_tuples([key])
+            value = self._sanitize_column(key, value)
+            self._mgr.append_block(items, value)
+
     def _iset_item(self, loc: int, value):
         self._ensure_valid_index(value)
 

pandas/core/internals/managers.py

@@ -1204,6 +1204,26 @@ def insert(self, loc: int, item: Label, value, allow_duplicates: bool = False):
         if len(self.blocks) > 100:
             self._consolidate_inplace()
 
+    def append_block(self, items, values):
+        base, size = len(self.items), len(items)
+
+        new_axis = self.items.append(items)
+        block = make_block(
+            values=values, ndim=self.ndim, placement=slice(base, base + size)
+        )
+
+        blk_no = len(self.blocks)
+        self._blklocs = np.append(self.blklocs, range(size))
+        self._blknos = np.append(self.blknos, size * (blk_no,))
+
+        self.axes[0] = new_axis
+        self.blocks += (block,)
+
+        self._known_consolidated = False
+
+        if len(self.blocks) > 100:
+            self._consolidate_inplace()
+
     def reindex_axis(
         self,
         new_index,

pandas/tests/frame/indexing/test_indexing.py

@@ -157,7 +157,7 @@ def test_setitem_list(self, float_frame):
         tm.assert_series_equal(float_frame["B"], data["A"], check_names=False)
         tm.assert_series_equal(float_frame["A"], data["B"], check_names=False)
 
-        msg = "Columns must be same length as key"
+        msg = "Key must be same length as columns or top level of MultiIndex"
         with pytest.raises(ValueError, match=msg):
             data[["A"]] = float_frame[["A", "B"]]
         newcolumndata = range(len(data.index) - 1)

pandas/tests/indexing/multiindex/test_multiindex.py

@@ -91,3 +91,71 @@ def test_multiindex_get_loc_list_raises(self):
         msg = "unhashable type"
         with pytest.raises(TypeError, match=msg):
             idx.get_loc([])
+
+    def test_multiindex_frame_assign(self):
+        df0 = pd.DataFrame({"a": [0, 1, 2, 3], "b": [3, 4, 5, 6]})
+        df1 = pd.concat({"x": df0, "y": df0}, axis=1)
+        df2 = pd.concat({"q": df1, "r": df1}, axis=1)
+
+        # level one assign
+        result = df2.copy()
+        result["m"] = result["q"] + result["r"]
+        expected = pd.concat({"q": df1, "r": df1, "m": 2 * df1}, axis=1)
+        tm.assert_frame_equal(result, expected)
+
+        # level one assign - multiple
+        result = df2.copy()
+        result[["m", "n"]] = 2 * result[["q", "r"]]
+        expected = pd.concat({"q": df1, "r": df1, "m": 2 * df1, "n": 2 * df1}, axis=1)
+        tm.assert_frame_equal(result, expected)
+
+        # level two assign
+        result = df2.copy()
+        result["m", "x"] = df2["q", "x"] + df2["q", "y"]
+        expected = pd.concat(
+            {"q": df1, "r": df1, "m": pd.concat({"x": 2 * df0}, axis=1)}, axis=1
+        )
+        tm.assert_frame_equal(result, expected)
+
+        # level two assign - multiple (seems like getitem is not caught up here)
+        result = df2.copy()
+        result[[("m", "x"), ("n", "y")]] = 2 * df2["q"]
+        expected = pd.concat(
+            {
+                "q": df1,
+                "r": df1,
+                "m": pd.concat({"x": 2 * df0}, axis=1),
+                "n": pd.concat({"y": 2 * df0}, axis=1),
+            },
+            axis=1,
+        )
+        tm.assert_frame_equal(result, expected)
+
+        # level three assign
+        result = df2.copy()
+        result["m", "x", "a"] = df2["q", "x", "a"] + df2["q", "x", "b"]
+        expected = pd.concat(
+            {
+                "q": df1,
+                "r": df1,
+                "m": pd.concat(
+                    {"x": pd.concat({"a": df0["a"] + df0["b"]}, axis=1)}, axis=1
+                ),
+            },
+            axis=1,
+        )
+        tm.assert_frame_equal(result, expected)
+
+        # level three assign - multiple
+        result = df2.copy()
+        result[[("m", "x", "a"), ("n", "y", "b")]] = 2 * df2["q", "x"]
+        expected = pd.concat(
+            {
+                "q": df1,
+                "r": df1,
+                "m": pd.concat({"x": pd.concat({"a": 2 * df0["a"]}, axis=1)}, axis=1),
+                "n": pd.concat({"y": pd.concat({"b": 2 * df0["b"]}, axis=1)}, axis=1),
+            },
+            axis=1,
+        )
+        tm.assert_frame_equal(result, expected)