TST: added tests for key option to series and frame sorting functions

Author: jacobaustin123

pandas/core/series.py

@@ -3134,6 +3134,7 @@ def sort_index(
         # Validate the axis parameter
         self._get_axis_number(axis)
         index = self.index
+        true_index = index
         if key is not None:
             index = index.map(key)
 
@@ -3166,16 +3167,12 @@ def sort_index(
                 else:
                     return self.copy()
 
-            if key is not None:
-                key_func = np.vectorize(key)
-                index = key_func(index)
-
             indexer = nargsort(
                 index, kind=kind, ascending=ascending, na_position=na_position
             )
 
         indexer = ensure_platform_int(indexer)
-        new_index = index.take(indexer)
+        new_index = true_index.take(indexer)
         new_index = new_index._sort_levels_monotonic()
 
         new_values = self._values.take(indexer)

pandas/tests/frame/test_sorting.py

@@ -500,6 +500,64 @@ def test_sort_index_categorical_index(self):
         expected = df.iloc[[2, 3, 0, 1, 5, 4]]
         tm.assert_frame_equal(result, expected)
 
+    def test_sort_index_key(self):
+        df = DataFrame(np.arange(6, dtype="int64"), index=list("aaBBca"))
+
+        result = df.sort_index()
+        expected = df.iloc[[2, 3, 0, 1, 5, 4]]
+        assert_frame_equal(result, expected)
+
+        result = df.sort_index(key=str.lower)
+        expected = df.iloc[[0, 1, 5, 2, 3, 4]]
+        assert_frame_equal(result, expected)
+
+        result = df.sort_index(key=str.lower, ascending=False)
+        expected = df.iloc[[4, 2, 3, 0, 1, 5]]
+        assert_frame_equal(result, expected)
+
+    @pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float32', 'float64'])
+    def test_sort_index_key_int(self, dtype):
+        df = DataFrame(np.arange(6, dtype=dtype), index=np.arange(6, dtype=dtype))
+
+        result = df.sort_index()
+        assert_frame_equal(result, df)
+
+        result = df.sort_index(key=lambda x : -x)
+        expected = df.sort_index(ascending=False)
+        assert_frame_equal(result, expected)
+
+        result = df.sort_index(key=lambda x : 2 * x)
+        assert_frame_equal(result, df)
+
+    def test_sort_value_key(self):
+        df = DataFrame(np.array([0, 5, np.nan, 3, 2, np.nan]))
+
+        result = df.sort_values(0)
+        expected = df.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_frame_equal(result, expected)
+
+        result = df.sort_values(0, key=lambda x : x + 5)
+        expected = df.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_frame_equal(result, expected)
+
+        result = df.sort_values(0, key=lambda x : -x, ascending=False)
+        expected = df.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_frame_equal(result, expected)
+
+    def test_sort_value_key_nan(self):
+        df = DataFrame(np.array([["hello", "goodbye"], ["hello", "Hello"]]))
+
+        result = df.sort_values(1)
+        expected = df[::-1]
+        assert_frame_equal(result, expected)
+
+        result = df.sort_values([0, 1], key=str.lower)
+        assert_frame_equal(result, df)
+
+        result = df.sort_values([0, 1], key=str.lower, ascending=False)
+        expected = df.sort_values(1, key=str.lower, ascending=False)
+        assert_frame_equal(result, expected)
+
     def test_sort_index(self):
         # GH13496
 

pandas/tests/series/test_sorting.py

@@ -169,6 +169,62 @@ def test_sort_index_kind(self):
         index_sorted_series = series.sort_index(kind="heapsort")
         tm.assert_series_equal(expected_series, index_sorted_series)
 
+    def test_sort_index_key(self):
+        series = Series(np.arange(6, dtype="int64"), index=list("aaBBca"))
+
+        result = series.sort_index()
+        expected = series.iloc[[2, 3, 0, 1, 5, 4]]
+        assert_series_equal(result, expected)
+
+        result = series.sort_index(key=str.lower)
+        expected = series.iloc[[0, 1, 5, 2, 3, 4]]
+        assert_series_equal(result, expected)
+
+        result = series.sort_index(key=str.lower, ascending=False)
+        expected = series.iloc[[4, 2, 3, 0, 1, 5]]
+        assert_series_equal(result, expected)
+
+    @pytest.mark.parametrize('dtype', ['int8', 'int16', 'int32', 'int64', 'float32', 'float64'])
+    def test_sort_index_key_int(self, dtype):
+        series = Series(np.arange(6, dtype=dtype), index=np.arange(6, dtype=dtype))
+
+        result = series.sort_index()
+        assert_series_equal(result, series)
+
+        result = series.sort_index(key=lambda x : -x)
+        expected = series.sort_index(ascending=False)
+        assert_series_equal(result, expected)
+
+        result = series.sort_index(key=lambda x: 2 * x)
+        assert_series_equal(result, series)
+
+
+    def test_sort_value_key(self):
+        series = Series(np.array(["Hello", "goodbye"]))
+
+        result = series.sort_values(0)
+        expected = series
+        assert_series_equal(result, expected)
+
+        result = series.sort_values(0, key=str.lower)
+        expected = series[::-1]
+        assert_series_equal(result, expected)
+
+    def test_sort_value_key_nan(self):
+        series = Series(np.array([0, 5, np.nan, 3, 2, np.nan]))
+
+        result = series.sort_values(0)
+        expected = series.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_series_equal(result, expected)
+
+        result = series.sort_values(0, key=lambda x: x + 5)
+        expected = series.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_series_equal(result, expected)
+
+        result = series.sort_values(0, key=lambda x: -x, ascending=False)
+        expected = series.iloc[[0, 4, 3, 1, 2, 5]]
+        assert_series_equal(result, expected)
+
     def test_sort_index_na_position(self):
         series = Series(index=[3, 2, 1, 4, 3, np.nan])