REF: avoid special case in DTA/TDA.median, flesh out tests

pandas-dev · jreback · Nov 8, 2020 · Oct 26, 2020 · Oct 26, 2020 · Oct 26, 2020
commit 65fe4e460b607c41d62408fb664d35e21ca4f52b
diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py
@@ -1359,21 +1359,20 @@ def median(self, axis: Optional[int] = None, skipna: bool = True, *args, **kwarg
         if axis is not None and abs(axis) >= self.ndim:
             raise ValueError("abs(axis) must be less than ndim")
 
-        if self.size == 0:
-            if self.ndim == 1 or axis is None:
-                return NaT
-            shape = list(self.shape)
-            del shape[axis]
-            shape = [1 if x == 0 else x for x in shape]
-            result = np.empty(shape, dtype="i8")
-            result.fill(iNaT)
+        if is_period_dtype(self.dtype):
+            # pass datetime64 values to nanops to get correct NaT semantics
+            result = nanops.nanmedian(
+                self._ndarray.view("M8[ns]"), axis=axis, skipna=skipna
+            )
+            result = result.view("i8")
+            if axis is None or self.ndim == 1:
+                return self._box_func(result)
             return self._from_backing_data(result)
 
-        mask = self.isna()
-        result = nanops.nanmedian(self.asi8, axis=axis, skipna=skipna, mask=mask)
+        result = nanops.nanmedian(self._ndarray, axis=axis, skipna=skipna)
         if axis is None or self.ndim == 1:
             return self._box_func(result)
-        return self._from_backing_data(result.astype("i8"))
+        return self._from_backing_data(result)
 
 
 class DatelikeOps(DatetimeLikeArrayMixin):

diff --git a/pandas/core/nanops.py b/pandas/core/nanops.py
@@ -339,7 +339,12 @@ def _wrap_results(result, dtype: DtypeObj, fill_value=None):
             assert not isna(fill_value), "Expected non-null fill_value"
             if result == fill_value:
                 result = np.nan
-            result = Timestamp(result, tz=tz)
+            if tz is not None:
+                result = Timestamp(result, tz=tz)
+            elif isna(result):
+                result = np.datetime64("NaT", "ns")
+            else:
+                result = np.int64(result).view("datetime64[ns]")
         else:
             # If we have float dtype, taking a view will give the wrong result
             result = result.astype(dtype)

diff --git a/pandas/tests/arrays/test_datetimelike.py b/pandas/tests/arrays/test_datetimelike.py
@@ -4,7 +4,7 @@
 import pytest
 import pytz
 
-from pandas._libs import OutOfBoundsDatetime, Timestamp
+from pandas._libs import NaT, OutOfBoundsDatetime, Timestamp
 from pandas.compat.numpy import np_version_under1p18
 
 import pandas as pd
@@ -456,6 +456,54 @@ def test_shift_fill_int_deprecated(self):
         expected[1:] = arr[:-1]
         tm.assert_equal(result, expected)
 
+    def test_median(self, arr1d):
+        arr = arr1d
+        if len(arr) % 2 == 0:
+            # make it easier to define `expected`
+            arr = arr[:-1]
+
+        expected = arr[len(arr) // 2]
+
+        result = arr.median()
+        assert type(result) is type(expected)
+        assert result == expected
+
+        arr[len(arr) // 2] = NaT
+        if not isinstance(expected, Period):
+            expected = arr[len(arr) // 2 - 1 : len(arr) // 2 + 2].mean()
+
+        assert arr.median(skipna=False) is NaT
+
+        result = arr.median()
+        assert type(result) is type(expected)
+        assert result == expected
+
+        assert arr[:0].median() is NaT
+        assert arr[:0].median(skipna=False) is NaT
+
+        # 2d Case
+        arr2 = arr.reshape(-1, 1)
+
+        result = arr2.median(axis=None)
+        assert type(result) is type(expected)
+        assert result == expected
+
+        assert arr2.median(axis=None, skipna=False) is NaT
+
+        result = arr2.median(axis=0)
+        expected2 = type(arr)._from_sequence([expected], dtype=arr.dtype)
+        tm.assert_equal(result, e
8000
xpected2)
+
+        result = arr2.median(axis=0, skipna=False)
+        expected2 = type(arr)._from_sequence([NaT], dtype=arr.dtype)
+        tm.assert_equal(result, expected2)
+
+        result = arr2.median(axis=1)
+        tm.assert_equal(result, arr)
+
+        result = arr2.median(axis=1, skipna=False)
+        tm.assert_equal(result, arr)
+
 
 class TestDatetimeArray(SharedTests):
     index_cls = pd.DatetimeIndex
@@ -465,7 +513,7 @@ class TestDatetimeArray(SharedTests):
     @pytest.fixture
     def arr1d(self, tz_naive_fixture, freqstr):
         tz = tz_naive_fixture
-        dti = pd.date_range("2016-01-01 01:01:00", periods=3, freq=freqstr, tz=tz)
+        dti = pd.date_range("2016-01-01 01:01:00", periods=5, freq=freqstr, tz=tz)
         dta = dti._data
         return dta
 

diff --git a/pandas/tests/arrays/test_datetimes.py b/pandas/tests/arrays/test_datetimes.py
@@ -514,7 +514,7 @@ def test_median_empty(self, skipna, tz):
         tm.assert_equal(result, expected)
 
         result = arr.median(axis=1, skipna=skipna)
-        expected = type(arr)._from_sequence([pd.NaT], dtype=arr.dtype)
+        expected = type(arr)._from_sequence([], dtype=arr.dtype)
         tm.assert_equal(result, expected)
 
     def test_median(self, arr1d):