REF: implement is_exact_match (#39149)

Author: jbrockmendel

pandas/core/indexers.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from pandas._typing import Any, AnyArrayLike
+from pandas._typing import Any, AnyArrayLike, ArrayLike
 
 from pandas.core.dtypes.common import (
     is_array_like,
@@ -270,6 +270,27 @@ def maybe_convert_indices(indices, n: int):
 # Unsorted
 
 
+def is_exact_shape_match(target: ArrayLike, value: ArrayLike) -> bool:
+    """
+    Is setting this value into this target overwriting the entire column?
+
+    Parameters
+    ----------
+    target : np.ndarray or ExtensionArray
+    value : np.ndarray or ExtensionArray
+
+    Returns
+    -------
+    bool
+    """
+    return (
+        len(value.shape) > 0
+        and len(target.shape) > 0
+        and value.shape[0] == target.shape[0]
+        and value.size == target.size
+    )
+
+
 def length_of_indexer(indexer, target=None) -> int:
     """
     Return the expected length of target[indexer]

pandas/core/indexing.py

@@ -33,6 +33,7 @@
 from pandas.core.construction import array as pd_array
 from pandas.core.indexers import (
     check_array_indexer,
+    is_exact_shape_match,
     is_list_like_indexer,
     length_of_indexer,
 )
@@ -1815,6 +1816,9 @@ def _setitem_single_column(self, loc: int, value, plane_indexer):
         # GH#6149 (null slice), GH#10408 (full bounds)
         if com.is_null_slice(pi) or com.is_full_slice(pi, len(self.obj)):
             ser = value
+        elif is_array_like(value) and is_exact_shape_match(ser, value):
+            ser = value
+
         else:
             # set the item, possibly having a dtype change
             ser = ser.copy()

pandas/core/internals/blocks.py

@@ -75,6 +75,7 @@
 from pandas.core.indexers import (
     check_setitem_lengths,
     is_empty_indexer,
+    is_exact_shape_match,
     is_scalar_indexer,
 )
 import pandas.core.missing as missing
@@ -903,15 +904,7 @@ def setitem(self, indexer, value):
 
         # coerce if block dtype can store value
         values = self.values
-        if self._can_hold_element(value):
-            # We only get here for non-Extension Blocks, so _try_coerce_args
-            #  is only relevant for DatetimeBlock and TimedeltaBlock
-            if self.dtype.kind in ["m", "M"]:
-                arr = self.array_values().T
-                arr[indexer] = value
-                return self
-
-        else:
+        if not self._can_hold_element(value):
             # current dtype cannot store value, coerce to common dtype
             # TODO: can we just use coerce_to_target_dtype for all this
             if hasattr(value, "dtype"):
@@ -932,6 +925,11 @@ def setitem(self, indexer, value):
 
             return self.astype(dtype).setitem(indexer, value)
 
+        if self.dtype.kind in ["m", "M"]:
+            arr = self.array_values().T
+            arr[indexer] = value
+            return self
+
         # value must be storable at this moment
         if is_extension_array_dtype(getattr(value, "dtype", None)):
             # We need to be careful not to allow through strings that
@@ -947,11 +945,7 @@ def setitem(self, indexer, value):
 
         # length checking
         check_setitem_lengths(indexer, value, values)
-        exact_match = (
-            len(arr_value.shape)
-            and arr_value.shape[0] == values.shape[0]
-            and arr_value.size == values.size
-        )
+        exact_match = is_exact_shape_match(values, arr_value)
         if is_empty_indexer(indexer, arr_value):
             # GH#8669 empty indexers
             pass