CLN: de-kludge quantile, make interpolate_with_fill understand datetime64

pandas/core/internals/blocks.py

@@ -1223,7 +1223,6 @@ def _interpolate_with_fill(
             fill_value=fill_value,
             dtype=self.dtype,
         )
-        values = self._try_coerce_result(values)
 
         blocks = [self.make_block_same_class(values, ndim=self.ndim)]
         return self._maybe_downcast(blocks, downcast)
@@ -1526,18 +1525,7 @@ def quantile(self, qs, interpolation="linear", axis=0):
         # We should always have ndim == 2 becase Series dispatches to DataFrame
         assert self.ndim == 2
 
-        if self.is_datetimetz:
-            # TODO: cleanup this special case.
-            # We need to operate on i8 values for datetimetz
-            # but `Block.get_values()` returns an ndarray of objects
-            # right now. We need an API for "values to do numeric-like ops on"
-            values = self.values.view("M8[ns]")
-
-            # TODO: NonConsolidatableMixin shape
-            # Usual shape inconsistencies for ExtensionBlocks
-            values = values[None, :]
-        else:
-            values = self.get_values()
+        values = self.get_values()
 
         is_empty = values.shape[axis] == 0
         orig_scalar = not is_list_like(qs)
@@ -1576,7 +1564,6 @@ def quantile(self, qs, interpolation="linear", axis=0):
             result = lib.item_from_zerodim(result)
 
         ndim = getattr(result, "ndim", None) or 0
-        result = self._try_coerce_result(result)
         return make_block(result, placement=np.arange(len(result)), ndim=ndim)
 
     def _replace_coerce(
@@ -1710,7 +1697,6 @@ def putmask(self, mask, new, align=True, inplace=False, axis=0, transpose=False)
         mask = _safe_reshape(mask, new_values.shape)
 
         new_values[mask] = new
-        new_values = self._try_coerce_result(new_values)
         return [self.make_block(values=new_values)]
 
     def _try_cast_result(self, result, dtype=None):
@@ -2293,13 +2279,6 @@ def _try_coerce_args(self, other):
 
         return other
 
-    def _try_coerce_result(self, result):
-        """ reverse of try_coerce_args """
-        if isinstance(result, np.ndarray) and result.dtype.kind == "i":
-            # needed for _interpolate_with_ffill
-            result = result.view("M8[ns]")
-        return result
-
     def to_native_types(
         self, slicer=None, na_rep=None, date_format=None, quoting=None, **kwargs
     ):
@@ -2477,15 +2456,7 @@ def _try_coerce_result(self, result):
                 result = self._holder._from_sequence(
                     result.astype(np.int64), freq=None, dtype=self.values.dtype
                 )
-            elif result.dtype == "M8[ns]":
-                # otherwise we get here via quantile and already have M8[ns]
-                result = self._holder._simple_new(
-                    result, freq=None, dtype=self.values.dtype
-                )
 
-        elif isinstance(result, np.datetime64):
-            # also for post-quantile
-            result = self._box_func(result)
         return result
 
     @property
@@ -2564,6 +2535,19 @@ def equals(self, other):
             return False
         return (self.values.view("i8") == other.values.view("i8")).all()
 
+    def quantile(self, qs, interpolation="linear", axis=0):
+        naive = self.values.view("M8[ns]")
+
+        # kludge for 2D block with 1D values
+        naive = naive.reshape(self.shape)
+
+        blk = self.make_block(naive)
+        res_blk = blk.quantile(qs, interpolation=interpolation, axis=axis)
+
+        # ravel is kludge for 2D block with 1D values, assumes column-like
+        aware = self._holder(res_blk.values.ravel(), dtype=self.dtype)
+        return self.make_block_same_class(aware, ndim=res_blk.ndim)
+
 
 class TimeDeltaBlock(DatetimeLikeBlockMixin, IntBlock):
     __slots__ = ()
@@ -2639,10 +2623,6 @@ def _try_coerce_args(self, other):
 
         return other
 
-    def _try_coerce_result(self, result):
-        """ reverse of try_coerce_args / try_operate """
-        return result
-
     def should_store(self, value):
         return issubclass(
             value.dtype.type, np.timedelta64

pandas/core/internals/managers.py

@@ -908,7 +908,7 @@ def fast_xs(self, loc):
             # Such assignment may incorrectly coerce NaT to None
             # result[blk.mgr_locs] = blk._slice((slice(None), loc))
             for i, rl in enumerate(blk.mgr_locs):
-                result[rl] = blk._try_coerce_result(blk.iget((i, loc)))
+                result[rl] = blk.iget((i, loc))
 
         if is_extension_array_dtype(dtype):
             result = dtype.construct_array_type()._from_sequence(result, dtype=dtype)

pandas/core/missing.py

@@ -463,14 +463,27 @@ def interpolate_2d(
     Perform an actual interpolation of values, values will be make 2-d if
     needed fills inplace, returns the result.
     """
+    if is_datetime64tz_dtype(values):
+        naive = values.view("M8[ns]")
+        result = interpolate_2d(
+            naive,
+            method=method,
+            axis=axis,
+            limit=limit,
+            fill_value=fill_value,
+            dtype=dtype,
+        )
+        return type(values)._from_sequence(result, dtype=values.dtype)
+
+    orig_values = values
 
     transf = (lambda x: x) if axis == 0 else (lambda x: x.T)
 
     # reshape a 1 dim if needed
     ndim = values.ndim
     if values.ndim == 1:
         if axis != 0:  # pragma: no cover
-            raise AssertionError("cannot interpolate on a ndim == 1 with " "axis != 0")
+            raise AssertionError("cannot interpolate on a ndim == 1 with axis != 0")
         values = values.reshape(tuple((1,) + values.shape))
 
     if fill_value is None:
@@ -490,6 +503,10 @@ def interpolate_2d(
     if ndim == 1:
         values = values[0]
 
+    if orig_values.dtype.kind == "M":
+        # convert float back to datetime64
+        values = values.astype(orig_values.dtype)
+
     return values