BUG: loc setitem with a partial MultiIndex key sets NaN (GH#43719)

@@ -219,6 +219,7 @@ MultiIndex - Bug in :meth:`MultiIndex.append`, :func:`concat`, :meth:`DataFrame.join` and :func:`merge` emitting a spurious ``RuntimeWarning`` when :class:`MultiIndex` level values could not be sorted (:issue:`61477`) - Bug in :meth:`MultiIndex.from_tuples` with an object ``ndarray``, :meth:`MultiIndex.append` and :meth:`Index.map` truncating or raising ``IndexError`` on tuples of different lengths; :meth:`MultiIndex.from_tuples` now raises a ``ValueError`` for an ``ndarray`` of tuples of different lengths (:issue:`50293`) - Bug in :meth:`Series.map`, :meth:`Index.get_indexer`, :meth:`DataFrame.reindex` and :attr:`DataFrame.loc` with a :class:`MultiIndex` raising ``IndexError`` or matching on a prefix when looking up tuples whose length differs from the number of levels (:issue:`50293`)+- Bug in :meth:`DataFrame.loc` and :meth:`Series.loc` setitem with a partial :class:`MultiIndex` key (e.g. ``df.loc[:, "a"] = df.loc[:, "a"] * 2``) setting ``NaN`` (or raising) instead of the value (:issue:`10440`, :issue:`15310`, :issue:`40186`, :issue:`43719`, :issue:`45751`, :issue:`46837`) -  I/O

@@ -932,6 +932,71 @@ def _maybe_mask_setitem_value(self, indexer, value):          return indexer, value +    @final+    def _maybe_restore_dropped_levels(self, key, indexer, value):+        """+        Relabel a Series/DataFrame value whose axis lacks the MultiIndex levels+        matched by the leading scalars of that axis's key.++        ``obj.loc[partial_key]`` drops those levels, so without this+        ``obj.loc[partial_key] = obj.loc[partial_key] * 2`` aligns to all-NaN.+        """+        if not isinstance(value, (ABCSeries, ABCDataFrame)):+            return value++        if self.axis is not None:+            key = _tupleize_axis_indexer(self.ndim, self.axis, key)+        # mirror how _get_setitem_indexer split the key between the axes+        if isinstance(indexer, tuple):+            keys = key+        else:+            keys, indexer = (key,), (indexer,)++        if value.ndim == 2:+            axes = list(range(len(indexer)))+        elif self.ndim == 1 or not (+            # a dict is a missing scalar key, see convert_missing_indexer+            is_scalar(indexer[0]) or isinstance(indexer[0], dict)+        ):+            # a Series value aligns on the rows unless they are selected by a scalar+            axes = [0]+        elif len(indexer) > 1:+            axes = [1]+        else:+            return value++        for axis in axes:+            ax = self.obj._get_axis(axis)+            idx = indexer[axis]+            value_axis = axis if value.ndim == 2 else 0+            value_ax = value._get_axis(value_axis)+            ndropped = ax.nlevels - value_ax.nlevels+            if (+                not isinstance(ax, MultiIndex)+                or not 0 < ndropped <= _count_leading_scalars(keys[axis])+                or not (+                    isinstance(idx, slice)+                    or (is_list_like_indexer(idx) and getattr(idx, "ndim", 1) == 1)+                )+                or isinstance(idx, dict)+            ):+                continue+            target = ax[idx]+            if not len(target) or not target.is_unique or not value_ax.is_unique:+                continue+            dropped = target._drop_level_numbers(list(range(ndropped))[::-1])+            # dropped repeats labels if the key matched several, see+            #  test_loc_setitem_partial_key_matching_several_labels+            positions = value_ax.get_indexer(dropped)+            found = positions != -1+            if not found.any():+                # e.g. the value is already labeled like the target+                continue+            # labels unmatched on either side are left to the usual alignment+            value = value.take(positions[found], axis=value_axis)+            value = value.set_axis(target[found], axis=value_axis)+        return value+     @final     def _ensure_listlike_indexer(self, key, axis=None) -> None:         """@@ -1004,6 +1069,8 @@ def __setitem__(self, key, value) -> None:         orig_nrows = self.obj.shape[0]         indexer = self._get_setitem_indexer(key)  # may alter self.obj         self._has_valid_setitem_indexer(key)+        if self.name == "loc":+            value = self._maybe_restore_dropped_levels(key, indexer, value)          # Capture dtype info cheaply between potential column expansion         # (above) and potential row expansion (below).@@ -3662,6 +3729,20 @@ def _tuplify(ndim: int, loc: Hashable) -> tuple[Hashable | slice, ...]:     return tuple(_tup)  +def _count_leading_scalars(key) -> int:+    """+    Number of leading MultiIndex levels pinned to a single label by ``key``.+    """+    if not isinstance(key, tuple):+        return int(is_scalar(key))+    count = 0+    for item in key:+        if not is_scalar(item):+            break+        count += 1+    return count++ def _tupleize_axis_indexer(ndim: int, axis: AxisInt, key) -> tuple[Any, ...]:     """     If we have an axis, adapt the given key to be axis-independent.

@@ -139,10 +139,9 @@ def test_multiindex_setitem(self):         df.loc[["bar"]] *= 2         tm.assert_frame_equal(df.loc[["bar"]], expected) -        # raise because these have differing levels-        msg = "cannot align on a multi-index with out specifying the join levels"-        with pytest.raises(TypeError, match=msg):-            df.loc["bar"] *= 2+        # GH#43719 the levels dropped by the partial key are restored to align+        df.loc["bar"] *= 2+        tm.assert_frame_equal(df.loc[["bar"]], expected * 2)      def test_multiindex_setitem2(self):         # from SO@@ -728,3 +727,163 @@ def test_iloc_setitem_unique_key_on_duplicated_multiindex_columns(     df.iloc[:, 0] = value     expected = pd.DataFrame(expected_values, columns=columns)     tm.assert_frame_equal(df, expected)+++@pytest.mark.parametrize("rows", [slice(None), [0, 1], np.array([True, False])])+@pytest.mark.parametrize("key", ["a", ("a",)])+@pytest.mark.parametrize("dtype", ["float32", "int64", "Int64"])+def test_loc_setitem_partial_column_key_frame_value(rows, key, dtype):+    # GH#43719 the value's columns lack the level matched by the partial key+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=mi, dtype=dtype)+    expected = df.copy()+    expected.loc[rows, ("a", slice(None))] = df.loc[:, ("a", slice(None))] * 10++    df.loc[rows, key] = df.loc[:, key] * 10+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_partial_column_key_frame_value_aligns():+    # GH#43719 the value is aligned by label, not position+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame([[1.0, 2.0, 3.0]], columns=mi)+    df.loc[:, "a"] = df.loc[:, "a"][["y", "x"]] * 10+    expected = pd.DataFrame([[10.0, 20.0, 3.0]], columns=mi)+    tm.assert_frame_equal(df, expected)+++@pytest.mark.parametrize("key", ["a", ("a", "b")])+def test_loc_setitem_partial_column_key_three_levels(key):+    # GH#43719+    mi = pd.MultiIndex.from_tuples(+        [("a", "b", "x"), ("a", "b", "y"), ("a", "c", "x"), ("d", "b", "x")]+    )+    df = pd.DataFrame([[1.0, 2.0, 3.0, 4.0]], columns=mi)+    expected = df.copy()+    expected[key] = df[key] * 10++    df.loc[:, key] = df.loc[:, key] * 10+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_partial_row_key_value():+    # GH#43719 same for the rows of a DataFrame or Series+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame({"A": [1.0, 2.0, 3.0], "B": [4.0, 5.0, 6.0]}, index=mi)+    expected = pd.DataFrame({"A": [10.0, 20.0, 3.0], "B": [4.0, 5.0, 6.0]}, index=mi)++    result = df.copy()+    result.loc["a"] = result.loc["a"] * [10, 1]+    tm.assert_frame_equal(result, expected)++    result = df.copy()+    result.loc["a", "A"] = result.loc["a", "A"] * 10+    tm.assert_frame_equal(result, expected)++    for key in ["a", ("a", slice(None))]:+        ser = df["A"].copy()+        ser.loc[key] = ser.loc[key] * 10+        tm.assert_series_equal(ser, expected["A"])+++def test_loc_setitem_partial_row_key_value_from_other_key():+    # GH#10440 the value is aligned on the levels not matched by the key+    mi = pd.MultiIndex.from_product([["p", "q"], ["x", "y"]])+    df = pd.DataFrame({"A": [1.0, 2.0, 3.0, 4.0]}, index=mi)+    df.loc["q", "A"] = df.loc["p", "A"][["y", "x"]]+    expected = pd.DataFrame({"A": [1.0, 2.0, 1.0, 2.0]}, index=mi)+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_partial_column_key_series_value():+    # GH#43719 a Series value aligned along the columns+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], columns=mi)+    df.loc[0, "a"] = df.loc[0, "a"] * 10+    expected = pd.DataFrame([[10.0, 20.0, 3.0], [4.0, 5.0, 6.0]], columns=mi)+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_scalar_and_list_column_key_series_value():+    # GH#15310+    mi = pd.MultiIndex.from_product([["a", "b"], ["x", "y"]])+    df = pd.DataFrame(np.zeros((2, 4)), columns=mi)+    df.loc[0, ("a", ["y", "x"])] = pd.Series({"x": 1.0, "y": 2.0})+    expected = pd.DataFrame([[1.0, 2.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]], columns=mi)+    tm.assert_frame_equal(df, expected)+++@pytest.mark.parametrize("dtype", ["string", "Int64", "datetime64[ns, UTC]"])+def test_loc_setitem_partial_column_key_frame_value_missing_labels(dtype):+    # GH#43719 labels missing from the value are filled with the dtype's NA+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=mi).astype(dtype)+    expected = df.copy()+    expected.loc[:, ("a", slice(None))] = df.loc[:, [("a", "x")]]++    df.loc[:, "a"] = df.loc[:, "a"][["x"]]+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_partial_row_key_value_not_relabeled():+    # GH#43719 a value whose labels are not the remaining levels is left alone+    mi = pd.MultiIndex.from_tuples(+        [("a", "b", "x"), ("a", "b", "y"), ("a", "c", "x"), ("d", "b", "x")]+    )+    ser = pd.Series([1.0, 2.0, 3.0, 4.0], index=mi)+    value = pd.Series([10.0, 20.0], index=mi[[0, 2]].droplevel(2))+    ser.loc["a"] = value+    expected = pd.Series([10.0, 10.0, 20.0, 4.0], index=mi)+    tm.assert_series_equal(ser, expected)++    ser = pd.Series([1.0, 2.0, 3.0], index=mi[[0, 1, 3]].droplevel(1))+    value = ser.loc[["a"]] * 10+    value.index = value.index.to_flat_index()+    ser.loc["a"] = value+    expected = pd.Series([10.0, 20.0, 3.0], index=ser.index)+    tm.assert_series_equal(ser, expected)+++@pytest.mark.parametrize("rows", [slice(None), [True, True]])+def test_loc_setitem_non_partial_key_value_not_relabeled(rows):+    # GH#43719 only the levels matched by a partial key are restored, so the+    #  value is not aligned on "x"/"y" even though every row is under "g"+    mi = pd.MultiIndex.from_product([["g"], ["x", "y"]])+    ser = pd.Series([1.0, 2.0], index=mi)+    ser.loc[rows] = pd.Series([10.0, 20.0], index=["x", "y"])+    expected = pd.Series([np.nan, np.nan], index=mi)+    tm.assert_series_equal(ser, expected)+++def test_loc_setitem_partial_column_key_with_enlargement():+    # GH#43719+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    df = pd.DataFrame([[1.0, 2.0, 3.0]], columns=mi)+    df.loc[1, "a"] = df.loc[0, "a"] * 10+    expected = pd.DataFrame([[1.0, 2.0, 3.0], [10.0, 20.0, np.nan]], columns=mi)+    tm.assert_frame_equal(df, expected)+++def test_loc_setitem_partial_row_key_value_extra_labels():+    # GH#43719 value labels not under the key are ignored+    mi = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "x")])+    ser = pd.Series([1.0, 2.0, 3.0], index=mi)+    ser.loc["a"] = pd.Series({"x": 10.0, "z": 30.0})+    expected = pd.Series([10.0, np.nan, 3.0], index=mi)+    tm.assert_series_equal(ser, expected)+++@pytest.mark.parametrize(+    "first_level, key",+    [+        (pd.to_datetime(["2020-01-01", "2020-01-02", "2021-01-01"]), "2020"),+        (pd.IntervalIndex.from_tuples([(0, 2), (1, 3), (4, 5)]), 1.5),+    ],+)+def test_loc_setitem_partial_key_matching_several_labels(first_level, key):+    # GH#43719 the value is broadcast to each first-level label the key matches+    mi = pd.MultiIndex.from_product([first_level, ["x", "y"]])+    ser = pd.Series(np.arange(6.0), index=mi)+    ser.loc[key] = pd.Series({"x": 10.0, "y": 20.0})+    expected = pd.Series([10.0, 20.0, 10.0, 20.0, 4.0, 5.0], index=mi)+    tm.assert_series_equal(ser, expected)