Skip to content

Commit ba2e043

Browse files
TST: added tests for key option to series and frame sorting functions
1 parent 00a5dcb commit ba2e043

File tree

3 files changed

+116
-5
lines changed

3 files changed

+116
-5
lines changed

pandas/core/series.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3137,6 +3137,7 @@ def sort_index(
31373137
# Validate the axis parameter
31383138
self._get_axis_number(axis)
31393139
index = self.index
3140+
true_index = index
31403141
if key is not None:
31413142
index = index.map(key)
31423143

@@ -3169,16 +3170,12 @@ def sort_index(
31693170
else:
31703171
return self.copy()
31713172

3172-
if key is not None:
3173-
key_func = np.vectorize(key)
3174-
index = key_func(index)
3175-
31763173
indexer = nargsort(
31773174
index, kind=kind, ascending=ascending, na_position=na_position
31783175
)
31793176

31803177
indexer = ensure_platform_int(indexer)
3181-
new_index = index.take(indexer)
3178+
new_index = true_index.take(indexer)
31823179
new_index = new_index._sort_levels_monotonic()
31833180

31843181
new_values = self._values.take(indexer)

pandas/tests/frame/test_sorting.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -568,6 +568,64 @@ def test_sort_index_categorical_index(self):
568568
expected = df.iloc[[2, 3, 0, 1, 5, 4]]
569569
tm.assert_frame_equal(result, expected)
570570

571+
def test_sort_index_key(self):
572+
df = DataFrame(np.arange(6, dtype="int64"), index=list("aaBBca"))
573+
574+
result = df.sort_index()
575+
expected = df.iloc[[2, 3, 0, 1, 5, 4]]
576+
assert_frame_equal(result, expected)
577+
578+
result = df.sort_index(key=str.lower)
579+
expected = df.iloc[[0, 1, 5, 2, 3, 4]]
580+
assert_frame_equal(result, expected)
581+
582+
result = df.sort_index(key=str.lower, ascending=False)
583+
expected = df.iloc[[4, 2, 3, 0, 1, 5]]
584+
assert_frame_equal(result, expected)
585+
586+
@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float32', 'float64'])
587+
def test_sort_index_key_int(self, dtype):
588+
df = DataFrame(np.arange(6, dtype=dtype), index=np.arange(6, dtype=dtype))
589+
590+
result = df.sort_index()
591+
assert_frame_equal(result, df)
592+
593+
result = df.sort_index(key=lambda x : -x)
594+
expected = df.sort_index(ascending=False)
595+
assert_frame_equal(result, expected)
596+
597+
result = df.sort_index(key=lambda x : 2 * x)
598+
assert_frame_equal(result, df)
599+
600+
def test_sort_value_key(self):
601+
df = DataFrame(np.array([0, 5, np.nan, 3, 2, np.nan]))
602+
603+
result = df.sort_values(0)
604+
expected = df.iloc[[0, 4, 3, 1, 2, 5]]
605+
assert_frame_equal(result, expected)
606+
607+
result = df.sort_values(0, key=lambda x : x + 5)
608+
expected = df.iloc[[0, 4, 3, 1, 2, 5]]
609+
assert_frame_equal(result, expected)
610+
611+
result = df.sort_values(0, key=lambda x : -x, ascending=False)
612+
expected = df.iloc[[0, 4, 3, 1, 2, 5]]
613+
assert_frame_equal(result, expected)
614+
615+
def test_sort_value_key_nan(self):
616+
df = DataFrame(np.array([["hello", "goodbye"], ["hello", "Hello"]]))
617+
618+
result = df.sort_values(1)
619+
expected = df[::-1]
620+
assert_frame_equal(result, expected)
621+
622+
result = df.sort_values([0, 1], key=str.lower)
623+
assert_frame_equal(result, df)
624+
625+
result = df.sort_values([0, 1], key=str.lower, ascending=False)
626+
expected = df.sort_values(1, key=str.lower, ascending=False)
627+
assert_frame_equal(result, expected)
628+
571629
def test_sort_index(self):
572630
# GH13496
573631

pandas/tests/series/test_sorting.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,62 @@ def test_sort_index_kind(self):
169169
index_sorted_series = series.sort_index(kind="heapsort")
170170
tm.assert_series_equal(expected_series, index_sorted_series)
171171

172+
def test_sort_index_key(self):
173+
series = Series(np.arange(6, dtype="int64"), index=list("aaBBca"))
174+
175+
result = series.sort_index()
176+
expected = series.iloc[[2, 3, 0, 1, 5, 4]]
177+
assert_series_equal(result, expected)
178+
179+
result = series.sort_index(key=str.lower)
180+
expected = series.iloc[[0, 1, 5, 2, 3, 4]]
181+
assert_series_equal(result, expected)
182+
183+
result = series.sort_index(key=str.lower, ascending=False)
184+
expected = series.iloc[[4, 2, 3, 0, 1, 5]]
185+
assert_series_equal(result, expected)
186+
187+
@pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float32', 'float64'])
188+
def test_sort_index_key_int(self, dtype):
189+
series = Series(np.arange(6, dtype=dtype), index=np.arange(6, dtype=dtype))
190+
191+
result = series.sort_index()
192+
assert_series_equal(result, series)
193+
194+
result = series.sort_index(key=lambda x : -x)
195+
expected = series.sort_index(ascending=False)
196+
assert_series_equal(result, expected)
197+
198+
result = series.sort_index(key=lambda x: 2 * x)
199+
assert_series_equal(result, series)
200+
201+
202+
def test_sort_value_key(self):
203+
series = Series(np.array(["Hello", "goodbye"]))
204+
205+
result = series.sort_values(0)
206+
expected = series
207+
assert_series_equal(result, expected)
208+
209+
result = series.sort_values(0, key=str.lower)
210+
expected = series[::-1]
211+
assert_series_equal(result, expected)
212+
213+
def test_sort_value_key_nan(self):
214+
series = Series(np.array([0, 5, np.nan, 3, 2, np.nan]))
215+
216+
result = series.sort_values(0)
217+
expected = series.iloc[[0, 4, 3, 1, 2, 5]]
218+
assert_series_equal(result, expected)
219+
220+
result = series.sort_values(0, key=lambda x: x + 5)
221+
expected = series.iloc[[0, 4, 3, 1, 2, 5]]
222+
assert_series_equal(result, expected)
223+
224+
result = series.sort_values(0, key=lambda x: -x, ascending=False)
225+
expected = series.iloc[[0, 4, 3, 1, 2, 5]]
226+
assert_series_equal(result, expected)
227+
172228
def test_sort_index_na_position(self):
173229
series = Series(index=[3, 2, 1, 4, 3, np.nan])
174230

0 commit comments

Comments
 (0)