Skip to content

Commit 7f910b5

Browse files
authored
CLN: ensure we pass correct type to DTI/TDI shallow_copy (#37171)
1 parent 0a5a3a7 commit 7f910b5

File tree

4 files changed

+17
-8
lines changed

4 files changed

+17
-8
lines changed

pandas/core/algorithms.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -696,6 +696,8 @@ def factorize(
696696

697697
# return original tenor
698698
if isinstance(original, ABCIndexClass):
699+
if original.dtype.kind in ["m", "M"] and isinstance(uniques, np.ndarray):
700+
uniques = type(original._data)._simple_new(uniques, dtype=original.dtype)
699701
uniques = original._shallow_copy(uniques, name=None)
700702
elif isinstance(original, ABCSeries):
701703
from pandas import Index
@@ -1650,7 +1652,8 @@ def take_nd(
16501652
"""
16511653
mask_info = None
16521654

1653-
if is_extension_array_dtype(arr):
1655+
if isinstance(arr, ABCExtensionArray):
1656+
# Check for EA to catch DatetimeArray, TimedeltaArray
16541657
return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill)
16551658

16561659
arr = extract_array(arr)
@@ -2043,7 +2046,7 @@ def safe_sort(
20432046
"Only list-like objects are allowed to be passed to safe_sort as values"
20442047
)
20452048

2046-
if not isinstance(values, np.ndarray) and not is_extension_array_dtype(values):
2049+
if not isinstance(values, (np.ndarray, ABCExtensionArray)):
20472050
# don't convert to string types
20482051
dtype, _ = infer_dtype_from_array(values)
20492052
values = np.asarray(values, dtype=dtype)

pandas/core/indexes/base.py

+8-3
Original file line numberDiff line numberDiff line change
@@ -2664,6 +2664,11 @@ def _union(self, other, sort):
26642664
return self._shallow_copy(result)
26652665

26662666
def _wrap_setop_result(self, other, result):
2667+
if isinstance(self, (ABCDatetimeIndex, ABCTimedeltaIndex)) and isinstance(
2668+
result, np.ndarray
2669+
):
2670+
result = type(self._data)._simple_new(result, dtype=self.dtype)
2671+
26672672
name = get_op_result_name(self, other)
26682673
if isinstance(result, Index):
26692674
if result.name != name:
@@ -2740,10 +2745,10 @@ def intersection(self, other, sort=False):
27402745
indexer = algos.unique1d(Index(rvals).get_indexer_non_unique(lvals)[0])
27412746
indexer = indexer[indexer != -1]
27422747

2743-
result = other.take(indexer)
2748+
result = other.take(indexer)._values
27442749

27452750
if sort is None:
2746-
result = algos.safe_sort(result.values)
2751+
result = algos.safe_sort(result)
27472752

27482753
return self._wrap_setop_result(other, result)
27492754

@@ -2800,7 +2805,7 @@ def difference(self, other, sort=None):
28002805
indexer = indexer.take((indexer != -1).nonzero()[0])
28012806

28022807
label_diff = np.setdiff1d(np.arange(this.size), indexer, assume_unique=True)
2803-
the_diff = this.values.take(label_diff)
2808+
the_diff = this._values.take(label_diff)
28042809
if sort is None:
28052810
try:
28062811
the_diff = algos.safe_sort(the_diff)

pandas/core/indexes/datetimelike.py

-3
Original file line numberDiff line numberDiff line change
@@ -696,9 +696,6 @@ def _shallow_copy(self, values=None, name: Label = lib.no_default):
696696
name = self.name if name is lib.no_default else name
697697

698698
if values is not None:
699-
# TODO: We would rather not get here
700-
if isinstance(values, np.ndarray):
701-
values = type(self._data)(values, dtype=self.dtype)
702699
return self._simple_new(values, name=name)
703700

704701
result = self._simple_new(self._data, name=name)

pandas/tests/indexes/test_common.py

+4
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,10 @@ def test_get_unique_index(self, index):
320320
vals[0] = np.nan
321321

322322
vals_unique = vals[:2]
323+
if index.dtype.kind in ["m", "M"]:
324+
# i.e. needs_i8_conversion but not period_dtype, as above
325+
vals = type(index._data)._simple_new(vals, dtype=index.dtype)
326+
vals_unique = type(index._data)._simple_new(vals_unique, dtype=index.dtype)
323327
idx_nan = index._shallow_copy(vals)
324328
idx_unique_nan = index._shallow_copy(vals_unique)
325329
assert idx_unique_nan.is_unique is True

0 commit comments

Comments
 (0)