ENH: 2D compat for DTA tz_localize, to_period

pandas/core/arrays/_mixins.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from functools import wraps
 from typing import Any, Optional, Sequence, Type, TypeVar, Union
 
 import numpy as np
@@ -27,6 +28,26 @@
 )
 
 
+def ravel_compat(meth):
+    """
+    Decorator to ravel a 2D array before passing it to a cython operation,
+    then reshape the result to our own shape.
+    """
+
+    @wraps(meth)
+    def method(self, *args, **kwargs):
+        if self.ndim == 1:
+            return meth(self, *args, **kwargs)
+
+        flags = self._ndarray.flags
+        flat = self.ravel("K")
+        result = meth(flat, *args, **kwargs)
+        order = "F" if flags.f_contiguous else "C"
+        return result.reshape(self.shape, order=order)
+
+    return method
+
+
 class NDArrayBackedExtensionArray(ExtensionArray):
     """
     ExtensionArray that is backed by a single NumPy ndarray.

pandas/core/arrays/datetimelike.py

@@ -64,7 +64,7 @@
 from pandas.core import nanops, ops
 from pandas.core.algorithms import checked_add_with_arr, isin, unique1d, value_counts
 from pandas.core.arraylike import OpsMixin
-from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
+from pandas.core.arrays._mixins import NDArrayBackedExtensionArray, ravel_compat
 import pandas.core.common as com
 from pandas.core.construction import array, extract_array
 from pandas.core.indexers import check_array_indexer, check_setitem_lengths
@@ -679,6 +679,9 @@ def value_counts(self, dropna: bool = False):
         -------
         Series
         """
+        if self.ndim != 1:
+            raise NotImplementedError
+
         from pandas import Index, Series
 
         if dropna:
@@ -694,6 +697,7 @@ def value_counts(self, dropna: bool = False):
         )
         return Series(result._values, index=index, name=result.name)
 
+    @ravel_compat
     def map(self, mapper):
         # TODO(GH-23179): Add ExtensionArray.map
         # Need to figure out if we want ExtensionArray.map first.
@@ -820,6 +824,9 @@ def freq(self, value):
             value = to_offset(value)
             self._validate_frequency(self, value)
 
+            if self.ndim > 1:
+                raise ValueError("Cannot set freq with ndim > 1")
+
         self._freq = value
 
     @property
@@ -918,7 +925,7 @@ def _is_monotonic_decreasing(self) -> bool:
 
     @property
     def _is_unique(self) -> bool:
-        return len(unique1d(self.asi8)) == len(self)
+        return len(unique1d(self.asi8.ravel("K"))) == self.size
 
     # ------------------------------------------------------------------
     # Arithmetic Methods

pandas/core/arrays/datetimes.py

@@ -612,14 +612,15 @@ def astype(self, dtype, copy=True):
     # -----------------------------------------------------------------
     # Rendering Methods
 
+    @dtl.ravel_compat
     def _format_native_types(self, na_rep="NaT", date_format=None, **kwargs):
         from pandas.io.formats.format import get_format_datetime64_from_values
 
         fmt = get_format_datetime64_from_values(self, date_format)
 
         return tslib.format_array_from_datetime(
-            self.asi8.ravel(), tz=self.tz, format=fmt, na_rep=na_rep
-        ).reshape(self.shape)
+            self.asi8, tz=self.tz, format=fmt, na_rep=na_rep
+        )
 
     # -----------------------------------------------------------------
     # Comparison Methods
@@ -819,6 +820,7 @@ def tz_convert(self, tz):
         dtype = tz_to_dtype(tz)
         return self._simple_new(self.asi8, dtype=dtype, freq=self.freq)
 
+    @dtl.ravel_compat
     def tz_localize(self, tz, ambiguous="raise", nonexistent="raise"):
         """
         Localize tz-naive Datetime Array/Index to tz-aware
@@ -1051,6 +1053,7 @@ def normalize(self):
         new_values = normalize_i8_timestamps(self.asi8, self.tz)
         return type(self)(new_values)._with_freq("infer").tz_localize(self.tz)
 
+    @dtl.ravel_compat
     def to_period(self, freq=None):
         """
         Cast to PeriodArray/Index at a particular frequency.

pandas/core/arrays/period.py

@@ -562,6 +562,7 @@ def _formatter(self, boxed: bool = False):
             return str
         return "'{}'".format
 
+    @dtl.ravel_compat
     def _format_native_types(self, na_rep="NaT", date_format=None, **kwargs):
         """
         actually format my specific types

pandas/core/arrays/timedeltas.py

@@ -400,11 +400,12 @@ def _formatter(self, boxed=False):
 
         return get_format_timedelta64(self, box=True)
 
+    @dtl.ravel_compat
     def _format_native_types(self, na_rep="NaT", date_format=None, **kwargs):
         from pandas.io.formats.format import get_format_timedelta64
 
         formatter = get_format_timedelta64(self._data, na_rep)
-        return np.array([formatter(x) for x in self._data.ravel()]).reshape(self.shape)
+        return np.array([formatter(x) for x in self._data])
 
     # ----------------------------------------------------------------
     # Arithmetic Methods

pandas/tests/arrays/test_datetimelike.py

@@ -720,6 +720,15 @@ def test_to_period(self, datetime_index, freqstr):
         #  an EA-specific tm.assert_ function
         tm.assert_index_equal(pd.Index(result), pd.Index(expected))
 
+    def test_to_period_2d(self, arr1d):
+        arr2d = arr1d.reshape(1, -1)
+
+        warn = None if arr1d.tz is None else UserWarning
+        with tm.assert_produces_warning(warn):
+            result = arr2d.to_period("D")
+            expected = arr1d.to_period("D").reshape(1, -1)
+        tm.assert_period_array_equal(result, expected)
+
     @pytest.mark.parametrize("propname", pd.DatetimeIndex._bool_ops)
     def test_bool_properties(self, arr1d, propname):
         # in this case _bool_ops is just `is_leap_year`

pandas/tests/arrays/test_datetimes.py

@@ -449,6 +449,17 @@ def test_shift_requires_tzmatch(self):
         with pytest.raises(ValueError, match=msg):
             dta.shift(1, fill_value=fill_value)
 
+    def test_tz_localize_t2d(self):
+        dti = pd.date_range("1994-05-12", periods=12, tz="US/Pacific")
+        dta = dti._data.reshape(3, 4)
+        result = dta.tz_localize(None)
+
+        expected = dta.ravel().tz_localize(None).reshape(dta.shape)
+        tm.assert_datetime_array_equal(result, expected)
+
+        roundtrip = expected.tz_localize("US/Pacific")
+        tm.assert_datetime_array_equal(roundtrip, dta)
+
 
 class TestSequenceToDT64NS:
     def test_tz_dtype_mismatch_raises(self):