You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
316 lines
11 KiB
316 lines
11 KiB
1 year ago
|
import numpy as np
|
||
|
import pytest
|
||
|
|
||
|
from pandas import (
|
||
|
DatetimeIndex,
|
||
|
Index,
|
||
|
NaT,
|
||
|
PeriodIndex,
|
||
|
TimedeltaIndex,
|
||
|
timedelta_range,
|
||
|
)
|
||
|
import pandas._testing as tm
|
||
|
|
||
|
|
||
|
def check_freq_ascending(ordered, orig, ascending):
|
||
|
"""
|
||
|
Check the expected freq on a PeriodIndex/DatetimeIndex/TimedeltaIndex
|
||
|
when the original index is generated (or generate-able) with
|
||
|
period_range/date_range/timedelta_range.
|
||
|
"""
|
||
|
if isinstance(ordered, PeriodIndex):
|
||
|
assert ordered.freq == orig.freq
|
||
|
elif isinstance(ordered, (DatetimeIndex, TimedeltaIndex)):
|
||
|
if ascending:
|
||
|
assert ordered.freq.n == orig.freq.n
|
||
|
else:
|
||
|
assert ordered.freq.n == -1 * orig.freq.n
|
||
|
|
||
|
|
||
|
def check_freq_nonmonotonic(ordered, orig):
|
||
|
"""
|
||
|
Check the expected freq on a PeriodIndex/DatetimeIndex/TimedeltaIndex
|
||
|
when the original index is _not_ generated (or generate-able) with
|
||
|
period_range/date_range//timedelta_range.
|
||
|
"""
|
||
|
if isinstance(ordered, PeriodIndex):
|
||
|
assert ordered.freq == orig.freq
|
||
|
elif isinstance(ordered, (DatetimeIndex, TimedeltaIndex)):
|
||
|
assert ordered.freq is None
|
||
|
|
||
|
|
||
|
class TestSortValues:
|
||
|
@pytest.fixture(params=[DatetimeIndex, TimedeltaIndex, PeriodIndex])
|
||
|
def non_monotonic_idx(self, request):
|
||
|
if request.param is DatetimeIndex:
|
||
|
return DatetimeIndex(["2000-01-04", "2000-01-01", "2000-01-02"])
|
||
|
elif request.param is PeriodIndex:
|
||
|
dti = DatetimeIndex(["2000-01-04", "2000-01-01", "2000-01-02"])
|
||
|
return dti.to_period("D")
|
||
|
else:
|
||
|
return TimedeltaIndex(
|
||
|
["1 day 00:00:05", "1 day 00:00:01", "1 day 00:00:02"]
|
||
|
)
|
||
|
|
||
|
def test_argmin_argmax(self, non_monotonic_idx):
|
||
|
assert non_monotonic_idx.argmin() == 1
|
||
|
assert non_monotonic_idx.argmax() == 0
|
||
|
|
||
|
def test_sort_values(self, non_monotonic_idx):
|
||
|
idx = non_monotonic_idx
|
||
|
ordered = idx.sort_values()
|
||
|
assert ordered.is_monotonic_increasing
|
||
|
ordered = idx.sort_values(ascending=False)
|
||
|
assert ordered[::-1].is_monotonic_increasing
|
||
|
|
||
|
ordered, dexer = idx.sort_values(return_indexer=True)
|
||
|
assert ordered.is_monotonic_increasing
|
||
|
tm.assert_numpy_array_equal(dexer, np.array([1, 2, 0], dtype=np.intp))
|
||
|
|
||
|
ordered, dexer = idx.sort_values(return_indexer=True, ascending=False)
|
||
|
assert ordered[::-1].is_monotonic_increasing
|
||
|
tm.assert_numpy_array_equal(dexer, np.array([0, 2, 1], dtype=np.intp))
|
||
|
|
||
|
def check_sort_values_with_freq(self, idx):
|
||
|
ordered = idx.sort_values()
|
||
|
tm.assert_index_equal(ordered, idx)
|
||
|
check_freq_ascending(ordered, idx, True)
|
||
|
|
||
|
ordered = idx.sort_values(ascending=False)
|
||
|
expected = idx[::-1]
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
check_freq_ascending(ordered, idx, False)
|
||
|
|
||
|
ordered, indexer = idx.sort_values(return_indexer=True)
|
||
|
tm.assert_index_equal(ordered, idx)
|
||
|
tm.assert_numpy_array_equal(indexer, np.array([0, 1, 2], dtype=np.intp))
|
||
|
check_freq_ascending(ordered, idx, True)
|
||
|
|
||
|
ordered, indexer = idx.sort_values(return_indexer=True, ascending=False)
|
||
|
expected = idx[::-1]
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
tm.assert_numpy_array_equal(indexer, np.array([2, 1, 0], dtype=np.intp))
|
||
|
check_freq_ascending(ordered, idx, False)
|
||
|
|
||
|
@pytest.mark.parametrize("freq", ["D", "H"])
|
||
|
def test_sort_values_with_freq_timedeltaindex(self, freq):
|
||
|
# GH#10295
|
||
|
idx = timedelta_range(start=f"1{freq}", periods=3, freq=freq).rename("idx")
|
||
|
|
||
|
self.check_sort_values_with_freq(idx)
|
||
|
|
||
|
@pytest.mark.parametrize(
|
||
|
"idx",
|
||
|
[
|
||
|
DatetimeIndex(
|
||
|
["2011-01-01", "2011-01-02", "2011-01-03"], freq="D", name="idx"
|
||
|
),
|
||
|
DatetimeIndex(
|
||
|
["2011-01-01 09:00", "2011-01-01 10:00", "2011-01-01 11:00"],
|
||
|
freq="H",
|
||
|
name="tzidx",
|
||
|
tz="Asia/Tokyo",
|
||
|
),
|
||
|
],
|
||
|
)
|
||
|
def test_sort_values_with_freq_datetimeindex(self, idx):
|
||
|
self.check_sort_values_with_freq(idx)
|
||
|
|
||
|
@pytest.mark.parametrize("freq", ["D", "2D", "4D"])
|
||
|
def test_sort_values_with_freq_periodindex(self, freq):
|
||
|
# here with_freq refers to being period_range-like
|
||
|
idx = PeriodIndex(
|
||
|
["2011-01-01", "2011-01-02", "2011-01-03"], freq=freq, name="idx"
|
||
|
)
|
||
|
self.check_sort_values_with_freq(idx)
|
||
|
|
||
|
@pytest.mark.parametrize(
|
||
|
"idx",
|
||
|
[
|
||
|
PeriodIndex(["2011", "2012", "2013"], name="pidx", freq="A"),
|
||
|
Index([2011, 2012, 2013], name="idx"), # for compatibility check
|
||
|
],
|
||
|
)
|
||
|
def test_sort_values_with_freq_periodindex2(self, idx):
|
||
|
# here with_freq indicates this is period_range-like
|
||
|
self.check_sort_values_with_freq(idx)
|
||
|
|
||
|
def check_sort_values_without_freq(self, idx, expected):
|
||
|
ordered = idx.sort_values(na_position="first")
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
if not idx.isna().any():
|
||
|
ordered = idx.sort_values()
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
ordered = idx.sort_values(ascending=False)
|
||
|
tm.assert_index_equal(ordered, expected[::-1])
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
ordered, indexer = idx.sort_values(return_indexer=True, na_position="first")
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
|
||
|
exp = np.array([0, 4, 3, 1, 2], dtype=np.intp)
|
||
|
tm.assert_numpy_array_equal(indexer, exp)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
if not idx.isna().any():
|
||
|
ordered, indexer = idx.sort_values(return_indexer=True)
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
|
||
|
exp = np.array([0, 4, 3, 1, 2], dtype=np.intp)
|
||
|
tm.assert_numpy_array_equal(indexer, exp)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
ordered, indexer = idx.sort_values(return_indexer=True, ascending=False)
|
||
|
tm.assert_index_equal(ordered, expected[::-1])
|
||
|
|
||
|
exp = np.array([2, 1, 3, 0, 4], dtype=np.intp)
|
||
|
tm.assert_numpy_array_equal(indexer, exp)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
def test_sort_values_without_freq_timedeltaindex(self):
|
||
|
# GH#10295
|
||
|
|
||
|
idx = TimedeltaIndex(
|
||
|
["1 hour", "3 hour", "5 hour", "2 hour ", "1 hour"], name="idx1"
|
||
|
)
|
||
|
expected = TimedeltaIndex(
|
||
|
["1 hour", "1 hour", "2 hour", "3 hour", "5 hour"], name="idx1"
|
||
|
)
|
||
|
self.check_sort_values_without_freq(idx, expected)
|
||
|
|
||
|
@pytest.mark.parametrize(
|
||
|
"index_dates,expected_dates",
|
||
|
[
|
||
|
(
|
||
|
["2011-01-01", "2011-01-03", "2011-01-05", "2011-01-02", "2011-01-01"],
|
||
|
["2011-01-01", "2011-01-01", "2011-01-02", "2011-01-03", "2011-01-05"],
|
||
|
),
|
||
|
(
|
||
|
["2011-01-01", "2011-01-03", "2011-01-05", "2011-01-02", "2011-01-01"],
|
||
|
["2011-01-01", "2011-01-01", "2011-01-02", "2011-01-03", "2011-01-05"],
|
||
|
),
|
||
|
(
|
||
|
[NaT, "2011-01-03", "2011-01-05", "2011-01-02", NaT],
|
||
|
[NaT, NaT, "2011-01-02", "2011-01-03", "2011-01-05"],
|
||
|
),
|
||
|
],
|
||
|
)
|
||
|
def test_sort_values_without_freq_datetimeindex(
|
||
|
self, index_dates, expected_dates, tz_naive_fixture
|
||
|
):
|
||
|
tz = tz_naive_fixture
|
||
|
|
||
|
# without freq
|
||
|
idx = DatetimeIndex(index_dates, tz=tz, name="idx")
|
||
|
expected = DatetimeIndex(expected_dates, tz=tz, name="idx")
|
||
|
|
||
|
self.check_sort_values_without_freq(idx, expected)
|
||
|
|
||
|
@pytest.mark.parametrize(
|
||
|
"idx,expected",
|
||
|
[
|
||
|
(
|
||
|
PeriodIndex(
|
||
|
[
|
||
|
"2011-01-01",
|
||
|
"2011-01-03",
|
||
|
"2011-01-05",
|
||
|
"2011-01-02",
|
||
|
"2011-01-01",
|
||
|
],
|
||
|
freq="D",
|
||
|
name="idx1",
|
||
|
),
|
||
|
PeriodIndex(
|
||
|
[
|
||
|
"2011-01-01",
|
||
|
"2011-01-01",
|
||
|
"2011-01-02",
|
||
|
"2011-01-03",
|
||
|
"2011-01-05",
|
||
|
],
|
||
|
freq="D",
|
||
|
name="idx1",
|
||
|
),
|
||
|
),
|
||
|
(
|
||
|
PeriodIndex(
|
||
|
[
|
||
|
"2011-01-01",
|
||
|
"2011-01-03",
|
||
|
"2011-01-05",
|
||
|
"2011-01-02",
|
||
|
"2011-01-01",
|
||
|
],
|
||
|
freq="D",
|
||
|
name="idx2",
|
||
|
),
|
||
|
PeriodIndex(
|
||
|
[
|
||
|
"2011-01-01",
|
||
|
"2011-01-01",
|
||
|
"2011-01-02",
|
||
|
"2011-01-03",
|
||
|
"2011-01-05",
|
||
|
],
|
||
|
freq="D",
|
||
|
name="idx2",
|
||
|
),
|
||
|
),
|
||
|
(
|
||
|
PeriodIndex(
|
||
|
[NaT, "2011-01-03", "2011-01-05", "2011-01-02", NaT],
|
||
|
freq="D",
|
||
|
name="idx3",
|
||
|
),
|
||
|
PeriodIndex(
|
||
|
[NaT, NaT, "2011-01-02", "2011-01-03", "2011-01-05"],
|
||
|
freq="D",
|
||
|
name="idx3",
|
||
|
),
|
||
|
),
|
||
|
(
|
||
|
PeriodIndex(
|
||
|
["2011", "2013", "2015", "2012", "2011"], name="pidx", freq="A"
|
||
|
),
|
||
|
PeriodIndex(
|
||
|
["2011", "2011", "2012", "2013", "2015"], name="pidx", freq="A"
|
||
|
),
|
||
|
),
|
||
|
(
|
||
|
# For compatibility check
|
||
|
Index([2011, 2013, 2015, 2012, 2011], name="idx"),
|
||
|
Index([2011, 2011, 2012, 2013, 2015], name="idx"),
|
||
|
),
|
||
|
],
|
||
|
)
|
||
|
def test_sort_values_without_freq_periodindex(self, idx, expected):
|
||
|
# here without_freq means not generateable by period_range
|
||
|
self.check_sort_values_without_freq(idx, expected)
|
||
|
|
||
|
def test_sort_values_without_freq_periodindex_nat(self):
|
||
|
# doesn't quite fit into check_sort_values_without_freq
|
||
|
idx = PeriodIndex(["2011", "2013", "NaT", "2011"], name="pidx", freq="D")
|
||
|
expected = PeriodIndex(["NaT", "2011", "2011", "2013"], name="pidx", freq="D")
|
||
|
|
||
|
ordered = idx.sort_values(na_position="first")
|
||
|
tm.assert_index_equal(ordered, expected)
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
ordered = idx.sort_values(ascending=False)
|
||
|
tm.assert_index_equal(ordered, expected[::-1])
|
||
|
check_freq_nonmonotonic(ordered, idx)
|
||
|
|
||
|
|
||
|
def test_order_stability_compat():
|
||
|
# GH#35922. sort_values is stable both for normal and datetime-like Index
|
||
|
pidx = PeriodIndex(["2011", "2013", "2015", "2012", "2011"], name="pidx", freq="A")
|
||
|
iidx = Index([2011, 2013, 2015, 2012, 2011], name="idx")
|
||
|
ordered1, indexer1 = pidx.sort_values(return_indexer=True, ascending=False)
|
||
|
ordered2, indexer2 = iidx.sort_values(return_indexer=True, ascending=False)
|
||
|
tm.assert_numpy_array_equal(indexer1, indexer2)
|