317 lines
11 KiB
Python
317 lines
11 KiB
Python
import operator
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from pandas.core.dtypes.common import is_list_like
|
|
|
|
import pandas as pd
|
|
from pandas import (
|
|
Categorical,
|
|
Index,
|
|
Interval,
|
|
IntervalIndex,
|
|
Period,
|
|
Series,
|
|
Timedelta,
|
|
Timestamp,
|
|
date_range,
|
|
period_range,
|
|
timedelta_range,
|
|
)
|
|
import pandas._testing as tm
|
|
from pandas.core.arrays import (
|
|
BooleanArray,
|
|
IntervalArray,
|
|
)
|
|
from pandas.tests.arithmetic.common import get_upcast_box
|
|
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
(Index([0, 2, 4, 4]), Index([1, 3, 5, 8])),
|
|
(Index([0.0, 1.0, 2.0, np.nan]), Index([1.0, 2.0, 3.0, np.nan])),
|
|
(
|
|
timedelta_range("0 days", periods=3).insert(3, pd.NaT),
|
|
timedelta_range("1 day", periods=3).insert(3, pd.NaT),
|
|
),
|
|
(
|
|
date_range("20170101", periods=3).insert(3, pd.NaT),
|
|
date_range("20170102", periods=3).insert(3, pd.NaT),
|
|
),
|
|
(
|
|
date_range("20170101", periods=3, tz="US/Eastern").insert(3, pd.NaT),
|
|
date_range("20170102", periods=3, tz="US/Eastern").insert(3, pd.NaT),
|
|
),
|
|
],
|
|
ids=lambda x: str(x[0].dtype),
|
|
)
|
|
def left_right_dtypes(request):
|
|
"""
|
|
Fixture for building an IntervalArray from various dtypes
|
|
"""
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture
|
|
def interval_array(left_right_dtypes):
|
|
"""
|
|
Fixture to generate an IntervalArray of various dtypes containing NA if possible
|
|
"""
|
|
left, right = left_right_dtypes
|
|
return IntervalArray.from_arrays(left, right)
|
|
|
|
|
|
def create_categorical_intervals(left, right, closed="right"):
|
|
return Categorical(IntervalIndex.from_arrays(left, right, closed))
|
|
|
|
|
|
def create_series_intervals(left, right, closed="right"):
|
|
return Series(IntervalArray.from_arrays(left, right, closed))
|
|
|
|
|
|
def create_series_categorical_intervals(left, right, closed="right"):
|
|
return Series(Categorical(IntervalIndex.from_arrays(left, right, closed)))
|
|
|
|
|
|
class TestComparison:
|
|
@pytest.fixture(params=[operator.eq, operator.ne])
|
|
def op(self, request):
|
|
return request.param
|
|
|
|
@pytest.fixture(
|
|
params=[
|
|
IntervalArray.from_arrays,
|
|
IntervalIndex.from_arrays,
|
|
create_categorical_intervals,
|
|
create_series_intervals,
|
|
create_series_categorical_intervals,
|
|
],
|
|
ids=[
|
|
"IntervalArray",
|
|
"IntervalIndex",
|
|
"Categorical[Interval]",
|
|
"Series[Interval]",
|
|
"Series[Categorical[Interval]]",
|
|
],
|
|
)
|
|
def interval_constructor(self, request):
|
|
"""
|
|
Fixture for all pandas native interval constructors.
|
|
To be used as the LHS of IntervalArray comparisons.
|
|
"""
|
|
return request.param
|
|
|
|
def elementwise_comparison(self, op, interval_array, other):
|
|
"""
|
|
Helper that performs elementwise comparisons between `array` and `other`
|
|
"""
|
|
other = other if is_list_like(other) else [other] * len(interval_array)
|
|
expected = np.array([op(x, y) for x, y in zip(interval_array, other)])
|
|
if isinstance(other, Series):
|
|
return Series(expected, index=other.index)
|
|
return expected
|
|
|
|
def test_compare_scalar_interval(self, op, interval_array):
|
|
# matches first interval
|
|
other = interval_array[0]
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
# matches on a single endpoint but not both
|
|
other = Interval(interval_array.left[0], interval_array.right[1])
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
def test_compare_scalar_interval_mixed_closed(self, op, closed, other_closed):
|
|
interval_array = IntervalArray.from_arrays(range(2), range(1, 3), closed=closed)
|
|
other = Interval(0, 1, closed=other_closed)
|
|
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
def test_compare_scalar_na(
|
|
self, op, interval_array, nulls_fixture, box_with_array, request
|
|
):
|
|
box = box_with_array
|
|
|
|
if box is pd.DataFrame:
|
|
if interval_array.dtype.subtype.kind not in "iuf":
|
|
mark = pytest.mark.xfail(
|
|
reason="raises on DataFrame.transpose (would be fixed by EA2D)"
|
|
)
|
|
request.node.add_marker(mark)
|
|
|
|
obj = tm.box_expected(interval_array, box)
|
|
result = op(obj, nulls_fixture)
|
|
|
|
if nulls_fixture is pd.NA:
|
|
# GH#31882
|
|
exp = np.ones(interval_array.shape, dtype=bool)
|
|
expected = BooleanArray(exp, exp)
|
|
else:
|
|
expected = self.elementwise_comparison(op, interval_array, nulls_fixture)
|
|
|
|
if not (box is Index and nulls_fixture is pd.NA):
|
|
# don't cast expected from BooleanArray to ndarray[object]
|
|
xbox = get_upcast_box(obj, nulls_fixture, True)
|
|
expected = tm.box_expected(expected, xbox)
|
|
|
|
tm.assert_equal(result, expected)
|
|
|
|
rev = op(nulls_fixture, obj)
|
|
tm.assert_equal(rev, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"other",
|
|
[
|
|
0,
|
|
1.0,
|
|
True,
|
|
"foo",
|
|
Timestamp("2017-01-01"),
|
|
Timestamp("2017-01-01", tz="US/Eastern"),
|
|
Timedelta("0 days"),
|
|
Period("2017-01-01", "D"),
|
|
],
|
|
)
|
|
def test_compare_scalar_other(self, op, interval_array, other):
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
def test_compare_list_like_interval(self, op, interval_array, interval_constructor):
|
|
# same endpoints
|
|
other = interval_constructor(interval_array.left, interval_array.right)
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_equal(result, expected)
|
|
|
|
# different endpoints
|
|
other = interval_constructor(
|
|
interval_array.left[::-1], interval_array.right[::-1]
|
|
)
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_equal(result, expected)
|
|
|
|
# all nan endpoints
|
|
other = interval_constructor([np.nan] * 4, [np.nan] * 4)
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_equal(result, expected)
|
|
|
|
def test_compare_list_like_interval_mixed_closed(
|
|
self, op, interval_constructor, closed, other_closed
|
|
):
|
|
interval_array = IntervalArray.from_arrays(range(2), range(1, 3), closed=closed)
|
|
other = interval_constructor(range(2), range(1, 3), closed=other_closed)
|
|
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_equal(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"other",
|
|
[
|
|
(
|
|
Interval(0, 1),
|
|
Interval(Timedelta("1 day"), Timedelta("2 days")),
|
|
Interval(4, 5, "both"),
|
|
Interval(10, 20, "neither"),
|
|
),
|
|
(0, 1.5, Timestamp("20170103"), np.nan),
|
|
(
|
|
Timestamp("20170102", tz="US/Eastern"),
|
|
Timedelta("2 days"),
|
|
"baz",
|
|
pd.NaT,
|
|
),
|
|
],
|
|
)
|
|
def test_compare_list_like_object(self, op, interval_array, other):
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
def test_compare_list_like_nan(self, op, interval_array, nulls_fixture):
|
|
other = [nulls_fixture] * 4
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
|
|
tm.assert_equal(result, expected)
|
|
|
|
@pytest.mark.parametrize(
|
|
"other",
|
|
[
|
|
np.arange(4, dtype="int64"),
|
|
np.arange(4, dtype="float64"),
|
|
date_range("2017-01-01", periods=4),
|
|
date_range("2017-01-01", periods=4, tz="US/Eastern"),
|
|
timedelta_range("0 days", periods=4),
|
|
period_range("2017-01-01", periods=4, freq="D"),
|
|
Categorical(list("abab")),
|
|
Categorical(date_range("2017-01-01", periods=4)),
|
|
pd.array(list("abcd")),
|
|
pd.array(["foo", 3.14, None, object()], dtype=object),
|
|
],
|
|
ids=lambda x: str(x.dtype),
|
|
)
|
|
def test_compare_list_like_other(self, op, interval_array, other):
|
|
result = op(interval_array, other)
|
|
expected = self.elementwise_comparison(op, interval_array, other)
|
|
tm.assert_numpy_array_equal(result, expected)
|
|
|
|
@pytest.mark.parametrize("length", [1, 3, 5])
|
|
@pytest.mark.parametrize("other_constructor", [IntervalArray, list])
|
|
def test_compare_length_mismatch_errors(self, op, other_constructor, length):
|
|
interval_array = IntervalArray.from_arrays(range(4), range(1, 5))
|
|
other = other_constructor([Interval(0, 1)] * length)
|
|
with pytest.raises(ValueError, match="Lengths must match to compare"):
|
|
op(interval_array, other)
|
|
|
|
@pytest.mark.parametrize(
|
|
"constructor, expected_type, assert_func",
|
|
[
|
|
(IntervalIndex, np.array, tm.assert_numpy_array_equal),
|
|
(Series, Series, tm.assert_series_equal),
|
|
],
|
|
)
|
|
def test_index_series_compat(self, op, constructor, expected_type, assert_func):
|
|
# IntervalIndex/Series that rely on IntervalArray for comparisons
|
|
breaks = range(4)
|
|
index = constructor(IntervalIndex.from_breaks(breaks))
|
|
|
|
# scalar comparisons
|
|
other = index[0]
|
|
result = op(index, other)
|
|
expected = expected_type(self.elementwise_comparison(op, index, other))
|
|
assert_func(result, expected)
|
|
|
|
other = breaks[0]
|
|
result = op(index, other)
|
|
expected = expected_type(self.elementwise_comparison(op, index, other))
|
|
assert_func(result, expected)
|
|
|
|
# list-like comparisons
|
|
other = IntervalArray.from_breaks(breaks)
|
|
result = op(index, other)
|
|
expected = expected_type(self.elementwise_comparison(op, index, other))
|
|
assert_func(result, expected)
|
|
|
|
other = [index[0], breaks[0], "foo"]
|
|
result = op(index, other)
|
|
expected = expected_type(self.elementwise_comparison(op, index, other))
|
|
assert_func(result, expected)
|
|
|
|
@pytest.mark.parametrize("scalars", ["a", False, 1, 1.0, None])
|
|
def test_comparison_operations(self, scalars):
|
|
# GH #28981
|
|
expected = Series([False, False])
|
|
s = Series([Interval(0, 1), Interval(1, 2)], dtype="interval")
|
|
result = s == scalars
|
|
tm.assert_series_equal(result, expected)
|