diff --git a/pandas/core/arrays/_mixins.py b/pandas/core/arrays/_mixins.py index 02214ff51b02a..b6938931e86af 100644 --- a/pandas/core/arrays/_mixins.py +++ b/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. diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index be9864731842d..ee1323b71f146 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/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 diff --git a/pandas/core/arrays/datetimes.py b/pandas/core/arrays/datetimes.py index 5fdfa62c393c4..b072ac3cec52e 100644 --- a/pandas/core/arrays/datetimes.py +++ b/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. diff --git a/pandas/core/arrays/period.py b/pandas/core/arrays/period.py index 257baf20ce911..40dd475e6b6f2 100644 --- a/pandas/core/arrays/period.py +++ b/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 diff --git a/pandas/core/arrays/timedeltas.py b/pandas/core/arrays/timedeltas.py index 93c9567380f7f..fe4eaa4b4bf19 100644 --- a/pandas/core/arrays/timedeltas.py +++ b/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 diff --git a/pandas/tests/arrays/test_datetimelike.py b/pandas/tests/arrays/test_datetimelike.py index c489aa5867632..7c093ebe00959 100644 --- a/pandas/tests/arrays/test_datetimelike.py +++ b/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` diff --git a/pandas/tests/arrays/test_datetimes.py b/pandas/tests/arrays/test_datetimes.py index 4addc0536848f..c8db0157ba219 100644 --- a/pandas/tests/arrays/test_datetimes.py +++ b/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):