Sort on multiple keys with datasets.Dataset.sort()
COLLABORATOR
Hi! 

`Dataset.sort` calls `df.sort_values` internally, and `df.sort_values` brings all the "sort" columns in memory, so sorting on multiple keys could be very expensive. This makes me think that maybe we can replace `df.sort_values` with `pyarrow.compute.sort_indices` - the latter can also sort on multiple keys and currently loads the data into memory; however, there is a plan to eventually implement "memory-map" friendly kernels for the Arrow compute ops (using the Acero execution engine). 

So to address this issue, you should replace `df.sort_values` with `pyarrow.compute.sort_indices` in `Dataset.sort` and adjust the signature of this function (deprecate the `kind` parameter, etc.).

PS: Feel free to ping us if you need some additional help/pointers
2023-01-19T16:40:48Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1397280202
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5TSM3K
https://api.github.com/repos/huggingface/datasets/issues/comments/1397280202/reactions
2023-01-30T13:12:40Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1397280202
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
CONTRIBUTOR
@mariosasko If I understand the code right, using `pyarrow.compute.sort_indices` would also require changes to the `select` method if it is meant to sort multiple keys. That's because `select` only accepts 1D input for `indices`, not an iterable or similar which would be required for multiple keys unless you want some looping over selects. Doesn't seem that straight-forward but I might be missing something here... 
2023-02-01T14:02:38Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412108992
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UKxLA
https://api.github.com/repos/huggingface/datasets/issues/comments/1412108992/reactions
2023-02-01T14:02:38Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412108992
https://avatars.githubusercontent.com/u/7805682?v=4
https://api.github.com/users/MichlF/events{/privacy}
https://api.github.com/users/MichlF/followers
https://api.github.com/users/MichlF/following{/other_user}
https://api.github.com/users/MichlF/gists{/gist_id}
https://github.com/MichlF
MichlF
MDQ6VXNlcjc4MDU2ODI=
https://api.github.com/users/MichlF/orgs
https://api.github.com/users/MichlF/received_events
https://api.github.com/users/MichlF/repos
https://api.github.com/users/MichlF/starred{/owner}{/repo}
https://api.github.com/users/MichlF/subscriptions
User
https://api.github.com/users/MichlF
public
COLLABORATOR
@MichlF No, it doesn't require modifying select because sorting on multiple keys also returns a 1D array.

It's easier to understand with an example:
```python
>>> import pyarrow as pa
>>> import pyarrow.compute as pc
>>> table = pa.table({
...     "name": ["John", "Eve", "Peter", "John"],
...     "surname": ["Johnson", "Smith", "Smith", "Doe"],
...     "age": [20, 40, 30, 50],
... })
>>> indices = pc.sort_indices(table, sort_keys=[("name", "ascending"), ("surname", "ascending")])
>>> print(indices)
[
  1,
  3,
  0,
  2
]
```


2023-02-01T15:37:06Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412264761
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5ULXM5
https://api.github.com/repos/huggingface/datasets/issues/comments/1412264761/reactions
2023-02-01T15:37:06Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412264761
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
CONTRIBUTOR
Thanks for clarifying.
I can prepare a PR to address this issue. This would be my first PR here so I have a few maybe silly questions but:
- What is the preferred input type of `sort_keys` for the sort method? A sequence with name, order tuples like pyarrow's `sort_indices` requires?
- What about backwards compatability: is it supposed to also accept the old way of calling sort() or should both `column` and `kind` be deprecated?
- If `sort_keys` is provided in the same format as for pyarrow's `sort_indices` - i.e. along with order for each column -, `reverse` doesn't make much sense either and should be deprecated as well I assume.
2023-02-01T16:28:49Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412349666
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5ULr7i
https://api.github.com/repos/huggingface/datasets/issues/comments/1412349666/reactions
2023-02-01T16:31:39Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412349666
https://avatars.githubusercontent.com/u/7805682?v=4
https://api.github.com/users/MichlF/events{/privacy}
https://api.github.com/users/MichlF/followers
https://api.github.com/users/MichlF/following{/other_user}
https://api.github.com/users/MichlF/gists{/gist_id}
https://github.com/MichlF
MichlF
MDQ6VXNlcjc4MDU2ODI=
https://api.github.com/users/MichlF/orgs
https://api.github.com/users/MichlF/received_events
https://api.github.com/users/MichlF/repos
https://api.github.com/users/MichlF/starred{/owner}{/repo}
https://api.github.com/users/MichlF/subscriptions
User
https://api.github.com/users/MichlF
public
COLLABORATOR
I think we can have the following signature:
```python
def sort(
        self,
        column_names: Union[str, Sequence[str]],
        reverse: Union[bool, Sequence[bool]] = False,
        kind="deprecated",
        null_placement: str = "last",
        keep_in_memory: bool = False,
        load_from_cache_file: bool = True,
        indices_cache_file_name: Optional[str] = None,
        writer_batch_size: Optional[int] = 1000,
        new_fingerprint: Optional[str] = None,
    ) -> "Dataset":
``` 

So we should:
* rename`column` to `column_names`. `column` is a positional argument, so it's OK to rename it (not marked as positional-only with "/", but still should be fine)
* deprecate `kind`
* keep `reverse` instead of introducing `sort_keys`, but we should allow passing a list of booleans that defines the sort order of each column from `column_names` to it (`reverse = False` would be equal to `[False] * len(column_names)` and `reverse = True` to `[True] * len(column_names)`)
2023-02-01T17:07:41Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412410080
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UL6rg
https://api.github.com/repos/huggingface/datasets/issues/comments/1412410080/reactions
2023-02-01T17:07:41Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412410080
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
CONTRIBUTOR
I am pretty much done with the PR. Just one clarification: `Sequence` in `arrow_dataset.py` is a custom dataclass from `features.py` instead of the `type.hinting` class `Sequence` from Python. Do you suggest using that custom `Sequence` class somehow ? Otherwise signature currently reads instead:
```Python
    def sort(
        self,
        column_names: Union[str, List[str]],
        reverse: Union[bool, List[bool]] = False,
        kind = "deprecated",
        null_placement: str = "last",
        keep_in_memory: bool = False,
        load_from_cache_file: bool = True,
        indices_cache_file_name: Optional[str] = None,
        writer_batch_size: Optional[int] = 1000,
        new_fingerprint: Optional[str] = None,
    )
```

Also, to maintain backwards compatibility, I added conditionals for `null_placement`, because pyarrow's `null_placement` only accepts `at_start` and `at_end`, and not `last` and `first`.
If that is all good, I think I can open the PR.
2023-02-01T18:38:49Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412539769
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UMaV5
https://api.github.com/repos/huggingface/datasets/issues/comments/1412539769/reactions
2023-02-01T18:41:13Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412539769
https://avatars.githubusercontent.com/u/7805682?v=4
https://api.github.com/users/MichlF/events{/privacy}
https://api.github.com/users/MichlF/followers
https://api.github.com/users/MichlF/following{/other_user}
https://api.github.com/users/MichlF/gists{/gist_id}
https://github.com/MichlF
MichlF
MDQ6VXNlcjc4MDU2ODI=
https://api.github.com/users/MichlF/orgs
https://api.github.com/users/MichlF/received_events
https://api.github.com/users/MichlF/repos
https://api.github.com/users/MichlF/starred{/owner}{/repo}
https://api.github.com/users/MichlF/subscriptions
User
https://api.github.com/users/MichlF
public
COLLABORATOR
I meant `typing.Sequence` (`datasets.Sequence` is a feature type). 

Regarding `null_placement`,  I think we can support both `at_start` and `at_end`, and  `last` and `first` (for backward compatibility; convert internally to `at_end` and `at_start` respectively).
2023-02-01T20:24:48Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412677247
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UM75_
https://api.github.com/repos/huggingface/datasets/issues/comments/1412677247/reactions
2023-02-01T20:24:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412677247
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
CONTRIBUTOR
> I meant typing.Sequence (datasets.Sequence is a feature type).

Sorry, I actually meant `typing.Sequence` and not `type.hinting`. However, the issue is still that `dataset.Sequence` is imported in `arrow_dataset.py` so I cannot import and use `typing.Sequence` for the `sort`'s signature without overwriting the `dataset.Sequence` import. The latter is used in the `align_labels_with_mapping` method so it's a necessary import for `arrow_dataset.py`. 
To import `typing.Sequence` as something else than `Sequence` to avoid overwriting may only be confusing and doesn't seem good practice!? The other solution is to keep `List` type hinting as in the signature I posted in my previous post but this excludes other Sequence types and may cause problems further down the line.
Please advise,
Thanks for all the clarifications!
2023-02-01T20:52:03Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1412704980
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UNCrU
https://api.github.com/repos/huggingface/datasets/issues/comments/1412704980/reactions
2023-02-01T21:00:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1412704980
https://avatars.githubusercontent.com/u/7805682?v=4
https://api.github.com/users/MichlF/events{/privacy}
https://api.github.com/users/MichlF/followers
https://api.github.com/users/MichlF/following{/other_user}
https://api.github.com/users/MichlF/gists{/gist_id}
https://github.com/MichlF
MichlF
MDQ6VXNlcjc4MDU2ODI=
https://api.github.com/users/MichlF/orgs
https://api.github.com/users/MichlF/received_events
https://api.github.com/users/MichlF/repos
https://api.github.com/users/MichlF/starred{/owner}{/repo}
https://api.github.com/users/MichlF/subscriptions
User
https://api.github.com/users/MichlF
public
COLLABORATOR
You can avoid the name collision by renaming `typing.Sequence` to `Sequence_` when importing:
```python
from typing import Sequence as Sequence_
```
2023-02-03T13:42:23Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1415886489
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5UZLaZ
https://api.github.com/repos/huggingface/datasets/issues/comments/1415886489/reactions
2023-02-03T13:42:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1415886489
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
COLLABORATOR
Resolved via #5502 
2023-02-24T16:15:10Z
https://github.com/huggingface/datasets/issues/5425#issuecomment-1443932528
https://api.github.com/repos/huggingface/datasets/issues/5425
IC_kwDODunzps5WEKlw
https://api.github.com/repos/huggingface/datasets/issues/comments/1443932528/reactions
2023-02-24T16:15:10Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1443932528
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
### Feature request

From discussion on forum: https://discuss.huggingface.co/t/datasets-dataset-sort-does-not-preserve-ordering/29065/1

`sort()` does not preserve ordering, and it does not support sorting on multiple columns, nor a key function.

The suggested solution: 
> ... having something similar to pandas and be able to specify multiple columns for sorting. We’re already using pandas under the hood to do the sorting in datasets.

The suggested workaround:
> convert your dataset to pandas and use `df.sort_values()`




### Motivation

Preserved ordering when sorting is very handy when one needs to sort on multiple columns, A and B, so that e.g. whenever A is equal for two or more rows, B is kept sorted. 
Having a parameter to do this in 🤗datasets would be cleaner than going through pandas and back, and it wouldn't add much complexity to the library.

Alternatives:
- the possibility to specify multiple keys to sort by with decreasing priority (suggested solution),
- the ability to provide a key function for sorting, so that one can manually specify the sorting criteria.

### Your contribution

I'll be happy to contribute by submitting a PR. Will get documented on `CONTRIBUTING.MD`.
Would love to get thoughts on this, if anyone has anything to add.