Implement BlockManager.iset (#32350)

Author: jbrockmendel

pandas/core/internals/managers.py

@@ -8,7 +8,7 @@
 import numpy as np
 
 from pandas._libs import Timedelta, Timestamp, internals as libinternals, lib
-from pandas._typing import DtypeObj
+from pandas._typing import DtypeObj, Label
 from pandas.util._validators import validate_bool_kwarg
 
 from pandas.core.dtypes.cast import (
@@ -1001,7 +1001,25 @@ def delete(self, item):
         )
         self._rebuild_blknos_and_blklocs()
 
-    def set(self, item, value):
+    def set(self, item: Label, value):
+        """
+        Set new item in-place.
+
+        Notes
+        -----
+        Does not consolidate.
+        Adds new Block if not contained in the current items Index.
+        """
+        try:
+            loc = self.items.get_loc(item)
+        except KeyError:
+            # This item wasn't present, just insert at end
+            self.insert(len(self.items), item, value)
+            return
+
+        self.iset(loc, value)
+
+    def iset(self, loc: Union[int, slice, np.ndarray], value):
         """
         Set new item in-place. Does not consolidate. Adds new Block if not
         contained in the current set of items
@@ -1034,13 +1052,6 @@ def value_getitem(placement):
                     "Shape of new values must be compatible with manager shape"
                 )
 
-        try:
-            loc = self.items.get_loc(item)
-        except KeyError:
-            # This item wasn't present, just insert at end
-            self.insert(len(self.items), item, value)
-            return
-
         if isinstance(loc, int):
             loc = [loc]
 
@@ -1086,7 +1097,7 @@ def value_getitem(placement):
             unfit_mgr_locs = np.concatenate(unfit_mgr_locs)
             unfit_count = len(unfit_mgr_locs)
 
-            new_blocks = []
+            new_blocks: List[Block] = []
             if value_is_extension_type:
                 # This code (ab-)uses the fact that sparse blocks contain only
                 # one item.
@@ -1145,6 +1156,9 @@ def insert(self, loc: int, item, value, allow_duplicates: bool = False):
         # insert to the axis; this could possibly raise a TypeError
         new_axis = self.items.insert(loc, item)
 
+        if value.ndim == self.ndim - 1 and not is_extension_array_dtype(value):
+            value = _safe_reshape(value, (1,) + value.shape)
+
         block = make_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))
 
         for blkno, count in _fast_count_smallints(self._blknos[loc:]):