PERF: new_block_2d (#44121)

Author: jbrockmendel

pandas/core/internals/blocks.py

@@ -1289,7 +1289,8 @@ def _unstack(
         new_values = new_values.T[mask]
         new_placement = new_placement[mask]
 
-        blocks = [new_block(new_values, placement=new_placement, ndim=2)]
+        bp = BlockPlacement(new_placement)
+        blocks = [new_block_2d(new_values, placement=bp)]
         return blocks, mask
 
     @final
@@ -1318,7 +1319,7 @@ def quantile(
         assert is_list_like(qs)  # caller is responsible for this
 
         result = quantile_compat(self.values, np.asarray(qs._values), interpolation)
-        return new_block(result, placement=self._mgr_locs, ndim=2)
+        return new_block_2d(result, placement=self._mgr_locs)
 
 
 class EABackedBlock(Block):
@@ -1941,6 +1942,17 @@ def get_block_type(dtype: DtypeObj):
     return cls
 
 
+def new_block_2d(values: ArrayLike, placement: BlockPlacement):
+    # new_block specialized to case with
+    #  ndim=2
+    #  isinstance(placement, BlockPlacement)
+    #  check_ndim/ensure_block_shape already checked
+    klass = get_block_type(values.dtype)
+
+    values = maybe_coerce_values(values)
+    return klass(values, ndim=2, placement=placement)
+
+
 def new_block(values, placement, *, ndim: int) -> Block:
     # caller is responsible for ensuring values is NOT a PandasArray
 

pandas/core/internals/concat.py

@@ -48,7 +48,7 @@
 )
 from pandas.core.internals.blocks import (
     ensure_block_shape,
-    new_block,
+    new_block_2d,
 )
 from pandas.core.internals.managers import BlockManager
 
@@ -224,11 +224,11 @@ def concatenate_managers(
                 # _is_uniform_join_units ensures a single dtype, so
                 #  we can use np.concatenate, which is more performant
                 #  than concat_compat
-                values = np.concatenate(vals, axis=blk.ndim - 1)
+                values = np.concatenate(vals, axis=1)
             else:
                 # TODO(EA2D): special-casing not needed with 2D EAs
                 values = concat_compat(vals, axis=1)
-                values = ensure_block_shape(values, blk.ndim)
+                values = ensure_block_shape(values, ndim=2)
 
             values = ensure_wrapped_if_datetimelike(values)
 
@@ -240,7 +240,7 @@ def concatenate_managers(
         if fastpath:
             b = blk.make_block_same_class(values, placement=placement)
         else:
-            b = new_block(values, placement=placement, ndim=len(axes))
+            b = new_block_2d(values, placement=placement)
 
         blocks.append(b)
 

pandas/core/internals/construction.py

@@ -79,8 +79,9 @@
     SingleArrayManager,
 )
 from pandas.core.internals.blocks import (
+    BlockPlacement,
     ensure_block_shape,
-    new_block,
+    new_block_2d,
 )
 from pandas.core.internals.managers import (
     BlockManager,
@@ -370,14 +371,16 @@ def ndarray_to_mgr(
         if any(x is not y for x, y in zip(obj_columns, maybe_datetime)):
             dvals_list = [ensure_block_shape(dval, 2) for dval in maybe_datetime]
             block_values = [
-                new_block(dvals_list[n], placement=n, ndim=2)
+                new_block_2d(dvals_list[n], placement=BlockPlacement(n))
                 for n in range(len(dvals_list))
             ]
         else:
-            nb = new_block(values, placement=slice(len(columns)), ndim=2)
+            bp = BlockPlacement(slice(len(columns)))
+            nb = new_block_2d(values, placement=bp)
             block_values = [nb]
     else:
-        nb = new_block(values, placement=slice(len(columns)), ndim=2)
+        bp = BlockPlacement(slice(len(columns)))
+        nb = new_block_2d(values, placement=bp)
         block_values = [nb]
 
     if len(columns) == 0:

pandas/core/internals/managers.py

@@ -72,6 +72,7 @@
     extend_blocks,
     get_block_type,
     new_block,
+    new_block_2d,
 )
 from pandas.core.internals.ops import (
     blockwise_all,
@@ -824,6 +825,7 @@ def _slice_take_blocks_ax0(
     def _make_na_block(
         self, placement: BlockPlacement, fill_value=None, use_na_proxy: bool = False
     ) -> Block:
+        # Note: we only get here with self.ndim == 2
 
         if use_na_proxy:
             assert fill_value is None
@@ -844,7 +846,7 @@ def _make_na_block(
         # Tuple[Any, Any]]"
         block_values = np.empty(block_shape, dtype=dtype)  # type: ignore[arg-type]
         block_values.fill(fill_value)
-        return new_block(block_values, placement=placement, ndim=block_values.ndim)
+        return new_block_2d(block_values, placement=placement)
 
     def take(self: T, indexer, axis: int = 1, verify: bool = True) -> T:
         """
@@ -1106,40 +1108,38 @@ def value_getitem(placement):
             )
 
         if unfit_val_locs:
-            unfit_mgr_locs = np.concatenate(unfit_mgr_locs)
-            unfit_count = len(unfit_mgr_locs)
+            unfit_idxr = np.concatenate(unfit_mgr_locs)
+            unfit_count = len(unfit_idxr)
 
             new_blocks: list[Block] = []
             if value_is_extension_type:
                 # This code (ab-)uses the fact that EA blocks contain only
                 # one item.
                 # TODO(EA2D): special casing unnecessary with 2D EAs
                 new_blocks.extend(
-                    new_block(
+                    new_block_2d(
                         values=value,
-                        ndim=self.ndim,
-                        placement=slice(mgr_loc, mgr_loc + 1),
+                        placement=BlockPlacement(slice(mgr_loc, mgr_loc + 1)),
                     )
-                    for mgr_loc in unfit_mgr_locs
+                    for mgr_loc in unfit_idxr
                 )
 
-                self._blknos[unfit_mgr_locs] = np.arange(unfit_count) + len(self.blocks)
-                self._blklocs[unfit_mgr_locs] = 0
+                self._blknos[unfit_idxr] = np.arange(unfit_count) + len(self.blocks)
+                self._blklocs[unfit_idxr] = 0
 
             else:
                 # unfit_val_locs contains BlockPlacement objects
                 unfit_val_items = unfit_val_locs[0].append(unfit_val_locs[1:])
 
                 new_blocks.append(
-                    new_block(
+                    new_block_2d(
                         values=value_getitem(unfit_val_items),
-                        ndim=self.ndim,
-                        placement=unfit_mgr_locs,
+                        placement=BlockPlacement(unfit_idxr),
                     )
                 )
 
-                self._blknos[unfit_mgr_locs] = len(self.blocks)
-                self._blklocs[unfit_mgr_locs] = np.arange(unfit_count)
+                self._blknos[unfit_idxr] = len(self.blocks)
+                self._blklocs[unfit_idxr] = np.arange(unfit_count)
 
             self.blocks += tuple(new_blocks)
 
@@ -1161,10 +1161,15 @@ def insert(self, loc: int, item: Hashable, value: ArrayLike) -> None:
 
         if value.ndim == 2:
             value = value.T
+            if len(value) > 1:
+                raise ValueError(
+                    f"Expected a 1D array, got an array with shape {value.T.shape}"
+                )
         else:
             value = ensure_block_shape(value, ndim=self.ndim)
 
-        block = new_block(values=value, ndim=self.ndim, placement=slice(loc, loc + 1))
+        bp = BlockPlacement(slice(loc, loc + 1))
+        block = new_block_2d(values=value, placement=bp)
 
         self._insert_update_mgr_locs(loc)
         self._insert_update_blklocs_and_blknos(loc)
@@ -1968,7 +1973,7 @@ def _form_blocks(arrays: list[ArrayLike], consolidate: bool) -> list[Block]:
 def _tuples_to_blocks_no_consolidate(tuples) -> list[Block]:
     # tuples produced within _form_blocks are of the form (placement, array)
     return [
-        new_block(ensure_block_shape(x[1], ndim=2), placement=x[0], ndim=2)
+        new_block_2d(ensure_block_shape(x[1], ndim=2), placement=BlockPlacement(x[0]))
         for x in tuples
     ]
 
@@ -2035,7 +2040,7 @@ def _merge_blocks(
         new_mgr_locs = new_mgr_locs[argsort]
 
         bp = BlockPlacement(new_mgr_locs)
-        return [new_block(new_values, placement=bp, ndim=2)]
+        return [new_block_2d(new_values, placement=bp)]
 
     # can't consolidate --> no merge
     return blocks

pandas/tests/frame/indexing/test_insert.py

@@ -100,9 +100,7 @@ def test_insert_EA_no_warning(self):
     def test_insert_frame(self):
         # GH#42403
         df = DataFrame({"col1": [1, 2], "col2": [3, 4]})
-        msg = (
-            "Expected a 1D array, got an array with shape "
-            r"\(2, 2\)|Wrong number of items passed 2, placement implies 1"
-        )
+
+        msg = r"Expected a 1D array, got an array with shape \(2, 2\)"
         with pytest.raises(ValueError, match=msg):
             df.insert(1, "newcol", df)

pandas/tests/frame/test_constructors.py

@@ -222,12 +222,7 @@ def test_constructor_cast_failure(self):
         df["foo"] = np.ones((4, 2)).tolist()
 
         # this is not ok
-        msg = "|".join(
-            [
-                "Wrong number of items passed 2, placement implies 1",
-                "Expected a 1D array, got an array with shape \\(4, 2\\)",
-            ]
-        )
+        msg = "Expected a 1D array, got an array with shape \\(4, 2\\)"
         with pytest.raises(ValueError, match=msg):
             df["test"] = np.ones((4, 2))