ENH: Added key option to df/series.sort_values(key=...) and df/series.sort_index(key=...) sorting
State: closed

- [x] closes #3942
- [x] tests added / passed
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry

Added a key parameter to the DataFrame/Series.sort_values and sort_index functions matching Python sorted semantics for allowing custom sorting orders. Address open issue https://github.com/pandas-dev/pandas/issues/3942. 


Labels:
name: Enhancement
name: Reshaping
description: Concat, Merge/Join, Stack/Unstack, Explode

Comments:
From typing import Callable at the top of the module then just Callable for the annotation

Sent from my iPhone

> On Jul 6, 2019, at 10:41 AM, Jacob Austin <notifications@github.com> wrote:
> 
> @ja3067 commented on this pull request.
> 
> In pandas/core/frame.py:
> 
> > @@ -4977,6 +4977,7 @@ def sort_values(
>          inplace=False,
>          kind="quicksort",
>          na_position="last",
> +        key=None
> Would you just like key : typing.Callable?
> 
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub, or mute the thread.

Hello @jacobaustin123! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:



































There are currently no PEP 8 issues detected in this Pull Request. Cheers! :beers: 

##### Comment last updated at 2020-04-27 02:12:02 UTC
Just pushed a new commit that adds support for Index.sort_values(key=...) and adds a key flag to nargsort(...) and lexsort_indexer(...). I've also added static type hints. I still use Index.map to apply the key function for sort_index functions, because it seems simpler and semantically better to apply the key to the values and not the codes. 
@ja3067 is this still active?
I will make the requested changes this weekend. The only decision I think needs to be made is how `sort_values(key=...)` and `sort_index(key=...)` should behave for multiple levels or columns. Should the same function be applied to every level/column, or should the key function be expected to operate on a tuple? Likewise, how should the key behave for `sort_remaining`? Should it be applied to everything? 
> The only decision I think needs to be made is how `sort_values(key=...)` and `sort_index(key=...)` should behave for multiple levels or columns.

Without a `key` is just operates on the levels individually, no? If not mistaken something should be the same here
> Without a `key` is just operates on the levels individually, no? If not mistaken something should be the same here

@WillAyd Agreed. In the case, say, of one numeric column and one string column, should we for instance support a dictionary of {column_name : Callable} so that different keys can be passed, or the key can be passed only to the primary column. Same question for sort_index.
I would leave that enhancement to a separate PR if requested
@ja3067 is this still active?
@WillAyd I'll make the necessary updates this weekend. Sorry for the delay.
@WillAyd Ok. The final PR is almost done. Last question is about the `key` for `sortindex` for a MultiIndex when `level` is specified. There are two options:
1. the key function is applied to the entire MultiIndex tuple like the map function does currently. You could do something like `key=lambda x : x[0].lower()`
2. the key function is applied to each level in the MultiIindex separately. The easiest way to do this would be to add a method to the MultiIndex called `idx.maplevel(self, mapper, level=0)` that applies a map function to a set of levels. However, this would probably require modifying the Index API. Would this be OK/worthwhile? This would be a little less flexible, but maybe more reasonable. 

What do you think?
@WillAyd I've incorporated the review changes, and it passes tests now. Should be good.
Can you merge master and repush? Not sure what 37 error was may be resolved
@WillAyd Passes all tests. Also fixed some new merge conflicts that showed up. 
@jreback I've incorporated most of your changes locally. there are a few mypy issues I can't figure out that I've made comments about, and I want to know your thoughts on `ensure_key_mapped`. I'm digging pretty deep into the pandas core, and I want this to be "Pandathonic". 
also needs a rebase
@WillAyd @jreback @TomAugspurger I've made changes in response to all the comments. I consolidated all of the key code into `assert_key_mapped`, and added a bunch of documentation, tests, and docstrings. 

I `type: ignore[override]` the inheritance issue, but there's a separate `flake8/pyflakes` issue that raises spurious errors for that syntax (which is fixed in master), so we have to ignore `flake8` on those two lines until there's a new release. 

There's also a typing issue with the `DatetimeIndexOpsMixin` class. None of it has been type annotated, and I think that's what has prevented mypy errors so far. But when type hints are added, it starts raising issues likes `error: "DatetimeIndexOpsMixin" has no attribute "_get_attributes_dict"` (`datetimelike.py:277`), since it is intended to be "mixed in" to a class with those attributes. I added an `isinstance` assertion that makes mypy happy, but it's going to be a bigger issue in the future. 

Otherwise it seems good. I still use `Index::map` to simplify `assert_key_mapped` for various `Index` subclasses. Let me know if there are other changes I should make. I can rebase the whole thing and force push if it looks good.  
can you merge master and will look again
@jreback merged master. let me know what you think. would be great to merge this soon.
Just to confirm one behavior thing, for `DataFrame.sort_values(by=list_of_cols)` we currently call `key` on each scalar value, once for each value in every column in `by`.

```python
In [34]: df = pd.DataFrame({"A": ['a', "a", "B"], "B": ["a", "b", "c"], "C": [1, 2, 3]})

In [35]: def key(x):
    ...:     print(x)
    ...:     return x.lower()
    ...:

In [36]: df.sort_values(['A', 'B'], key=key)
a
a
B
a
b
c
Out[36]:
   A  B  C
0  a  a  1
1  a  b  2
2  B  c  3
```

Alternatively, we could pass a tuple of values to `key` like `('a', 'a')`, `('a', 'b')`, etc. In this case, `key` would expect a tuple of values, rather than a scalar for DataFrame.sort_values. It depends a bit on your mental model of how `sort_values` works.

Which do people think would be more useful in practice?
> Which do people think would be more useful in practice?

I've been wondering the same question; it extends to MultiIndices as well. In the current PR, `sort_values` applies to each scalar element, but `sort_index` takes a tuple of MultiIndex values. 

I think a scalar is much more intuitive for `sort_values`, even though it's a little less powerful. You can pass things like `str.lower` (probably the most common use case), instead of `lambda x : [y.lower() for y in x]`. A future PR could even make `key` accept a dictionary of Callables, to achieve the same effect? 

As for `sort_index`, the tuple approach is just more easier to implement. The current behavior is as follows:

```
>>> df = DataFrame(
...     {"a": [3, 1, 2], "b": [0, 0, 0], "c": [0, 1, 2], "d": list("abc")}
... ).set_index(list("abc"))
>>> df
       d
a b c
3 0 0  a
1 0 1  b
2 0 2  c
>>> df.sort_index(
...     level=["b", "a"], key=lambda x: x[0]
... )
       d
a b c
1 0 1  b
2 0 2  c
3 0 0  a
>>> df.sort_index(
... level=["b", "a"], key=lambda x: [-y for y in x]
... )
       d
a b c
3 0 0  a
2 0 2  c
1 0 1  b
```

In other words, `sort_index` takes a tuple and applies it using the Index.map function. I think this is the most sensible functionality, but it may be a bit inconsistent between the `sort_index` and `sort_values` implementation. What do other people think?
> I think a scalar is much more intuitive for sort_values, even though it's a little less powerful. You can pass things like str.lower (probably the most common use case),

It's not clear to me which behavior is easier. If you're sorting homogenous data, then yes a `key` taking a scalar is much easier. Especially if you want to reuse `key` across potentially different DataFrames.

But for heterogenous data, you would need something like

```python
def key(val):
    if isinstance(val, str): return val.lower()
    elif isinstance(val, int): return -val
```

vs. something like 

```python
def key(val):
    return val[0].lower(), -val[1]
```

Neither is obviously better in all situations. The right thing to do here isn't clear to me.
And I definitely don't want to delay merging the clear improvement for Index and Series. So if we aren't able to come to an agreement quickly (within the next day?) then perhaps we raise when a `key` is passed to DataFrame.sort_values?
> Neither is obviously better in all situations. The right thing to do here isn't clear to me.

I agree that heterogeneous data is tricky. If it were a tuple, it would just be a tuple of the columns passed to `by`? That would be ok since most of the time you'd only pass one column at a time to `sort_values`. I think that would be ok too – it's more or less the way `sort_index` works.

My other idea for this would be a dictionary based on keys, i.e. if the key is a dictionary, the keys of each column could be used to look up Callables in the dictionary to apply to each column. That might be inconsistent with the use of `dict` keys in other areas like `Index.map`. 
@ja3067 can u merge master
@jorisvandenbossche @jbrockmendel @WillAyd / others any thoughts on the API discussion I raised in https://github.com/pandas-dev/pandas/pull/27237#issuecomment-569814979? When sorting by multiple columns is it more useful for `key` to get each scalar value sequentially, or a tuple of values in that row?

I'm not sure, but it feels to me like getting a tuple is more closely aligned with the behavior of the builtin `sorted`. When I do `sorted([('a', 1), ('a', 0)], key=lambda x: x[1])`, the key function is seeing *elements of the sequence*. When sorting a DataFrame, my mental model is that we're sorting rows, so the elements of the sequence are rows, which can be represented as a tuple.
@TomAugspurger I don't have a strong opinion here, agree that matching the behavior of the builtin is desirable.  Can we avoid materializing all the tuples though?  Is there an option that lets us consistently vectorize the key application?
Is the application of `key` currently vectorized? I haven't looked closely at the implementation, but I see we're using `lib.map_infer_mask` (does applying a Python function in a for loop count as vectorized?).

Right now my preference is probably to remove `key` from DataFrame.sort_values (or keep it in the signature but raise a NotImplementedError) until we can figure this out.
I think the scalar approach makes sense

Sent from my iPhone

> On Jan 2, 2020, at 4:07 PM, Tom Augspurger <notifications@github.com> wrote:
> 
> ﻿
> Is the application of key currently vectorized? I haven't looked closely at the implementation, but I see we're using lib.map_infer_mask (does applying a Python function in a for loop count as vectorized?).
> 
> Right now my preference is probably to remove key from DataFrame.sort_values (or keep it in the signature but raise a NotImplementedError) until we can figure this out.
> 
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub, or unsubscribe.

I agree with https://github.com/pandas-dev/pandas/pull/27237#issuecomment-570239064, and tuple is useful here. In any event, @ja3067 let's merge this (pls merge master), then open an issue about this.