-
-
Notifications
You must be signed in to change notification settings - Fork 18.4k
REF: IntervalArray comparisons #37124
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 3 commits
a3b7c45
60197d2
cf20846
d3c1041
fcfe47d
3dfc008
fa6cecd
0d32acd
ff640ea
f501b4e
247ce90
84b0409
1ea125d
6e43726
badb99d
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,4 @@ | ||
import operator | ||
from operator import le, lt | ||
import textwrap | ||
from typing import TYPE_CHECKING, Optional, Tuple, Union, cast | ||
|
@@ -48,7 +49,7 @@ | |
from pandas.core.construction import array, extract_array | ||
from pandas.core.indexers import check_array_indexer | ||
from pandas.core.indexes.base import ensure_index | ||
from pandas.core.ops import unpack_zerodim_and_defer | ||
from pandas.core.ops import invalid_comparison, unpack_zerodim_and_defer | ||
|
||
if TYPE_CHECKING: | ||
from pandas import Index | ||
|
@@ -520,16 +521,15 @@ def __setitem__(self, key, value): | |
self._left[key] = value_left | ||
self._right[key] = value_right | ||
|
||
@unpack_zerodim_and_defer("__eq__") | ||
def __eq__(self, other): | ||
def _cmp_method(self, other, op): | ||
# ensure pandas array for list-like and eliminate non-interval scalars | ||
if is_list_like(other): | ||
if len(self) != len(other): | ||
raise ValueError("Lengths must match to compare") | ||
other = array(other) | ||
elif not isinstance(other, Interval): | ||
# non-interval scalar -> no matches | ||
return np.zeros(len(self), dtype=bool) | ||
return invalid_comparison(self, other, op) | ||
|
||
# determine the dtype of the elements we want to compare | ||
if isinstance(other, Interval): | ||
|
@@ -543,33 +543,87 @@ def __eq__(self, other): | |
# extract intervals if we have interval categories with matching closed | ||
if is_interval_dtype(other_dtype): | ||
if self.closed != other.categories.closed: | ||
return np.zeros(len(self), dtype=bool) | ||
return invalid_comparison(self, other, op) | ||
other = other.categories.take(other.codes) | ||
|
||
# interval-like -> need same closed and matching endpoints | ||
if is_interval_dtype(other_dtype): | ||
if self.closed != other.closed: | ||
return np.zeros(len(self), dtype=bool) | ||
return (self._left == other.left) & (self._right == other.right) | ||
return invalid_comparison(self, other, op) | ||
if isinstance(other, Interval): | ||
other = type(self)._from_sequence([other]) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can't we use |
||
if self._combined.dtype.kind in ["m", "M"]: | ||
# Need to repeat bc we do not broadcast length-1 | ||
# TODO: would be helpful to have a tile method to do | ||
# this without copies | ||
other = other.repeat(len(self)) | ||
else: | ||
other = type(self)(other) | ||
|
||
if op is operator.eq: | ||
jreback marked this conversation as resolved.
Show resolved
Hide resolved
|
||
return (self._combined[:, 0] == other._left) & ( | ||
self._combined[:, 1] == other._right | ||
) | ||
elif op is operator.ne: | ||
return (self._combined[:, 0] != other._left) | ( | ||
self._combined[:, 1] != other._right | ||
) | ||
elif op is operator.gt: | ||
return (self._combined[:, 0] > other._combined[:, 0]) | ( | ||
(self._combined[:, 0] == other._left) | ||
& (self._combined[:, 1] > other._right) | ||
) | ||
elif op is operator.ge: | ||
return (self == other) | (self > other) | ||
elif op is operator.lt: | ||
return (self._combined[:, 0] < other._combined[:, 0]) | ( | ||
(self._combined[:, 0] == other._left) | ||
& (self._combined[:, 1] < other._right) | ||
) | ||
else: | ||
# operator.lt | ||
return (self == other) | (self < other) | ||
|
||
# non-interval/non-object dtype -> no matches | ||
if not is_object_dtype(other_dtype): | ||
return np.zeros(len(self), dtype=bool) | ||
return invalid_comparison(self, other, op) | ||
|
||
# object dtype -> iteratively check for intervals | ||
result = np.zeros(len(self), dtype=bool) | ||
for i, obj in enumerate(other): | ||
# need object to be an Interval with same closed and endpoints | ||
if ( | ||
isinstance(obj, Interval) | ||
and self.closed == obj.closed | ||
and self._left[i] == obj.left | ||
and self._right[i] == obj.right | ||
): | ||
result[i] = True | ||
|
||
try: | ||
result = np.zeros(len(self), dtype=bool) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. you should do the np.zeros outside of the try/except There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. sure |
||
for i, obj in enumerate(other): | ||
result[i] = op(self[i], obj) | ||
except TypeError: | ||
# pd.NA | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you have an example (or test) that runs into this? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Ah, I see the removed special case below in the tests. Now, this is certainly in a messy state, but I am not sure the conversion to object dtype is needed here. Currently, we still return False (and not NA as for nullable dtypes) on comparisons with NA:
|
||
result = np.zeros(len(self), dtype=object) | ||
for i, obj in enumerate(other): | ||
result[i] = op(self[i], obj) | ||
return result | ||
|
||
@unpack_zerodim_and_defer("__eq__") | ||
def __eq__(self, other): | ||
return self._cmp_method(other, operator.eq) | ||
|
||
@unpack_zerodim_and_defer("__ne__") | ||
def __ne__(self, other): | ||
return self._cmp_method(other, operator.ne) | ||
|
||
@unpack_zerodim_and_defer("__gt__") | ||
def __gt__(self, other): | ||
return self._cmp_method(other, operator.gt) | ||
|
||
@unpack_zerodim_and_defer("__ge__") | ||
def __ge__(self, other): | ||
return self._cmp_method(other, operator.ge) | ||
|
||
@unpack_zerodim_and_defer("__lt__") | ||
def __lt__(self, other): | ||
return self._cmp_method(other, operator.lt) | ||
|
||
@unpack_zerodim_and_defer("__le__") | ||
def __le__(self, other): | ||
return self._cmp_method(other, operator.le) | ||
|
||
def fillna(self, value=None, method=None, limit=None): | ||
""" | ||
Fill NA/NaN values using the specified method. | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why is this no longer needed? or unrelated?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can you comment here @jbrockmendel
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In IntervalArray.repeat we call
self._combined.repeat(repeats, 0)
, which when _combined is DTA/TDA gets here withaxis=0
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Since there is still discussion about whether
_combined
is meant to stay (#37047), can you leave out this clean-up from this PR?(unless the clean-up is not tied to having
_combined
, but from your comment it seems so)