use Block.concat_same_type in concatenate_block_managers (concatting DataFrames)

Author: jorisvandenbossche

pandas/core/internals.py

@@ -312,13 +312,14 @@ def ftype(self):
     def merge(self, other):
         return _merge_blocks([self, other])
 
-    def concat_same_type(self, to_concat):
+    def concat_same_type(self, to_concat, placement=None):
         """
         Concatenate list of single blocks of the same type.
         """
-        values = np.concatenate([blk.values for blk in to_concat])
+        values = np.concatenate([blk.values for blk in to_concat],
+                                axis=self.ndim - 1)
         return self.make_block_same_class(
-            values, placement=slice(0, len(values), 1))
+            values, placement=placement or slice(0, len(values), 1))
 
     def reindex_axis(self, indexer, method=None, axis=1, fill_value=None,
                      limit=None, mask_info=None):
@@ -2415,7 +2416,7 @@ def to_native_types(self, slicer=None, na_rep='', quoting=None, **kwargs):
         # we are expected to return a 2-d ndarray
         return values.reshape(1, len(values))
 
-    def concat_same_type(self, to_concat):
+    def concat_same_type(self, to_concat, placement=None):
         """
         Concatenate list of single blocks of the same type.
         """
@@ -2424,9 +2425,9 @@ def concat_same_type(self, to_concat):
 
         if is_categorical_dtype(values.dtype):
             return self.make_block_same_class(
-                values, placement=slice(0, len(values), 1))
+                values, placement=placement or slice(0, len(values), 1))
         else:
-            return make_block(values, placement=slice(0, len(values), 1))
+            return make_block(values, placement=placement or slice(0, len(values), 1))
 
 
 class DatetimeBlock(DatetimeLikeBlockMixin, Block):
@@ -2705,7 +2706,7 @@ def shift(self, periods, axis=0, mgr=None):
         return [self.make_block_same_class(new_values,
                                            placement=self.mgr_locs)]
 
-    def concat_same_type(self, to_concat):
+    def concat_same_type(self, to_concat, placement=None):
         """
         Concatenate list of single blocks of the same type.
         """
@@ -2714,9 +2715,9 @@ def concat_same_type(self, to_concat):
 
         if is_datetimetz(values):
             return self.make_block_same_class(
-                values, placement=slice(0, len(values), 1))
+                values, placement=placement or slice(0, len(values), 1))
         else:
-            return make_block(values, placement=slice(0, len(values), 1))
+            return make_block(values, placement=placement or slice(0, len(values), 1))
 
 
 class SparseBlock(NonConsolidatableMixIn, Block):
@@ -2885,15 +2886,15 @@ def sparse_reindex(self, new_index):
         return self.make_block_same_class(values, sparse_index=new_index,
                                           placement=self.mgr_locs)
 
-    def concat_same_type(self, to_concat):
+    def concat_same_type(self, to_concat, placement=None):
         """
         Concatenate list of single blocks of the same type.
         """
         to_concat = [blk.values for blk in to_concat]
         values = _concat._concat_sparse(to_concat)
 
         return self.make_block_same_class(
-            values, placement=slice(0, len(values), 1))
+            values, placement=placement or slice(0, len(values), 1))
 
 
 def make_block(values, placement, klass=None, ndim=None, dtype=None,
@@ -5146,13 +5147,42 @@ def concatenate_block_managers(mgrs_indexers, axes, concat_axis, copy):
         [get_mgr_concatenation_plan(mgr, indexers)
          for mgr, indexers in mgrs_indexers], concat_axis)
 
-    blocks = [make_block(
-        concatenate_join_units(join_units, concat_axis, copy=copy),
-        placement=placement) for placement, join_units in concat_plan]
+    blocks = []
+
+    for placement, join_units in concat_plan:
+
+        if is_uniform_join_units(join_units):
+            b = join_units[0].block.concat_same_type(
+                [ju.block for ju in join_units], placement=placement)
+        else:
+            b = make_block(
+                concatenate_join_units(join_units, concat_axis, copy=copy),
+                placement=placement)
+        blocks.append(b)
 
     return BlockManager(blocks, axes)
 
 
+def is_uniform_join_units(join_units):
+    """
+    Check if the join units consist of blocks of uniform type that can
+    be concatenated using Block.concat_same_type instead of the generic
+    concatenate_join_units (which uses `_concat._concat_compat`).
+
+    """
+    return (
+        # all blocks need to have the same type
+        all([type(ju.block) is type(join_units[0].block) for ju in join_units])  # noqa
+        # no blocks that would get missing values (can lead to type upcasts)
+        and all([not ju.is_na for ju in join_units])
+        # no blocks with indexers (as then the dimensions do not fit)
+        and all([not ju.indexers for ju in join_units])
+        # disregard Panels
+        and all([ju.block.ndim <= 2 for ju in join_units])
+        # only use this path when there is something to concatenate
+        and len(join_units) > 1)
+
+
 def get_empty_dtype_and_na(join_units):
     """
     Return dtype and N/A values to use when concatenating specified units.

pandas/tests/internals/test_external_block.py

@@ -4,14 +4,25 @@
 import numpy as np
 
 import pandas as pd
-from pandas.core.internals import Block, BlockManager, SingleBlockManager
+from pandas.core.internals import Block, BlockManager, SingleBlockManager, NonConsolidatableMixIn
 
 
-class CustomBlock(Block):
+class CustomBlock(NonConsolidatableMixIn, Block):
+
+    _holder = np.ndarray
 
     def formatting_values(self):
         return np.array(["Val: {}".format(i) for i in self.values])
 
+    def concat_same_type(self, to_concat, placement=None):
+        """
+        Always concatenate disregarding self.ndim as the values are
+        always 1D in this custom Block
+        """
+        values = np.concatenate([blk.values for blk in to_concat])
+        return self.make_block_same_class(
+            values, placement=placement or slice(0, len(values), 1))
+
 
 def test_custom_repr():
     values = np.arange(3, dtype='int64')
@@ -23,7 +34,7 @@ def test_custom_repr():
     assert repr(s) == '0    Val: 0\n1    Val: 1\n2    Val: 2\ndtype: int64'
 
     # dataframe
-    block = CustomBlock(values.reshape(1, -1), placement=slice(0, 1))
+    block = CustomBlock(values, placement=slice(0, 1))
     blk_mgr = BlockManager([block], [['col'], range(3)])
     df = pd.DataFrame(blk_mgr)
     assert repr(df) == '      col\n0  Val: 0\n1  Val: 1\n2  Val: 2'
@@ -36,3 +47,15 @@ def test_concat_series():
 
     res = pd.concat([s, s])
     assert isinstance(res._data.blocks[0], CustomBlock)
+
+
+def test_concat_dataframe():
+    df = pd.DataFrame({'a': [1, 2, 3]})
+    blocks = df._data.blocks
+    values = np.arange(3, dtype='int64')
+    custom_block = CustomBlock(values, placement=slice(1, 2))
+    blocks = blocks + (custom_block, )
+    block_manager = BlockManager(blocks, [pd.Index(['a', 'b']), df.index])
+    df = pd.DataFrame(block_manager)
+    res = pd.concat([df, df])
+    assert isinstance(res._data.blocks[1], CustomBlock)