PERF: DataFrame.transpose with dt64tz

asv_bench/benchmarks/reshape.py

@@ -53,6 +53,42 @@ def time_unstack(self):
         self.df.unstack(1)
 
 
+class ReshapeExtensionDtype:
+
+    params = ["datetime64[ns, US/Pacific]", "Period[s]"]
+    param_names = ["dtype"]
+
+    def setup(self, dtype):
+        lev = pd.Index(list("ABCDEFGHIJ"))
+        ri = pd.Index(range(1000))
+        mi = MultiIndex.from_product([lev, ri], names=["foo", "bar"])
+
+        index = date_range("2016-01-01", periods=10000, freq="s", tz="US/Pacific")
+        if dtype == "Period[s]":
+            index = index.tz_localize(None).to_period("s")
+
+        ser = pd.Series(index, index=mi)
+        df = ser.unstack("bar")
+        # roundtrips -> df.stack().equals(ser)
+
+        self.ser = ser
+        self.df = df
+
+    def time_stack(self, dtype):
+        self.df.stack()
+
+    def time_unstack_fast(self, dtype):
+        # last level -> doesnt have to make copies
+        self.ser.unstack("bar")
+
+    def time_unstack_slow(self, dtype):
+        # first level -> must make copies
+        self.ser.unstack("foo")
+
+    def time_transpose(self, dtype):
+        self.df.T
+
+
 class Unstack:
 
     params = ["int", "category"]

pandas/core/array_algos/take.py

@@ -1,6 +1,10 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    cast,
+)
 
 import numpy as np
 
@@ -14,11 +18,15 @@
 from pandas.core.dtypes.common import (
     ensure_int64,
     ensure_platform_int,
+    is_strict_ea,
 )
 from pandas.core.dtypes.missing import na_value_for_dtype
 
 from pandas.core.construction import ensure_wrapped_if_datetimelike
 
+if TYPE_CHECKING:
+    from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
+
 
 def take_nd(
     arr: ArrayLike,
@@ -66,6 +74,13 @@ def take_nd(
     if not isinstance(arr, np.ndarray):
         # i.e. ExtensionArray,
         # includes for EA to catch DatetimeArray, TimedeltaArray
+        if not is_strict_ea(arr):
+            # i.e. DatetimeArray, TimedeltaArray
+            arr = cast("NDArrayBackedExtensionArray", arr)
+            return arr.take(
+                indexer, axis=axis, fill_value=fill_value, allow_fill=allow_fill
+            )
+
         return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
 
     arr = np.asarray(arr)

pandas/core/arrays/_mixins.py

@@ -20,7 +20,10 @@
     cache_readonly,
     doc,
 )
-from pandas.util._validators import validate_fillna_kwargs
+from pandas.util._validators import (
+    validate_bool_kwarg,
+    validate_fillna_kwargs,
+)
 
 from pandas.core.dtypes.common import is_dtype_equal
 from pandas.core.dtypes.missing import array_equivalent
@@ -35,6 +38,7 @@
 from pandas.core.arrays.base import ExtensionArray
 from pandas.core.construction import extract_array
 from pandas.core.indexers import check_array_indexer
+from pandas.core.sorting import nargminmax
 
 NDArrayBackedExtensionArrayT = TypeVar(
     "NDArrayBackedExtensionArrayT", bound="NDArrayBackedExtensionArray"
@@ -182,6 +186,22 @@ def equals(self, other) -> bool:
     def _values_for_argsort(self):
         return self._ndarray
 
+    # Signature of "argmin" incompatible with supertype "ExtensionArray"
+    def argmin(self, axis: int = 0, skipna: bool = True):  # type:ignore[override]
+        # override base class by adding axis keyword
+        validate_bool_kwarg(skipna, "skipna")
+        if not skipna and self.isna().any():
+            raise NotImplementedError
+        return nargminmax(self, "argmin", axis=axis)
+
+    # Signature of "argmax" incompatible with supertype "ExtensionArray"
+    def argmax(self, axis: int = 0, skipna: bool = True):  # type:ignore[override]
+        # override base class by adding axis keyword
+        validate_bool_kwarg(skipna, "skipna")
+        if not skipna and self.isna().any():
+            raise NotImplementedError
+        return nargminmax(self, "argmax", axis=axis)
+
     def copy(self: NDArrayBackedExtensionArrayT) -> NDArrayBackedExtensionArrayT:
         new_data = self._ndarray.copy()
         return self._from_backing_data(new_data)
@@ -278,7 +298,7 @@ def fillna(
 
         if mask.any():
             if method is not None:
-                func = missing.get_fill_func(method)
+                func = missing.get_fill_func(method, ndim=self.ndim)
                 new_values = func(self._ndarray.copy(), limit=limit, mask=mask)
                 # TODO: PandasArray didn't used to copy, need tests for this
                 new_values = self._from_backing_data(new_values)

pandas/core/dtypes/common.py

@@ -1495,6 +1495,31 @@ def is_extension_type(arr) -> bool:
     return False
 
 
+def is_strict_ea(obj):
+    """
+    ExtensionArray that does not support 2D, or more specifically that does
+    not use HybridBlock.
+    """
+    from pandas.core.arrays import (
+        DatetimeArray,
+        ExtensionArray,
+        TimedeltaArray,
+    )
+
+    return isinstance(obj, ExtensionArray) and not isinstance(
+        obj, (DatetimeArray, TimedeltaArray)
+    )
+
+
+def is_ea_dtype(dtype) -> bool:
+    """
+    Analogue to is_extension_array_dtype but excluding DatetimeTZDtype.
+    """
+    # Note: if other EA dtypes are ever held in HybridBlock, exclude those
+    #  here too.
+    return is_extension_array_dtype(dtype) and not is_datetime64tz_dtype(dtype)
+
+
 def is_extension_array_dtype(arr_or_dtype) -> bool:
     """
     Check if an object is a pandas extension array type.

pandas/core/dtypes/concat.py

@@ -108,11 +108,15 @@ def is_nonempty(x) -> bool:
         to_concat = non_empties
 
     kinds = {obj.dtype.kind for obj in to_concat}
+    _contains_datetime = any(kind in ["m", "M"] for kind in kinds)
 
     all_empty = not len(non_empties)
     single_dtype = len({x.dtype for x in to_concat}) == 1
     any_ea = any(is_extension_array_dtype(x.dtype) for x in to_concat)
 
+    if _contains_datetime:
+        return _concat_datetime(to_concat, axis=axis)
+
     if any_ea:
         # we ignore axis here, as internally concatting with EAs is always
         # for axis=0
@@ -124,10 +128,7 @@ def is_nonempty(x) -> bool:
             cls = type(to_concat[0])
             return cls._concat_same_type(to_concat)
         else:
-            return np.concatenate(to_concat)
-
-    elif any(kind in ["m", "M"] for kind in kinds):
-        return _concat_datetime(to_concat, axis=axis)
+            return np.concatenate(to_concat, axis=axis)
 
     elif all_empty:
         # we have all empties, but may need to coerce the result dtype to
@@ -344,14 +345,5 @@ def _concat_datetime(to_concat, axis=0):
         #  in Timestamp/Timedelta
         return _concatenate_2d([x.astype(object) for x in to_concat], axis=axis)
 
-    if axis == 1:
-        # TODO(EA2D): kludge not necessary with 2D EAs
-        to_concat = [x.reshape(1, -1) if x.ndim == 1 else x for x in to_concat]
-
     result = type(to_concat[0])._concat_same_type(to_concat, axis=axis)
-
-    if result.ndim == 2 and is_extension_array_dtype(result.dtype):
-        # TODO(EA2D): kludge not necessary with 2D EAs
-        assert result.shape[0] == 1
-        result = result[0]
     return result

pandas/core/frame.py

@@ -110,6 +110,7 @@
     is_datetime64_any_dtype,
     is_dict_like,
     is_dtype_equal,
+    is_ea_dtype,
     is_extension_array_dtype,
     is_float,
     is_float_dtype,
@@ -121,6 +122,7 @@
     is_object_dtype,
     is_scalar,
     is_sequence,
+    is_strict_ea,
     pandas_dtype,
 )
 from pandas.core.dtypes.missing import (
@@ -756,7 +758,28 @@ def _can_fast_transpose(self) -> bool:
         if len(blocks) != 1:
             return False
 
-        return not self._mgr.any_extension_types
+        dtype = blocks[0].dtype
+        # TODO(EA2D) special case would be unnecessary with 2D EAs
+        return not is_ea_dtype(dtype)
+
+    @property
+    def _values_compat(self) -> ArrayLike:
+        """
+        Analogue to ._values that may return a 2D ExtensionArray.
+        """
+        mgr = self._mgr
+        if isinstance(mgr, ArrayManager):
+            return self._values
+
+        blocks = mgr.blocks
+        if len(blocks) != 1:
+            return self._values
+
+        arr = blocks[0].values
+        if arr.ndim == 1:
+            # non-2D ExtensionArray
+            return self._values
+        return arr.T
 
     # ----------------------------------------------------------------------
     # Rendering Methods
@@ -3164,16 +3187,32 @@ def transpose(self, *args, copy: bool = False) -> DataFrame:
         # construct the args
 
         dtypes = list(self.dtypes)
-        if self._is_homogeneous_type and dtypes and is_extension_array_dtype(dtypes[0]):
+
+        if self._can_fast_transpose:
+            # Note: tests pass without this, but this improves perf quite a bit.
+            # error: "ArrayLike" has no attribute "T"
+            new_values = self._values_compat.T  # type:ignore[attr-defined]
+            if copy:
+                new_values = new_values.copy()
+
+            result = self._constructor(
+                new_values, index=self.columns, columns=self.index
+            )
+
+        elif (
+            self._is_homogeneous_type and dtypes and is_extension_array_dtype(dtypes[0])
+        ):
             # We have EAs with the same dtype. We can preserve that dtype in transpose.
             dtype = dtypes[0]
+
             arr_type = dtype.construct_array_type()
             values = self.values
 
             new_values = [arr_type._from_sequence(row, dtype=dtype) for row in values]
             result = self._constructor(
                 dict(zip(self.index, new_values)), index=self.columns
             )
+            # TODO: what if index is non-unique? (not specific to EA2D)
 
         else:
             new_values = self.values.T
@@ -9185,6 +9224,8 @@ def func(values: np.ndarray):
 
         def blk_func(values, axis=1):
             if isinstance(values, ExtensionArray):
+                if not is_strict_ea(values):
+                    return values._reduce(name, axis=1, skipna=skipna, **kwds)
                 return values._reduce(name, skipna=skipna, **kwds)
             else:
                 return op(values, axis=axis, skipna=skipna, **kwds)

pandas/core/groupby/generic.py

@@ -64,6 +64,7 @@
     is_interval_dtype,
     is_numeric_dtype,
     is_scalar,
+    is_strict_ea,
     needs_i8_conversion,
 )
 from pandas.core.dtypes.missing import (
@@ -81,10 +82,7 @@
     validate_func_kwargs,
 )
 from pandas.core.apply import GroupByApply
-from pandas.core.arrays import (
-    Categorical,
-    ExtensionArray,
-)
+from pandas.core.arrays import Categorical
 from pandas.core.base import (
     DataError,
     SpecificationError,
@@ -1128,7 +1126,7 @@ def py_fallback(values: ArrayLike) -> ArrayLike:
             obj: FrameOrSeriesUnion
 
             # call our grouper again with only this block
-            if isinstance(values, ExtensionArray) or values.ndim == 1:
+            if is_strict_ea(values) or values.ndim == 1:
                 # TODO(EA2D): special case not needed with 2D EAs
                 obj = Series(values)
             else:

pandas/core/groupby/ops.py

@@ -71,6 +71,7 @@
 )
 
 import pandas.core.algorithms as algorithms
+from pandas.core.arrays import ExtensionArray
 from pandas.core.base import SelectionMixin
 import pandas.core.common as com
 from pandas.core.frame import DataFrame
@@ -209,7 +210,9 @@ def apply(self, f: F, data: FrameOrSeries, axis: int = 0):
         result_values = None
 
         sdata: FrameOrSeries = splitter._get_sorted_data()
-        if sdata.ndim == 2 and np.any(sdata.dtypes.apply(is_extension_array_dtype)):
+        if sdata.ndim == 2 and any(
+            isinstance(x, ExtensionArray) for x in sdata._iter_column_arrays()
+        ):
             # calling splitter.fast_apply will raise TypeError via apply_frame_axis0
             #  if we pass EA instead of ndarray
             #  TODO: can we have a workaround for EAs backed by ndarray?