PERF: concat (#52291)

Author: jbrockmendel

pandas/core/dtypes/concat.py

@@ -41,6 +41,14 @@
     )
 
 
+def _is_nonempty(x, axis) -> bool:
+    # filter empty arrays
+    # 1-d dtypes always are included here
+    if x.ndim <= axis:
+        return True
+    return x.shape[axis] > 0
+
+
 def concat_compat(
     to_concat: Sequence[ArrayLike], axis: AxisInt = 0, ea_compat_axis: bool = False
 ) -> ArrayLike:
@@ -79,36 +87,34 @@ def concat_compat(
             # 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
-    def is_nonempty(x) -> bool:
-        if x.ndim <= axis:
-            return True
-        return x.shape[axis] > 0
+            return obj._concat_same_type(
+                # error: Unexpected keyword argument "axis" for "_concat_same_type"
+                # of "ExtensionArray"
+                to_concat_eas,
+                axis=axis,  # type: ignore[call-arg]
+            )
 
     # If all arrays are empty, there's nothing to convert, just short-cut to
     # the concatenation, #3121.
     #
     # Creating an empty array directly is tempting, but the winnings would be
     # marginal given that it would still require shape & dtype calculation and
     # np.concatenate which has them both implemented is compiled.
-    non_empties = [x for x in to_concat if is_nonempty(x)]
+    non_empties = [x for x in to_concat if _is_nonempty(x, axis)]
     if non_empties and axis == 0 and not ea_compat_axis:
         # ea_compat_axis see GH#39574
         to_concat = non_empties
 
     dtypes = {obj.dtype for obj in to_concat}
     kinds = {obj.dtype.kind for obj in to_concat}
     contains_datetime = any(
-        isinstance(dtype, (np.dtype, DatetimeTZDtype)) and dtype.kind in ["m", "M"]
+        isinstance(dtype, (np.dtype, DatetimeTZDtype)) and dtype.kind in "mM"
         for dtype in dtypes
     ) or any(isinstance(obj, ABCExtensionArray) and obj.ndim > 1 for obj in to_concat)
 
     all_empty = not len(non_empties)
-    single_dtype = len({x.dtype for x in to_concat}) == 1
-    any_ea = any(isinstance(x.dtype, ExtensionDtype) for x in to_concat)
+    single_dtype = len(dtypes) == 1
+    any_ea = any(isinstance(x, ExtensionDtype) for x in dtypes)
 
     if contains_datetime:
         return _concat_datetime(to_concat, axis=axis)
@@ -345,7 +351,7 @@ def _concat_datetime(to_concat: Sequence[ArrayLike], axis: AxisInt = 0) -> Array
 
     to_concat = [ensure_wrapped_if_datetimelike(x) for x in to_concat]
 
-    single_dtype = len({x.dtype for x in to_concat}) == 1
+    single_dtype = lib.dtypes_all_equal([x.dtype for x in to_concat])
 
     # multiple types, need to coerce to object
     if not single_dtype:

pandas/core/reshape/concat.py

@@ -480,9 +480,7 @@ def __init__(
         else:
             # filter out the empties if we have not multi-index possibilities
             # note to keep empty Series as it affect to result columns / name
-            non_empties = [
-                obj for obj in objs if sum(obj.shape) > 0 or isinstance(obj, ABCSeries)
-            ]
+            non_empties = [obj for obj in objs if sum(obj.shape) > 0 or obj.ndim == 1]
 
             if len(non_empties) and (
                 keys is None and names is None and levels is None and not self.intersect
@@ -495,19 +493,21 @@ def __init__(
         self.objs = objs
 
         # Standardize axis parameter to int
-        if isinstance(sample, ABCSeries):
+        if sample.ndim == 1:
             from pandas import DataFrame
 
             axis = DataFrame._get_axis_number(axis)
+            self._is_frame = False
+            self._is_series = True
         else:
             axis = sample._get_axis_number(axis)
+            self._is_frame = True
+            self._is_series = False
 
         # Need to flip BlockManager axis in the DataFrame special case
-        self._is_frame = isinstance(sample, ABCDataFrame)
         if self._is_frame:
             axis = sample._get_block_manager_axis(axis)
 
-        self._is_series = isinstance(sample, ABCSeries)
         if not 0 <= axis <= sample.ndim:
             raise AssertionError(
                 f"axis must be between 0 and {sample.ndim}, input was {axis}"
@@ -583,7 +583,8 @@ def get_result(self):
                 arrs = [ser._values for ser in self.objs]
 
                 res = concat_compat(arrs, axis=0)
-                result = cons(res, index=self.new_axes[0], name=name, dtype=res.dtype)
+                mgr = type(sample._mgr).from_array(res, index=self.new_axes[0])
+                result = cons(mgr, name=name, fastpath=True)
                 return result.__finalize__(self, method="concat")
 
             # combine as columns in a frame
@@ -666,7 +667,7 @@ def _get_concat_axis(self) -> Index:
                 num = 0
                 has_names = False
                 for i, x in enumerate(self.objs):
-                    if not isinstance(x, ABCSeries):
+                    if x.ndim != 1:
                         raise TypeError(
                             f"Cannot concatenate type 'Series' with "
                             f"object of type '{type(x).__name__}'"

pandas/tests/extension/test_numpy.py

@@ -346,18 +346,7 @@ def test_fillna_frame(self, data_missing):
 
 
 class TestReshaping(BaseNumPyTests, base.BaseReshapingTests):
-    @pytest.mark.parametrize(
-        "in_frame",
-        [
-            True,
-            pytest.param(
-                False,
-                marks=pytest.mark.xfail(reason="PandasArray inconsistently extracted"),
-            ),
-        ],
-    )
-    def test_concat(self, data, in_frame):
-        super().test_concat(data, in_frame)
+    pass
 
 
 class TestSetitem(BaseNumPyTests, base.BaseSetitemTests):