ENH: added df/series.sort_values(key=...) and df/series.sort_index(key=...) functionality

Author: jacobaustin123

pandas/core/frame.py

@@ -4722,6 +4722,7 @@ def sort_values(
         inplace=False,
         kind="quicksort",
         na_position="last",
+        key = None
     ):
         inplace = validate_bool_kwarg(inplace, "inplace")
         axis = self._get_axis_number(axis)
@@ -4735,7 +4736,12 @@ def sort_values(
         if len(by) > 1:
             from pandas.core.sorting import lexsort_indexer
 
-            keys = [self._get_label_or_level_values(x, axis=axis) for x in by]
+            if key is not None:
+                key_func = np.vectorize(key)
+                keys = [key_func(self._get_label_or_level_values(x, axis=axis)) for x in by]
+            else:
+                keys = [self._get_label_or_level_values(x, axis=axis) for x in by]
+
             indexer = lexsort_indexer(keys, orders=ascending, na_position=na_position)
             indexer = ensure_platform_int(indexer)
         else:
@@ -4744,6 +4750,10 @@ def sort_values(
             by = by[0]
             k = self._get_label_or_level_values(by, axis=axis)
 
+            if key is not None:
+                key_func = np.vectorize(key)
+                k = key_func(k)
+
             if isinstance(ascending, (tuple, list)):
                 ascending = ascending[0]
 
@@ -4772,6 +4782,7 @@ def sort_index(
         na_position="last",
         sort_remaining=True,
         by=None,
+        key=None
     ):
 
         # TODO: this can be combined with Series.sort_index impl as
@@ -4792,21 +4803,23 @@ def sort_index(
 
         axis = self._get_axis_number(axis)
         labels = self._get_axis(axis)
-
+        if key is not None:
+            labels = labels.map(key)
+            
         # make sure that the axis is lexsorted to start
         # if not we need to reconstruct to get the correct indexer
         labels = labels._sort_levels_monotonic()
         if level is not None:
-
             new_axis, indexer = labels.sortlevel(
                 level, ascending=ascending, sort_remaining=sort_remaining
             )
 
         elif isinstance(labels, ABCMultiIndex):
             from pandas.core.sorting import lexsort_indexer
 
+            codes = labels._get_codes_for_sorting()
             indexer = lexsort_indexer(
-                labels._get_codes_for_sorting(),
+                codes,
                 orders=ascending,
                 na_position=na_position,
             )

pandas/core/series.py

@@ -2838,6 +2838,7 @@ def sort_values(
         inplace=False,
         kind="quicksort",
         na_position="last",
+        key=None
     ):
         """
         Sort by the values.
@@ -2955,6 +2956,10 @@ def sort_values(
             )
 
         def _try_kind_sort(arr):
+            if key is not None:
+                key_func = np.vectorize(key)
+                arr = key_func(arr)
+
             # easier to ask forgiveness than permission
             try:
                 # if kind==mergesort, it can fail for object dtype
@@ -3014,6 +3019,7 @@ def sort_index(
         kind="quicksort",
         na_position="last",
         sort_remaining=True,
+        key=None
     ):
         """
         Sort Series by index labels.
@@ -3131,17 +3137,22 @@ def sort_index(
         # Validate the axis parameter
         self._get_axis_number(axis)
         index = self.index
-
+        if key is not None:
+            index = index.map(key)
+        
         if level is not None:
             new_index, indexer = index.sortlevel(
                 level, ascending=ascending, sort_remaining=sort_remaining
             )
+
         elif isinstance(index, MultiIndex):
             from pandas.core.sorting import lexsort_indexer
 
             labels = index._sort_levels_monotonic()
+            codes = labels._get_codes_for_sorting()
+
             indexer = lexsort_indexer(
-                labels._get_codes_for_sorting(),
+                codes,
                 orders=ascending,
                 na_position=na_position,
             )
@@ -3158,6 +3169,10 @@ 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
             )