REF: setitem_with_indexer DataFrame always go through split_path

pandas/core/frame.py

@@ -5313,7 +5313,9 @@ def _replace_columnwise(
                 target, value = mapping[ax[i]]
                 newobj = ser.replace(target, value, regex=regex)
 
-                res.iloc[:, i] = newobj
+                # If we had unique columns, we could just do
+                #  res[res.columns[i]] = newobj
+                res._iset_item_mgr(i, newobj._values)
 
         if inplace:
             return

pandas/core/indexing.py

@@ -11,6 +11,7 @@
 
 from pandas._libs.indexing import NDFrameIndexerBase
 from pandas._libs.lib import item_from_zerodim
+from pandas._typing import Shape
 from pandas.errors import (
     AbstractMethodError,
     InvalidIndexError,
@@ -634,12 +635,12 @@ def __call__(self, axis=None):
         new_self.axis = axis
         return new_self
 
-    def _get_setitem_indexer(self, key):
+    def _get_setitem_indexer(self, key, value):
         """
         Convert a potentially-label-based key into a positional indexer.
         """
         if self.name == "loc":
-            self._ensure_listlike_indexer(key)
+            self._ensure_listlike_indexer(key, value=value)
 
         if self.axis is not None:
             return self._convert_tuple(key)
@@ -677,9 +678,11 @@ def _ensure_listlike_indexer(self, key, axis=None, value=None):
         if self.ndim != 2:
             return
 
+        pi = None
         if isinstance(key, tuple) and len(key) > 1:
             # key may be a tuple if we are .loc
             # if length of key is > 1 set key to column part
+            pi = key[0]
             key = key[column_axis]
             axis = column_axis
 
@@ -693,17 +696,26 @@ def _ensure_listlike_indexer(self, key, axis=None, value=None):
             # GH#38148
             keys = self.obj.columns.union(key, sort=False)
 
-            self.obj._mgr = self.obj._mgr.reindex_axis(
-                keys, axis=0, consolidate=False, only_slice=True
-            )
+            if isinstance(value, ABCDataFrame) and com.is_null_slice(pi):
+                # We are setting obj.loc[:, new_keys] = newframe
+                # Setting these directly instead of reindexing keeps
+                #  us from converting integer dtypes to floats
+                new_keys = keys.difference(self.obj.columns)
+                self.obj[new_keys] = value[new_keys]
+
+            else:
+
+                self.obj._mgr = self.obj._mgr.reindex_axis(
+                    keys, axis=0, consolidate=False, only_slice=True
+                )
 
     def __setitem__(self, key, value):
         if isinstance(key, tuple):
             key = tuple(list(x) if is_iterator(x) else x for x in key)
             key = tuple(com.apply_if_callable(x, self.obj) for x in key)
         else:
             key = com.apply_if_callable(key, self.obj)
-        indexer = self._get_setitem_indexer(key)
+        indexer = self._get_setitem_indexer(key, value)
         self._has_valid_setitem_indexer(key)
 
         iloc = self if self.name == "iloc" else self.obj.iloc
@@ -1271,9 +1283,10 @@ def _convert_to_indexer(self, key, axis: int):
                 key = list(key)
 
             if com.is_bool_indexer(key):
+                # TODO: in this case should we do a .take on the value here?
+                # test_loc_setitem_all_false_boolean_two_blocks
                 key = check_bool_indexer(labels, key)
-                (inds,) = key.nonzero()
-                return inds
+                return key
             else:
                 return self._get_listlike_indexer(key, axis)[1]
         else:
@@ -1525,7 +1538,7 @@ def _convert_to_indexer(self, key, axis: int):
         """
         return key
 
-    def _get_setitem_indexer(self, key):
+    def _get_setitem_indexer(self, key, value):
         # GH#32257 Fall through to let numpy do validation
         if is_iterator(key):
             return list(key)
@@ -1547,32 +1560,6 @@ def _setitem_with_indexer(self, indexer, value, name="iloc"):
         """
         info_axis = self.obj._info_axis_number
 
-        # maybe partial set
-        take_split_path = not self.obj._mgr.is_single_block
-
-        # if there is only one block/type, still have to take split path
-        # unless the block is one-dimensional or it can hold the value
-        if (
-            not take_split_path
-            and getattr(self.obj._mgr, "blocks", False)
-            and self.ndim > 1
-        ):
-            # in case of dict, keys are indices
-            val = list(value.values()) if isinstance(value, dict) else value
-            blk = self.obj._mgr.blocks[0]
-            take_split_path = not blk._can_hold_element(val)
-
-        # if we have any multi-indexes that have non-trivial slices
-        # (not null slices) then we must take the split path, xref
-        # GH 10360, GH 27841
-        if isinstance(indexer, tuple) and len(indexer) == len(self.obj.axes):
-            for i, ax in zip(indexer, self.obj.axes):
-                if isinstance(ax, MultiIndex) and not (
-                    is_integer(i) or com.is_null_slice(i)
-                ):
-                    take_split_path = True
-                    break
-
         if isinstance(indexer, tuple):
             nindexer = []
             for i, idx in enumerate(indexer):
@@ -1666,7 +1653,7 @@ def _setitem_with_indexer(self, indexer, value, name="iloc"):
                 return
 
         # align and set the values
-        if take_split_path:
+        if self.ndim > 1:
             # We have to operate column-wise
             self._setitem_with_indexer_split_path(indexer, value, name)
         else:
@@ -1679,23 +1666,65 @@ def _setitem_with_indexer_split_path(self, indexer, value, name: str):
         # Above we only set take_split_path to True for 2D cases
         assert self.ndim == 2
 
+        orig = indexer
         if not isinstance(indexer, tuple):
             indexer = _tuplify(self.ndim, indexer)
         if len(indexer) > self.ndim:
             raise IndexError("too many indices for array")
         if isinstance(indexer[0], np.ndarray) and indexer[0].ndim > 2:
             raise ValueError(r"Cannot set values with ndim > 2")
-
         if (isinstance(value, ABCSeries) and name != "iloc") or isinstance(value, dict):
             from pandas import Series
 
             value = self._align_series(indexer, Series(value))
 
+        info_idx = indexer[1]
+        pi = indexer[0]
+        if (
+            isinstance(pi, ABCDataFrame)
+            and orig is pi
+            and hasattr(self.obj._mgr, "blocks")
+            and len(self.obj._mgr.blocks) == 1
+        ):
+            # FIXME: kludge
+            return self._setitem_single_block(orig, value, name)
+
+        from pandas.core.internals import ArrayManager
+
+        if (
+            com.is_null_slice(info_idx)
+            and is_scalar(value)
+            and not isinstance(pi, ABCDataFrame)
+            and not isinstance(self.obj._mgr, ArrayManager)
+        ):
+            # We can go directly through BlockManager.setitem without worrying
+            #  about alignment.
+            # TODO: do we need to do some kind of copy_with_setting check?
+            self.obj._mgr = self.obj._mgr.setitem(indexer=indexer, value=value)
+            return
+
+        if is_integer(info_idx) and not isinstance(self.obj._mgr, ArrayManager):
+            if is_integer(pi):
+                # We need to watch out for case where we are treating a listlike
+                #  as a scalar, e.g. test_setitem_iloc_scalar_single for JSONArray
+
+                mgr = self.obj._mgr
+                blkno = mgr.blknos[info_idx]
+                blkloc = mgr.blklocs[info_idx]
+                blk = mgr.blocks[blkno]
+
+                if blk._can_hold_element(value):
+                    # NB: we are assuming here that _can_hold_element is accurate
+                    # TODO: do we need to do some kind of copy_with_setting check?
+                    self.obj._check_is_chained_assignment_possible()
+                    blk.setitem_inplace((pi, blkloc), value)
+                    self.obj._maybe_update_cacher(clear=True)
+                    return
+
         # Ensure we have something we can iterate over
         info_axis = indexer[1]
         ilocs = self._ensure_iterable_column_indexer(info_axis)
 
-        pi = indexer[0]
         lplane_indexer = length_of_indexer(pi, self.obj.index)
         # lplane_indexer gives the expected length of obj[indexer[0]]
 
@@ -1711,7 +1740,7 @@ def _setitem_with_indexer_split_path(self, indexer, value, name: str):
 
             elif len(ilocs) == 1 and lplane_indexer == len(value) and not is_scalar(pi):
                 # We are setting multiple rows in a single column.
-                self._setitem_single_column(ilocs[0], value, pi)
+                self._setitem_iat_loc(ilocs[0], pi, value)
 
             elif len(ilocs) == 1 and 0 != lplane_indexer != len(value):
                 # We are trying to set N values into M entries of a single
@@ -1739,7 +1768,7 @@ def _setitem_with_indexer_split_path(self, indexer, value, name: str):
             elif len(ilocs) == len(value):
                 # We are setting multiple columns in a single row.
                 for loc, v in zip(ilocs, value):
-                    self._setitem_single_column(loc, v, pi)
+                    self._setitem_iat_loc(loc, pi, v)
 
             elif len(ilocs) == 1 and com.is_null_slice(pi) and len(self.obj) == 0:
                 # This is a setitem-with-expansion, see
@@ -1777,6 +1806,7 @@ def _setitem_with_indexer_2d_value(self, indexer, value):
 
         for i, loc in enumerate(ilocs):
             # setting with a list, re-coerces
+            # self._setitem_iat_loc(loc, pi, value[:, i].tolist())
             self._setitem_single_column(loc, value[:, i].tolist(), pi)
 
     def _setitem_with_indexer_frame_value(self, indexer, value: DataFrame, name: str):
@@ -1793,7 +1823,7 @@ def _setitem_with_indexer_frame_value(self, indexer, value: DataFrame, name: str
         if name == "iloc":
             for i, loc in enumerate(ilocs):
                 val = value.iloc[:, i]
-                self._setitem_single_column(loc, val, pi)
+                self._setitem_iat_loc(loc, pi, val)
 
         elif not unique_cols and value.columns.equals(self.obj.columns):
             # We assume we are already aligned, see
@@ -1810,12 +1840,21 @@ def _setitem_with_indexer_frame_value(self, indexer, value: DataFrame, name: str
                 else:
                     val = np.nan
 
-                self._setitem_single_column(loc, val, pi)
+                self._setitem_iat_loc(loc, pi, val)
 
         elif not unique_cols:
             raise ValueError("Setting with non-unique columns is not allowed.")
 
         else:
+            # TODO: not totally clear why we are requiring this
+            # Need so that we raise in test_multiindex_setitem
+            self._align_frame(indexer[0], value)
+
+            if com.is_bool_indexer(indexer[0]) and indexer[0].sum() == len(value):
+                # TODO: better place for this?
+                pi = indexer[0].nonzero()[0]
+                sub_indexer[0] = pi
+
             for loc in ilocs:
                 item = self.obj.columns[loc]
                 if item in value:
@@ -1826,7 +1865,7 @@ def _setitem_with_indexer_frame_value(self, indexer, value: DataFrame, name: str
                 else:
                     val = np.nan
 
-                self._setitem_single_column(loc, val, pi)
+                self._setitem_iat_loc(loc, pi, val)
 
     def _setitem_single_column(self, loc: int, value, plane_indexer):
         """
@@ -1839,6 +1878,7 @@ def _setitem_single_column(self, loc: int, value, plane_indexer):
             The indexer we use for setitem along axis=0.
         """
         pi = plane_indexer
+        pi, value = mask_setitem_value(pi, value, (len(self.obj),))
 
         ser = self.obj._ixs(loc, axis=1)
 
@@ -1877,6 +1917,35 @@ def _setitem_single_column(self, loc: int, value, plane_indexer):
         # reset the sliced object if unique
         self.obj._iset_item(loc, ser)
 
+    def _setitem_iat_loc(self, loc: int, pi, value):
+        # TODO: likely a BM method?
+        from pandas.core.internals import ArrayManager
+
+        if isinstance(self.obj._mgr, ArrayManager):
+            # TODO: implement this correctly for ArrayManager
+            return self._setitem_single_column(loc, value, pi)
+
+        mgr = self.obj._mgr
+        blkno = mgr.blknos[loc]
+        blkloc = mgr.blklocs[loc]
+        blk = mgr.blocks[blkno]
+        assert blk.mgr_locs[blkloc] == loc
+
+        if blk._can_hold_element(value):
+            # NB: we are assuming here that _can_hold_element is accurate
+            # TODO: do we need to do some kind of copy_with_setting check?
+            try:
+                self.obj._check_is_chained_assignment_possible()
+                blk.setitem_inplace((pi, blkloc), value)
+                self.obj._maybe_update_cacher(clear=True)
+            except ValueError:
+                if blk.is_extension:
+                    # FIXME: kludge bc _can_hold_element is wrong for EABLock
+                    return self._setitem_single_column(loc, value, pi)
+                raise
+        else:
+            self._setitem_single_column(loc, value, pi)
+
     def _setitem_single_block(self, indexer, value, name: str):
         """
         _setitem_with_indexer for the case when we have a single Block.
@@ -2442,3 +2511,34 @@ def need_slice(obj: slice) -> bool:
         or obj.stop is not None
         or (obj.step is not None and obj.step != 1)
     )
+
+
+def mask_setitem_value(indexer, value, shape: Shape):
+    """
+    Convert a boolean indexer to a positional indexer, masking `value` if necessary.
+    """
+    if com.is_bool_indexer(indexer):
+        indexer = np.asarray(indexer).nonzero()[0]
+        if is_list_like(value) and len(value) == shape[0]:
+            if not is_array_like(value):
+                value = [value[n] for n in indexer]
+            else:
+                value = value[indexer]
+
+    elif isinstance(indexer, tuple):
+        indexer = list(indexer)
+        for i, key in enumerate(indexer):
+            if com.is_bool_indexer(key):
+                new_key = np.asarray(key).nonzero()[0]
+                indexer[i] = new_key
+
+                if is_list_like(value) and len(value) == shape[i]:
+                    # FIXME: assuming value.ndim == 1 here?
+                    # FIXME: assuming non-i tuple member is scalar?
+                    if not is_array_like(value):
+                        value = [value[n] for n in new_key]
+                    else:
+                        value = value[new_key]
+
+        indexer = tuple(indexer)
+    return indexer, value