PERF: _form_blocks (#43144)

Author: jbrockmendel

pandas/core/internals/managers.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-from collections import defaultdict
 import itertools
 from typing import (
     Any,
     Callable,
-    DefaultDict,
     Hashable,
     Sequence,
     TypeVar,
@@ -67,9 +65,7 @@
 )
 from pandas.core.internals.blocks import (
     Block,
-    CategoricalBlock,
     DatetimeTZBlock,
-    ExtensionBlock,
     ensure_block_shape,
     extend_blocks,
     get_block_type,
@@ -1863,63 +1859,56 @@ def construction_error(
 # -----------------------------------------------------------------------
 
 
-def _form_blocks(arrays: list[ArrayLike], consolidate: bool) -> list[Block]:
-
-    items_dict: DefaultDict[str, list] = defaultdict(list)
-
-    for i, name_idx in enumerate(range(len(arrays))):
+def _grouping_func(tup):
+    # compat for numpy<1.21, in which comparing a np.dtype with an ExtensionDtype
+    # raises instead of returning False. Once earlier numpy versions are dropped,
+    # this can be simplified to `return tup[1].dtype`
+    dtype = tup[1].dtype
+    return isinstance(dtype, np.dtype), dtype
 
-        v = arrays[name_idx]
 
-        block_type = get_block_type(v)
-        items_dict[block_type.__name__].append((i, v))
+def _form_blocks(arrays: list[ArrayLike], consolidate: bool) -> list[Block]:
+    tuples = list(enumerate(arrays))
 
-    blocks: list[Block] = []
-    if len(items_dict["NumericBlock"]):
-        numeric_blocks = multi_blockify(
-            items_dict["NumericBlock"], consolidate=consolidate
-        )
-        blocks.extend(numeric_blocks)
+    if not consolidate:
+        nbs = _tuples_to_blocks_no_consolidate(tuples, dtype=None)
+        return nbs
 
-    if len(items_dict["DatetimeLikeBlock"]):
-        dtlike_blocks = multi_blockify(
-            items_dict["DatetimeLikeBlock"], consolidate=consolidate
-        )
-        blocks.extend(dtlike_blocks)
+    # group by dtype
+    grouper = itertools.groupby(tuples, _grouping_func)
 
-    if len(items_dict["DatetimeTZBlock"]):
-        dttz_blocks = [
-            DatetimeTZBlock(
-                ensure_block_shape(extract_array(array), 2),
-                placement=BlockPlacement(i),
-                ndim=2,
-            )
-            for i, array in items_dict["DatetimeTZBlock"]
-        ]
-        blocks.extend(dttz_blocks)
+    nbs = []
+    for (_, dtype), tup_block in grouper:
+        block_type = get_block_type(None, dtype)
 
-    if len(items_dict["ObjectBlock"]) > 0:
-        object_blocks = simple_blockify(
-            items_dict["ObjectBlock"], np.object_, consolidate=consolidate
-        )
-        blocks.extend(object_blocks)
+        if isinstance(dtype, np.dtype):
+            is_dtlike = dtype.kind in ["m", "M"]
 
-    if len(items_dict["CategoricalBlock"]) > 0:
-        cat_blocks = [
-            CategoricalBlock(array, placement=BlockPlacement(i), ndim=2)
-            for i, array in items_dict["CategoricalBlock"]
-        ]
-        blocks.extend(cat_blocks)
+            if issubclass(dtype.type, (str, bytes)):
+                dtype = np.dtype(object)
 
-    if len(items_dict["ExtensionBlock"]):
-        external_blocks = [
-            ExtensionBlock(array, placement=BlockPlacement(i), ndim=2)
-            for i, array in items_dict["ExtensionBlock"]
-        ]
+            values, placement = _stack_arrays(list(tup_block), dtype)
+            if is_dtlike:
+                values = ensure_wrapped_if_datetimelike(values)
+            blk = block_type(values, placement=BlockPlacement(placement), ndim=2)
+            nbs.append(blk)
 
-        blocks.extend(external_blocks)
+        elif is_1d_only_ea_dtype(dtype):
+            dtype_blocks = [
+                block_type(x[1], placement=BlockPlacement(x[0]), ndim=2)
+                for x in tup_block
+            ]
+            nbs.extend(dtype_blocks)
 
-    return blocks
+        else:
+            dtype_blocks = [
+                block_type(
+                    ensure_block_shape(x[1], 2), placement=BlockPlacement(x[0]), ndim=2
+                )
+                for x in tup_block
+            ]
+            nbs.extend(dtype_blocks)
+    return nbs
 
 
 def simple_blockify(tuples, dtype, consolidate: bool) -> list[Block]:
@@ -1970,11 +1959,16 @@ def _tuples_to_blocks_no_consolidate(tuples, dtype: DtypeObj | None) -> list[Blo
     if dtype is not None:
         return [
             new_block(
-                np.atleast_2d(x[1].astype(dtype, copy=False)), placement=x[0], ndim=2
+                ensure_block_shape(x[1].astype(dtype, copy=False), ndim=2),
+                placement=x[0],
+                ndim=2,
             )
             for x in tuples
         ]
-    return [new_block(np.atleast_2d(x[1]), placement=x[0], ndim=2) for x in tuples]
+    return [
+        new_block(ensure_block_shape(x[1], ndim=2), placement=x[0], ndim=2)
+        for x in tuples
+    ]
 
 
 def _stack_arrays(tuples, dtype: np.dtype):

pandas/tests/io/pytables/test_categorical.py

@@ -83,8 +83,9 @@ def test_categorical(setup_path):
         # Make sure the metadata is OK
         info = store.info()
         assert "/df2   " in info
-        # assert '/df2/meta/values_block_0/meta' in info
-        assert "/df2/meta/values_block_1/meta" in info
+        # df2._mgr.blocks[0] and df2._mgr.blocks[2] are Categorical
+        assert "/df2/meta/values_block_0/meta" in info
+        assert "/df2/meta/values_block_2/meta" in info
 
         # unordered
         _maybe_remove(store, "s2")