Skip to content

Commit bb2e215

Browse files
khemkaran10Khemkaran
andauthored
BUG: IntervalIndex.unique() only contains the first interval if all interval borders are negative (#61920)
Co-authored-by: Khemkaran <[email protected]>
1 parent e982ba2 commit bb2e215

File tree

2 files changed

+39
-23
lines changed

2 files changed

+39
-23
lines changed

pandas/core/arrays/interval.py

Lines changed: 14 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -2100,18 +2100,9 @@ def isin(self, values: ArrayLike) -> npt.NDArray[np.bool_]:
21002100
return np.zeros(self.shape, dtype=bool)
21012101

21022102
if self.dtype == values.dtype:
2103-
# GH#38353 instead of casting to object, operating on a
2104-
# complex128 ndarray is much more performant.
2105-
left = self._combined.view("complex128")
2106-
right = values._combined.view("complex128")
2107-
# error: Argument 1 to "isin" has incompatible type
2108-
# "Union[ExtensionArray, ndarray[Any, Any],
2109-
# ndarray[Any, dtype[Any]]]"; expected
2110-
# "Union[_SupportsArray[dtype[Any]],
2111-
# _NestedSequence[_SupportsArray[dtype[Any]]], bool,
2112-
# int, float, complex, str, bytes, _NestedSequence[
2113-
# Union[bool, int, float, complex, str, bytes]]]"
2114-
return np.isin(left, right).ravel() # type: ignore[arg-type]
2103+
left = self._combined
2104+
right = values._combined
2105+
return np.isin(left, right).ravel()
21152106

21162107
elif needs_i8_conversion(self.left.dtype) ^ needs_i8_conversion(
21172108
values.left.dtype
@@ -2127,24 +2118,29 @@ def _combined(self) -> IntervalSide:
21272118
# has no attribute "reshape" [union-attr]
21282119
left = self.left._values.reshape(-1, 1) # type: ignore[union-attr]
21292120
right = self.right._values.reshape(-1, 1) # type: ignore[union-attr]
2121+
# GH#38353 instead of casting to object, operating on a
2122+
# complex128 ndarray is much more performant.
21302123
if needs_i8_conversion(left.dtype):
21312124
# error: Item "ndarray[Any, Any]" of "Any | ndarray[Any, Any]" has
21322125
# no attribute "_concat_same_type"
21332126
comb = left._concat_same_type( # type: ignore[union-attr]
21342127
[left, right], axis=1
21352128
)
2129+
comb = comb.view("complex128")[:, 0]
21362130
else:
2137-
comb = np.concatenate([left, right], axis=1)
2131+
comb = (np.array(left.ravel(), dtype="complex128")) + (
2132+
1j * np.array(right.ravel(), dtype="complex128")
2133+
)
21382134
return comb
21392135

21402136
def _from_combined(self, combined: np.ndarray) -> IntervalArray:
21412137
"""
21422138
Create a new IntervalArray with our dtype from a 1D complex128 ndarray.
21432139
"""
2144-
nc = combined.view("i8").reshape(-1, 2)
21452140

21462141
dtype = self._left.dtype
21472142
if needs_i8_conversion(dtype):
2143+
nc = combined.view("i8").reshape(-1, 2)
21482144
assert isinstance(self._left, (DatetimeArray, TimedeltaArray))
21492145
new_left: DatetimeArray | TimedeltaArray | np.ndarray = type(
21502146
self._left
@@ -2155,18 +2151,13 @@ def _from_combined(self, combined: np.ndarray) -> IntervalArray:
21552151
)._from_sequence(nc[:, 1], dtype=dtype)
21562152
else:
21572153
assert isinstance(dtype, np.dtype)
2158-
new_left = nc[:, 0].view(dtype)
2159-
new_right = nc[:, 1].view(dtype)
2154+
new_left = np.real(combined).astype(dtype).ravel()
2155+
new_right = np.imag(combined).astype(dtype).ravel()
21602156
return self._shallow_copy(left=new_left, right=new_right)
21612157

21622158
def unique(self) -> IntervalArray:
2163-
# No overload variant of "__getitem__" of "ExtensionArray" matches argument
2164-
# type "Tuple[slice, int]"
2165-
nc = unique(
2166-
self._combined.view("complex128")[:, 0] # type: ignore[call-overload]
2167-
)
2168-
nc = nc[:, None]
2169-
return self._from_combined(nc)
2159+
nc = unique(self._combined)
2160+
return self._from_combined(np.asarray(nc)[:, None])
21702161

21712162

21722163
def _maybe_convert_platform_interval(values) -> ArrayLike:

pandas/tests/arrays/interval/test_interval.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,31 @@ def test_shift_datetime(self):
111111
with pytest.raises(TypeError, match=msg):
112112
a.shift(1, fill_value=np.timedelta64("NaT", "ns"))
113113

114+
def test_unique_with_negatives(self):
115+
# GH#61917
116+
idx_pos = IntervalIndex.from_tuples(
117+
[(3, 4), (3, 4), (2, 3), (2, 3), (1, 2), (1, 2)]
118+
)
119+
result = idx_pos.unique()
120+
expected = IntervalIndex.from_tuples([(3, 4), (2, 3), (1, 2)])
121+
tm.assert_index_equal(result, expected)
122+
123+
idx_neg = IntervalIndex.from_tuples(
124+
[(-4, -3), (-4, -3), (-3, -2), (-3, -2), (-2, -1), (-2, -1)]
125+
)
126+
result = idx_neg.unique()
127+
expected = IntervalIndex.from_tuples([(-4, -3), (-3, -2), (-2, -1)])
128+
tm.assert_index_equal(result, expected)
129+
130+
idx_mix = IntervalIndex.from_tuples(
131+
[(1, 2), (0, 1), (-1, 0), (-2, -1), (-3, -2), (-3, -2)]
132+
)
133+
result = idx_mix.unique()
134+
expected = IntervalIndex.from_tuples(
135+
[(1, 2), (0, 1), (-1, 0), (-2, -1), (-3, -2)]
136+
)
137+
tm.assert_index_equal(result, expected)
138+
114139

115140
class TestSetitem:
116141
def test_set_na(self, left_right_dtypes):

0 commit comments

Comments
 (0)