REGR: setitem with part of a MultiIndex raises (#54885)

Author: lukemanley

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

@@ -2453,12 +2453,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)