Backport PR #54885 on branch 2.1.x (REGR: setitem with part of a MultiIndex raises) (#54899)

Backport PR #54885: REGR: setitem with part of a MultiIndex raises

Co-authored-by: Luke Manley <[email protected]>

Author: meeseeksmachine

doc/source/whatsnew/v2.1.1.rst

@@ -13,7 +13,7 @@ including other versions of pandas.
 
 Fixed regressions
 ~~~~~~~~~~~~~~~~~
--
+- Fixed regression in :meth:`DataFrame.__setitem__` raising ``AssertionError`` when setting a :class:`Series` with a partial :class:`MultiIndex` (:issue:`54875`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_211.bug_fixes:

pandas/core/indexes/multi.py

@@ -2460,12 +2460,12 @@ def _reorder_ilevels(self, order) -> MultiIndex:
     def _recode_for_new_levels(
         self, new_levels, copy: bool = True
     ) -> Generator[np.ndarray, None, None]:
-        if len(new_levels) != self.nlevels:
+        if len(new_levels) > self.nlevels:
             raise AssertionError(
                 f"Length of new_levels ({len(new_levels)}) "
-                f"must be same as self.nlevels ({self.nlevels})"
+                f"must be <= self.nlevels ({self.nlevels})"
             )
-        for i in range(self.nlevels):
+        for i in range(len(new_levels)):
             yield recode_for_categories(
                 self.codes[i], self.levels[i], new_levels[i], copy=copy
             )

pandas/tests/indexing/multiindex/test_setitem.py

@@ -556,3 +556,21 @@ def test_frame_setitem_copy_no_write(
 
     result = df
     tm.assert_frame_equal(result, expected)
+
+
+def test_frame_setitem_partial_multiindex():
+    # GH 54875
+    df = DataFrame(
+        {
+            "a": [1, 2, 3],
+            "b": [3, 4, 5],
+            "c": 6,
+            "d": 7,
+        }
+    ).set_index(["a", "b", "c"])
+    ser = Series(8, index=df.index.droplevel("c"))
+    result = df.copy()
+    result["d"] = ser
+    expected = df.copy()
+    expected["d"] = 8
+    tm.assert_frame_equal(result, expected)