import numpy as np import pytest import pandas.util._test_decorators as td import pandas as pd from pandas import ( Index, Interval, IntervalIndex, Timedelta, Timestamp, date_range, timedelta_range, ) import pandas._testing as tm from pandas.core.arrays import IntervalArray @pytest.fixture( params=[ (Index([0, 2, 4]), Index([1, 3, 5])), (Index([0.0, 1.0, 2.0]), Index([1.0, 2.0, 3.0])), (timedelta_range("0 days", periods=3), timedelta_range("1 day", periods=3)), (date_range("20170101", periods=3), date_range("20170102", periods=3)), ( date_range("20170101", periods=3, tz="US/Eastern"), date_range("20170102", periods=3, tz="US/Eastern"), ), ], ids=lambda x: str(x[0].dtype), ) def left_right_dtypes(request): """ Fixture for building an IntervalArray from various dtypes """ return request.param class TestAttributes: @pytest.mark.parametrize( "left, right", [ (0, 1), (Timedelta("0 days"), Timedelta("1 day")), (Timestamp("2018-01-01"), Timestamp("2018-01-02")), ( Timestamp("2018-01-01", tz="US/Eastern"), Timestamp("2018-01-02", tz="US/Eastern"), ), ], ) @pytest.mark.parametrize("constructor", [IntervalArray, IntervalIndex]) def test_is_empty(self, constructor, left, right, closed): # GH27219 tuples = [(left, left), (left, right), np.nan] expected = np.array([closed != "both", False, False]) result = constructor.from_tuples(tuples, closed=closed).is_empty tm.assert_numpy_array_equal(result, expected) class TestMethods: @pytest.mark.parametrize("new_closed", ["left", "right", "both", "neither"]) def test_set_closed(self, closed, new_closed): # GH 21670 array = IntervalArray.from_breaks(range(10), closed=closed) result = array.set_closed(new_closed) expected = IntervalArray.from_breaks(range(10), closed=new_closed) tm.assert_extension_array_equal(result, expected) @pytest.mark.parametrize( "other", [ Interval(0, 1, closed="right"), IntervalArray.from_breaks([1, 2, 3, 4], closed="right"), ], ) def test_where_raises(self, other): ser = pd.Series(IntervalArray.from_breaks([1, 2, 3, 4], closed="left")) match = "'value.closed' is 'right', expected 'left'." with pytest.raises(ValueError, match=match): ser.where([True, False, True], other=other) def test_shift(self): # https://github.com/pandas-dev/pandas/issues/31495 a = IntervalArray.from_breaks([1, 2, 3]) result = a.shift() # int -> float expected = IntervalArray.from_tuples([(np.nan, np.nan), (1.0, 2.0)]) tm.assert_interval_array_equal(result, expected) def test_shift_datetime(self): a = IntervalArray.from_breaks(pd.date_range("2000", periods=4)) result = a.shift(2) expected = a.take([-1, -1, 0], allow_fill=True) tm.assert_interval_array_equal(result, expected) result = a.shift(-1) expected = a.take([1, 2, -1], allow_fill=True) tm.assert_interval_array_equal(result, expected) class TestSetitem: def test_set_na(self, left_right_dtypes): left, right = left_right_dtypes result = IntervalArray.from_arrays(left, right) if result.dtype.subtype.kind not in ["m", "M"]: msg = "'value' should be an interval type, got <.*NaTType'> instead." with pytest.raises(TypeError, match=msg): result[0] = pd.NaT if result.dtype.subtype.kind in ["i", "u"]: msg = "Cannot set float NaN to integer-backed IntervalArray" with pytest.raises(ValueError, match=msg): result[0] = np.NaN return result[0] = np.nan expected_left = Index([left._na_value] + list(left[1:])) expected_right = Index([right._na_value] + list(right[1:])) expected = IntervalArray.from_arrays(expected_left, expected_right) tm.assert_extension_array_equal(result, expected) def test_repr(): # GH 25022 arr = IntervalArray.from_tuples([(0, 1), (1, 2)]) result = repr(arr) expected = ( "\n" "[(0, 1], (1, 2]]\n" "Length: 2, closed: right, dtype: interval[int64]" ) assert result == expected # ---------------------------------------------------------------------------- # Arrow interaction pyarrow_skip = td.skip_if_no("pyarrow", min_version="0.15.1.dev") @pyarrow_skip def test_arrow_extension_type(): import pyarrow as pa from pandas.core.arrays._arrow_utils import ArrowIntervalType p1 = ArrowIntervalType(pa.int64(), "left") p2 = ArrowIntervalType(pa.int64(), "left") p3 = ArrowIntervalType(pa.int64(), "right") assert p1.closed == "left" assert p1 == p2 assert not p1 == p3 assert hash(p1) == hash(p2) assert not hash(p1) == hash(p3) @pyarrow_skip def test_arrow_array(): import pyarrow as pa from pandas.core.arrays._arrow_utils import ArrowIntervalType intervals = pd.interval_range(1, 5, freq=1).array result = pa.array(intervals) assert isinstance(result.type, ArrowIntervalType) assert result.type.closed == intervals.closed assert result.type.subtype == pa.int64() assert result.storage.field("left").equals(pa.array([1, 2, 3, 4], type="int64")) assert result.storage.field("right").equals(pa.array([2, 3, 4, 5], type="int64")) expected = pa.array([{"left": i, "right": i + 1} for i in range(1, 5)]) assert result.storage.equals(expected) # convert to its storage type result = pa.array(intervals, type=expected.type) assert result.equals(expected) # unsupported conversions with pytest.raises(TypeError, match="Not supported to convert IntervalArray"): pa.array(intervals, type="float64") with pytest.raises(TypeError, match="different 'subtype'"): pa.array(intervals, type=ArrowIntervalType(pa.float64(), "left")) @pyarrow_skip def test_arrow_array_missing(): import pyarrow as pa from pandas.core.arrays._arrow_utils import ArrowIntervalType arr = IntervalArray.from_breaks([0.0, 1.0, 2.0, 3.0]) arr[1] = None result = pa.array(arr) assert isinstance(result.type, ArrowIntervalType) assert result.type.closed == arr.closed assert result.type.subtype == pa.float64() # fields have missing values (not NaN) left = pa.array([0.0, None, 2.0], type="float64") right = pa.array([1.0, None, 3.0], type="float64") assert result.storage.field("left").equals(left) assert result.storage.field("right").equals(right) # structarray itself also has missing values on the array level vals = [ {"left": 0.0, "right": 1.0}, {"left": None, "right": None}, {"left": 2.0, "right": 3.0}, ] expected = pa.StructArray.from_pandas(vals, mask=np.array([False, True, False])) assert result.storage.equals(expected) @pyarrow_skip @pytest.mark.parametrize( "breaks", [[0.0, 1.0, 2.0, 3.0], pd.date_range("2017", periods=4, freq="D")], ids=["float", "datetime64[ns]"], ) def test_arrow_table_roundtrip(breaks): import pyarrow as pa from pandas.core.arrays._arrow_utils import ArrowIntervalType arr = IntervalArray.from_breaks(breaks) arr[1] = None df = pd.DataFrame({"a": arr}) table = pa.table(df) assert isinstance(table.field("a").type, ArrowIntervalType) result = table.to_pandas() assert isinstance(result["a"].dtype, pd.IntervalDtype) tm.assert_frame_equal(result, df) table2 = pa.concat_tables([table, table]) result = table2.to_pandas() expected = pd.concat([df, df], ignore_index=True) tm.assert_frame_equal(result, expected) @pyarrow_skip @pytest.mark.parametrize( "breaks", [[0.0, 1.0, 2.0, 3.0], pd.date_range("2017", periods=4, freq="D")], ids=["float", "datetime64[ns]"], ) def test_arrow_table_roundtrip_without_metadata(breaks): import pyarrow as pa arr = IntervalArray.from_breaks(breaks) arr[1] = None df = pd.DataFrame({"a": arr}) table = pa.table(df) # remove the metadata table = table.replace_schema_metadata() assert table.schema.metadata is None result = table.to_pandas() assert isinstance(result["a"].dtype, pd.IntervalDtype) tm.assert_frame_equal(result, df)