Sort on multiple keys with datasets.Dataset.sort()
Author: rocco-fortuna
State: closed
Labels: name: enhancement

name: good first issue

### 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.

Comments: 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

@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... 

@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
]
```



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.

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)`)

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.

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).

> 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!

You can avoid the name collision by renaming `typing.Sequence` to `Sequence_` when importing:
```python
from typing import Sequence as Sequence_
```

Resolved via #5502