PERF: DataFrame.transpose with dt64tz

pandas/core/array_algos/quantile.py

@@ -37,7 +37,8 @@ def quantile_compat(values: ArrayLike, qs, interpolation: str, axis: int) -> Arr
     -------
     np.ndarray or ExtensionArray
     """
-    if isinstance(values, np.ndarray):
+    if isinstance(values.dtype, np.dtype):
+        # i.e. np.ndarray, DatetimeArray, TimedeltaArray
         fill_value = na_value_for_dtype(values.dtype, compat=False)
         mask = isna(values)
         return quantile_with_mask(values, mask, fill_value, qs, interpolation, axis)
@@ -162,6 +163,10 @@ def quantile_ea_compat(
             assert result.shape == (1,), result.shape
             result = type(orig)._from_factorized(result, orig)
 
+        elif orig.ndim == 2:
+            # i.e. DatetimeArray
+            result = type(orig)._from_factorized(result, orig)
+
         else:
             assert result.shape == (1, len(qs)), result.shape
             result = type(orig)._from_factorized(result[0], orig)

pandas/core/array_algos/take.py

@@ -4,6 +4,7 @@
 from typing import (
     TYPE_CHECKING,
     Optional,
+    cast,
     overload,
 )
 
@@ -19,12 +20,14 @@
 from pandas.core.dtypes.common import (
     ensure_int64,
     ensure_platform_int,
+    is_1d_only_ea_obj,
 )
 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
     from pandas.core.arrays.base import ExtensionArray
 
 
@@ -98,6 +101,13 @@ def take_nd(
     if not isinstance(arr, np.ndarray):
         # i.e. ExtensionArray,
         # includes for EA to catch DatetimeArray, TimedeltaArray
+        if not is_1d_only_ea_obj(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"
@@ -185,6 +189,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)

pandas/core/dtypes/common.py

@@ -1499,6 +1499,33 @@ def is_extension_type(arr) -> bool:
     return False
 
 
+def is_1d_only_ea_obj(obj) -> bool:
+    """
+    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_1d_only_ea_dtype(dtype: Optional[DtypeObj]) -> bool:
+    """
+    Analogue to is_extension_array_dtype but excluding DatetimeTZDtype.
+    """
+    # Note: if other EA dtypes are ever held in HybridBlock, exclude those
+    #  here too.
+    # NB: need to check DatetimeTZDtype and not is_datetime64tz_dtype
+    #  to exclude ArrowTimestampUSDtype
+    return isinstance(dtype, ExtensionDtype) and not isinstance(dtype, DatetimeTZDtype)
+
+
 def is_extension_array_dtype(arr_or_dtype) -> bool:
     """
     Check if an object is a pandas extension array type.

pandas/core/dtypes/concat.py

@@ -112,11 +112,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
@@ -130,9 +134,6 @@ def is_nonempty(x) -> bool:
         else:
             return np.concatenate(to_concat)
 
-    elif any(kind in ["m", "M"] for kind in kinds):
-        return _concat_datetime(to_concat, axis=axis)
-
     elif all_empty:
         # we have all empties, but may need to coerce the result dtype to
         # object if we have non-numeric type operands (numpy would otherwise
@@ -348,14 +349,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

@@ -105,6 +105,8 @@
     ensure_int64,
     ensure_platform_int,
     infer_dtype_from_object,
+    is_1d_only_ea_dtype,
+    is_1d_only_ea_obj,
     is_bool_dtype,
     is_dataclass,
     is_datetime64_any_dtype,
@@ -213,6 +215,10 @@
         TimestampConvertibleTypes,
     )
 
+    from pandas.core.arrays import (
+        DatetimeArray,
+        TimedeltaArray,
+    )
     from pandas.core.groupby.generic import DataFrameGroupBy
     from pandas.core.resample import Resampler
 
@@ -884,7 +890,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_1d_only_ea_dtype(dtype)
+
+    @property
+    def _values_compat(self) -> Union[np.ndarray, DatetimeArray, TimedeltaArray]:
+        """
+        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
@@ -3314,7 +3341,20 @@ 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.
+            new_values = self._values_compat.T
+            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()
@@ -9471,6 +9511,10 @@ def func(values: np.ndarray):
 
         def blk_func(values, axis=1):
             if isinstance(values, ExtensionArray):
+                if not is_1d_only_ea_obj(values) and not isinstance(
+                    self._mgr, ArrayManager
+                ):
+                    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/ops.py

@@ -72,6 +72,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
@@ -267,7 +268,9 @@ def apply(self, f: F, data: FrameOrSeries, axis: int = 0):
         group_keys = self._get_group_keys()
         result_values = None
 
-        if data.ndim == 2 and np.any(data.dtypes.apply(is_extension_array_dtype)):
+        if data.ndim == 2 and any(
+            isinstance(x, ExtensionArray) for x in data._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?

pandas/core/internals/api.py

@@ -20,6 +20,7 @@
     Block,
     DatetimeTZBlock,
     check_ndim,
+    ensure_block_shape,
     extract_pandas_array,
     get_block_type,
 )
@@ -45,18 +46,23 @@ def make_block(
     # ExtensionDtype, None]"
     values, dtype = extract_pandas_array(values, dtype, ndim)  # type: ignore[arg-type]
 
+    needs_reshape = False
     if klass is None:
         dtype = dtype or values.dtype
         klass = get_block_type(values, dtype)
 
     elif klass is DatetimeTZBlock and not is_datetime64tz_dtype(values.dtype):
         # pyarrow calls get here
         values = DatetimeArray._simple_new(values, dtype=dtype)
+        needs_reshape = True
 
     if not isinstance(placement, BlockPlacement):
         placement = BlockPlacement(placement)
 
     ndim = _maybe_infer_ndim(values, placement, ndim)
+    if needs_reshape:
+        values = ensure_block_shape(values, ndim)
+
     check_ndim(values, placement, ndim)
     return klass(values, ndim=ndim, placement=placement)
 

pandas/core/internals/array_manager.py

@@ -130,6 +130,7 @@ def __init__(
         # Note: we are storing the axes in "_axes" in the (row, columns) order
         # which contrasts the order how it is stored in BlockManager
         self._axes = axes
+
         self.arrays = arrays
 
         if verify_integrity:
@@ -487,9 +488,9 @@ def apply_with_block(self: T, f, align_keys=None, swap_axis=True, **kwargs) -> T
             if isinstance(applied, list):
                 applied = applied[0]
             arr = applied.values
-            if self.ndim == 2:
-                if isinstance(arr, np.ndarray):
-                    arr = arr[0, :]
+            if self.ndim == 2 and arr.ndim == 2:
+                # 2D for np.ndarray or DatetimeArray/TimedeltaArray
+                arr = arr[0, :]
             result_arrays.append(arr)
 
         return type(self)(result_arrays, self._axes)