BUG: DTBlock/TDBlock.delete casting to ndarray (pandas-dev#40741)

Author: jbrockmendel

pandas/core/internals/blocks.py

@@ -278,6 +278,11 @@ def make_block_same_class(
         """ Wrap given values in a block of same type as self. """
         if placement is None:
             placement = self._mgr_locs
+
+        if values.dtype.kind == "m":
+            # TODO: remove this once fastparquet has stopped relying on it
+            values = ensure_wrapped_if_datetimelike(values)
+
         # We assume maybe_coerce_values has already been called
         return type(self)(values, placement=placement, ndim=self.ndim)
 
@@ -369,7 +374,6 @@ def set_inplace(self, locs, values):
         """
         self.values[locs] = values
 
-    @final
     def delete(self, loc) -> None:
         """
         Delete given loc(-s) from block in-place.
@@ -1851,6 +1855,19 @@ def fillna(
         new_values = values.fillna(value=value, limit=limit)
         return [self.make_block_same_class(values=new_values)]
 
+    def delete(self, loc) -> None:
+        """
+        Delete given loc(-s) from block in-place.
+        """
+        # This will be unnecessary if/when __array_function__ is implemented
+        self.values = self.values.delete(loc, axis=0)
+        self.mgr_locs = self._mgr_locs.delete(loc)
+        try:
+            self._cache.clear()
+        except AttributeError:
+            # _cache not yet initialized
+            pass
+
 
 class DatetimeLikeBlockMixin(NDArrayBackedExtensionBlock):
     """Mixin class for DatetimeBlock, DatetimeTZBlock, and TimedeltaBlock."""
@@ -2120,7 +2137,9 @@ def ensure_block_shape(values: ArrayLike, ndim: int = 1) -> ArrayLike:
             # TODO(EA2D): https://github.com/pandas-dev/pandas/issues/23023
             # block.shape is incorrect for "2D" ExtensionArrays
             # We can't, and don't need to, reshape.
-            values = np.asarray(values).reshape(1, -1)
+            values = cast(Union[np.ndarray, DatetimeArray, TimedeltaArray], values)
+            values = values.reshape(1, -1)
+
     return values
 
 

pandas/core/internals/concat.py

@@ -43,6 +43,7 @@
     DatetimeArray,
     ExtensionArray,
 )
+from pandas.core.construction import ensure_wrapped_if_datetimelike
 from pandas.core.internals.array_manager import ArrayManager
 from pandas.core.internals.blocks import (
     ensure_block_shape,
@@ -142,8 +143,9 @@ def concatenate_managers(
             else:
                 # TODO(EA2D): special-casing not needed with 2D EAs
                 values = concat_compat(vals)
-                if not isinstance(values, ExtensionArray):
-                    values = values.reshape(1, len(values))
+                values = ensure_block_shape(values, ndim=2)
+
+            values = ensure_wrapped_if_datetimelike(values)
 
             if blk.values.dtype == values.dtype:
                 # Fast-path

pandas/tests/internals/test_internals.py

@@ -30,6 +30,7 @@
 from pandas.core.arrays import (
     DatetimeArray,
     SparseArray,
+    TimedeltaArray,
 )
 from pandas.core.internals import (
     BlockManager,
@@ -300,6 +301,23 @@ def test_delete(self):
         with pytest.raises(IndexError, match=None):
             newb.delete(3)
 
+    def test_delete_datetimelike(self):
+        # dont use np.delete on values, as that will coerce from DTA/TDA to ndarray
+        arr = np.arange(20, dtype="i8").reshape(5, 4).view("m8[ns]")
+        df = DataFrame(arr)
+        blk = df._mgr.blocks[0]
+        assert isinstance(blk.values, TimedeltaArray)
+
+        blk.delete(1)
+        assert isinstance(blk.values, TimedeltaArray)
+
+        df = DataFrame(arr.view("M8[ns]"))
+        blk = df._mgr.blocks[0]
+        assert isinstance(blk.values, DatetimeArray)
+
+        blk.delete([1, 3])
+        assert isinstance(blk.values, DatetimeArray)
+
     def test_split(self):
         # GH#37799
         values = np.random.randn(3, 4)