PERF: concat_compat (#52316)

Author: jbrockmendel

pandas/_libs/lib.pyx

@@ -3060,6 +3060,9 @@ def dtypes_all_equal(list types not None) -> bool:
  """
  first = types[0]
  for t in types[1:]:
+ if t is first:
+ # Fastpath can provide a nice boost for EADtypes
+ continue
  try:
  if not t == first:
  return False

pandas/core/dtypes/concat.py

@@ -11,6 +11,8 @@
 
 import numpy as np
 
+from pandas._libs import lib
+
 from pandas.core.dtypes.astype import astype_array
 from pandas.core.dtypes.cast import (
  common_dtype_categorical_compat,
@@ -60,6 +62,24 @@ def concat_compat(
  -------
  a single array, preserving the combined dtypes
  """
+ if len(to_concat) and lib.dtypes_all_equal([obj.dtype for obj in to_concat]):
+ # fastpath!
+ obj = to_concat[0]
+ if isinstance(obj, np.ndarray):
+ to_concat_arrs = cast("Sequence[np.ndarray]", to_concat)
+ return np.concatenate(to_concat_arrs, axis=axis)
+
+ to_concat_eas = cast("Sequence[ExtensionArray]", to_concat)
+ if ea_compat_axis:
+ # We have 1D objects, that don't support axis keyword
+ return obj._concat_same_type(to_concat_eas)
+ elif axis == 0:
+ return obj._concat_same_type(to_concat_eas)
+ else:
+ # e.g. DatetimeArray
+ # NB: We are assuming here that ensure_wrapped_if_arraylike has
+ # been called where relevant.
+ return obj._concat_same_type(to_concat_eas, axis=axis)
 
  # filter empty arrays
  # 1-d dtypes always are included here

pandas/core/internals/concat.py

@@ -230,10 +230,12 @@ def concatenate_managers(
  # we can use np.concatenate, which is more performant
  # than concat_compat
  values = np.concatenate(vals, axis=1)
- else:
+ elif is_1d_only_ea_dtype(blk.dtype):
  # TODO(EA2D): special-casing not needed with 2D EAs
- values = concat_compat(vals, axis=1)
+ values = concat_compat(vals, axis=1, ea_compat_axis=True)
  values = ensure_block_shape(values, ndim=2)
+ else:
+ values = concat_compat(vals, axis=1)
 
  values = ensure_wrapped_if_datetimelike(values)
 
@@ -541,6 +543,9 @@ def get_reindexed_values(self, empty_dtype: DtypeObj, upcasted_na) -> ArrayLike:
  # if we did, the missing_arr.fill would cast to gibberish
  missing_arr = np.empty(self.shape, dtype=empty_dtype)
  missing_arr.fill(fill_value)
+
+ if empty_dtype.kind in "mM":
+ missing_arr = ensure_wrapped_if_datetimelike(missing_arr)
  return missing_arr
 
  if (not self.indexers) and (not self.block._can_consolidate):