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