|
3 | 3 | import numpy as np
|
4 | 4 | import pytest
|
5 | 5 |
|
6 |
| -from pandas import IntervalIndex, MultiIndex, Series |
| 6 | +from pandas import DatetimeIndex, IntervalIndex, MultiIndex, Series |
7 | 7 | import pandas._testing as tm
|
8 | 8 |
|
9 | 9 |
|
@@ -223,7 +223,66 @@ def test_sort_index_key_int(self):
|
223 | 223 | result = series.sort_index(key=lambda x: 2 * x)
|
224 | 224 | tm.assert_series_equal(result, series)
|
225 | 225 |
|
| 226 | + def test_sort_index_kind_key(self, sort_by_key): |
| 227 | + # GH #14444 & #13589: Add support for sort algo choosing |
| 228 | + series = Series(index=[3, 2, 1, 4, 3], dtype=object) |
| 229 | + expected_series = Series(index=[1, 2, 3, 3, 4], dtype=object) |
| 230 | + |
| 231 | + index_sorted_series = series.sort_index(kind="mergesort", key=sort_by_key) |
| 232 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 233 | + |
| 234 | + index_sorted_series = series.sort_index(kind="quicksort", key=sort_by_key) |
| 235 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 236 | + |
| 237 | + index_sorted_series = series.sort_index(kind="heapsort", key=sort_by_key) |
| 238 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 239 | + |
| 240 | + def test_sort_index_kind_neg_key(self): |
| 241 | + # GH #14444 & #13589: Add support for sort algo choosing |
| 242 | + series = Series(index=[3, 2, 1, 4, 3], dtype=object) |
| 243 | + expected_series = Series(index=[4, 3, 3, 2, 1], dtype=object) |
| 244 | + |
| 245 | + index_sorted_series = series.sort_index(kind="mergesort", key=lambda x: -x) |
| 246 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 247 | + |
| 248 | + index_sorted_series = series.sort_index(kind="quicksort", key=lambda x: -x) |
| 249 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 250 | + |
| 251 | + index_sorted_series = series.sort_index(kind="heapsort", key=lambda x: -x) |
| 252 | + tm.assert_series_equal(expected_series, index_sorted_series) |
| 253 | + |
| 254 | + def test_sort_index_na_position_key(self, sort_by_key): |
| 255 | + series = Series(index=[3, 2, 1, 4, 3, np.nan], dtype=object) |
| 256 | + expected_series_first = Series(index=[np.nan, 1, 2, 3, 3, 4], dtype=object) |
| 257 | + |
| 258 | + index_sorted_series = series.sort_index(na_position="first", key=sort_by_key) |
| 259 | + tm.assert_series_equal(expected_series_first, index_sorted_series) |
| 260 | + |
| 261 | + expected_series_last = Series(index=[1, 2, 3, 3, 4, np.nan], dtype=object) |
| 262 | + |
| 263 | + index_sorted_series = series.sort_index(na_position="last", key=sort_by_key) |
| 264 | + tm.assert_series_equal(expected_series_last, index_sorted_series) |
| 265 | + |
226 | 266 | def test_changes_length_raises(self):
|
227 | 267 | s = Series([1, 2, 3])
|
228 | 268 | with pytest.raises(ValueError, match="change the shape"):
|
229 | 269 | s.sort_index(key=lambda x: x[:1])
|
| 270 | + |
| 271 | + def test_sort_values_key_type(self): |
| 272 | + s = Series([1, 2, 3], DatetimeIndex(["2008-10-24", "2008-11-23", "2007-12-22"])) |
| 273 | + |
| 274 | + result = s.sort_index(key=lambda x: x.month) |
| 275 | + expected = s.iloc[[0, 1, 2]] |
| 276 | + tm.assert_series_equal(result, expected) |
| 277 | + |
| 278 | + result = s.sort_index(key=lambda x: x.day) |
| 279 | + expected = s.iloc[[2, 1, 0]] |
| 280 | + tm.assert_series_equal(result, expected) |
| 281 | + |
| 282 | + result = s.sort_index(key=lambda x: x.year) |
| 283 | + expected = s.iloc[[2, 0, 1]] |
| 284 | + tm.assert_series_equal(result, expected) |
| 285 | + |
| 286 | + result = s.sort_index(key=lambda x: x.month_name()) |
| 287 | + expected = s.iloc[[2, 1, 0]] |
| 288 | + tm.assert_series_equal(result, expected) |
0 commit comments