PERF: DataFrame.transpose with dt64tz

doc/source/whatsnew/v1.3.0.rst

@@ -673,6 +673,7 @@ Performance improvements
 - Performance improvement for concatenation of data with type :class:`CategoricalDtype` (:issue:`40193`)
 - Performance improvement in :meth:`.GroupBy.cummin` and :meth:`.GroupBy.cummax` with nullable data types (:issue:`37493`)
 - Performance improvement in :meth:`Series.nunique` with nan values (:issue:`40865`)
+- Performance improvement in :meth:`DataFrame.transpose`, :meth:`Series.unstack` with ``DatetimeTZDtype`` (:issue:`40149`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/array_algos/take.py

@@ -16,7 +16,10 @@
 from pandas._typing import ArrayLike
 
 from pandas.core.dtypes.cast import maybe_promote
-from pandas.core.dtypes.common import ensure_platform_int
+from pandas.core.dtypes.common import (
+ 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
@@ -91,12 +94,14 @@ def take_nd(
 
  if not isinstance(arr, np.ndarray):
  # i.e. ExtensionArray,
- if arr.ndim == 2:
- # e.g. DatetimeArray, TimedeltArray
+ # 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, fill_value=fill_value, allow_fill=allow_fill, axis=axis
  )
+
  return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
 
  arr = np.asarray(arr)

pandas/core/frame.py

@@ -101,6 +101,7 @@
  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,
@@ -139,7 +140,11 @@
 )
 from pandas.core.array_algos.take import take_2d_multi
 from pandas.core.arraylike import OpsMixin
-from pandas.core.arrays import ExtensionArray
+from pandas.core.arrays import (
+ DatetimeArray,
+ ExtensionArray,
+ TimedeltaArray,
+)
 from pandas.core.arrays.sparse import SparseFrameAccessor
 from pandas.core.construction import (
  extract_array,
@@ -852,6 +857,28 @@ def _can_fast_transpose(self) -> bool:
  # TODO(EA2D) special case would be unnecessary with 2D EAs
  return not is_1d_only_ea_dtype(dtype)
 
+ @property
+ def _values_compat(self) -> 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
+
+ # more generally, whatever we allow in NDArrayBackedExtensionBlock
+ arr = cast("DatetimeArray | TimedeltaArray", arr)
+ return arr.T
+
  # ----------------------------------------------------------------------
  # Rendering Methods
 
@@ -3292,7 +3319,18 @@ 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_vals = self._values_compat.T
+ if copy:
+ new_vals = new_vals.copy()
+
+ result = self._constructor(new_vals, 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()
@@ -9760,8 +9798,9 @@ def func(values: np.ndarray):
 
  def blk_func(values, axis=1):
  if isinstance(values, ExtensionArray):
- if values.ndim == 2:
- # i.e. DatetimeArray, TimedeltaArray
+ 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:

pandas/tests/frame/methods/test_transpose.py

@@ -90,3 +90,16 @@ def test_transpose_get_view(self, float_frame):
  dft.values[:, 5:10] = 5
 
  assert (float_frame.values[5:10] == 5).all()
+
+ @td.skip_array_manager_invalid_test
+ def test_transpose_get_view_dt64tzget_view(self):
+ dti = date_range("2016-01-01", periods=6, tz="US/Pacific")
+ arr = dti._data.reshape(3, 2)
+ df = DataFrame(arr)
+ assert df._mgr.nblocks == 1
+
+ result = df.T
+ assert result._mgr.nblocks == 1
+
+ rtrip = result._mgr.blocks[0].values
+ assert np.shares_memory(arr._data, rtrip._data)