Added functionality: sort datasets by multiple keys
State: closed

Added functionality implementation: sort datasets by multiple keys/columns as discussed in https://github.com/huggingface/datasets/issues/5425.

_The documentation is not available anymore as the PR was closed or merged._

> Thanks! I've left some comments.
> 
> We should also add some tests, mainly to make sure `reverse` behaves as expected. Let me know if you need help with that.

Thanks for the offer! I couldn't find any guidelines on how huggingface goes about testing, so it would indeed be great to get a few pointers on that. I assume I should expand on the `test_sort` function in `test_arrow_dataset.py` but since I am not very familiar with the `datasets` package, it isn't immediately for which cases I should test (i.e., expand on).

@MichlF 

Resolving a comment means that the comment has been addressed with the code change, so since this is not the case here, can you please "unresolve" the comments and address them adequately? 

> I assume I should expand on the `test_sort` function in `test_arrow_dataset.py`

Yes, that's correct. I think one test to check sorting on multiple keys and another one to check if an error is raised when `len(reverse)!=len(column_names)` should be enough.


> Yes, that's correct. I think one test to check sorting on multiple keys and another one to check if an error is raised when `len(reverse)!=len(column_names)` should be enough.

I have added the tests in https://github.com/huggingface/datasets/pull/5502/commits/0efa259732e822e94d67b96a70031a3daccedfc1 by keeping them in the same format of the tests of the old `sort` function. Let me know if they can be improved.


<details>
<summary>Show benchmarks</summary>

PyArrow==6.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.010170 / 0.011353 (-0.001183) | 0.005891 / 0.011008 (-0.005117) | 0.100416 / 0.038508 (0.061908) | 0.041309 / 0.023109 (0.018200) | 0.300813 / 0.275898 (0.024915) | 0.376679 / 0.323480 (0.053199) | 0.008806 / 0.007986 (0.000821) | 0.005964 / 0.004328 (0.001636) | 0.075862 / 0.004250 (0.071611) | 0.050370 / 0.037052 (0.013318) | 0.313365 / 0.258489 (0.054876) | 0.351184 / 0.293841 (0.057343) | 0.039556 / 0.128546 (-0.088991) | 0.012462 / 0.075646 (-0.063185) | 0.337141 / 0.419271 (-0.082130) | 0.049678 / 0.043533 (0.006145) | 0.298547 / 0.255139 (0.043408) | 0.317547 / 0.283200 (0.034347) | 0.113595 / 0.141683 (-0.028088) | 1.448467 / 1.452155 (-0.003688) | 1.501303 / 1.492716 (0.008587) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.011005 / 0.018006 (-0.007002) | 0.527430 / 0.000490 (0.526940) | 0.005073 / 0.000200 (0.004873) | 0.000100 / 0.000054 (0.000045) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.030377 / 0.037411 (-0.007034) | 0.116932 / 0.014526 (0.102406) | 0.124047 / 0.176557 (-0.052509) | 0.192358 / 0.737135 (-0.544777) | 0.130528 / 0.296338 (-0.165811) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.401158 / 0.215209 (0.185949) | 4.005854 / 2.077655 (1.928200) | 1.810365 / 1.504120 (0.306245) | 1.626490 / 1.541195 (0.085295) | 1.752591 / 1.468490 (0.284101) | 0.709065 / 4.584777 (-3.875712) | 3.893356 / 3.745712 (0.147643) | 3.655180 / 5.269862 (-1.614682) | 1.873660 / 4.565676 (-2.692017) | 0.085860 / 0.424275 (-0.338415) | 0.012671 / 0.007607 (0.005063) | 0.512804 / 0.226044 (0.286759) | 5.103426 / 2.268929 (2.834497) | 2.336148 / 55.444624 (-53.108477) | 2.000140 / 6.876477 (-4.876336) | 2.095155 / 2.142072 (-0.046918) | 0.848612 / 4.805227 (-3.956615) | 0.171840 / 6.500664 (-6.328824) | 0.064144 / 0.075469 (-0.011325) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.222106 / 1.841788 (-0.619682) | 15.828559 / 8.074308 (7.754251) | 14.995298 / 10.191392 (4.803906) | 0.172783 / 0.680424 (-0.507641) | 0.029296 / 0.534201 (-0.504905) | 0.447469 / 0.579283 (-0.131814) | 0.658615 / 0.434364 (0.224251) | 1.527607 / 0.540337 (0.987270) | 1.830018 / 1.386936 (0.443082) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.007922 / 0.011353 (-0.003431) | 0.005369 / 0.011008 (-0.005639) | 0.076580 / 0.038508 (0.038071) | 0.038770 / 0.023109 (0.015661) | 0.338995 / 0.275898 (0.063097) | 0.380865 / 0.323480 (0.057385) | 0.006489 / 0.007986 (-0.001497) | 0.004421 / 0.004328 (0.000093) | 0.074143 / 0.004250 (0.069893) | 0.054224 / 0.037052 (0.017171) | 0.348887 / 0.258489 (0.090397) | 0.395044 / 0.293841 (0.101203) | 0.037040 / 0.128546 (-0.091507) | 0.012547 / 0.075646 (-0.063099) | 0.087521 / 0.419271 (-0.331751) | 0.049918 / 0.043533 (0.006385) | 0.342428 / 0.255139 (0.087289) | 0.362216 / 0.283200 (0.079016) | 0.107204 / 0.141683 (-0.034479) | 1.509206 / 1.452155 (0.057052) | 1.596010 / 1.492716 (0.103293) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.246795 / 0.018006 (0.228788) | 0.505998 / 0.000490 (0.505509) | 0.000446 / 0.000200 (0.000246) | 0.000064 / 0.000054 (0.000009) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.031591 / 0.037411 (-0.005821) | 0.117595 / 0.014526 (0.103069) | 0.132500 / 0.176557 (-0.044056) | 0.202244 / 0.737135 (-0.534891) | 0.136624 / 0.296338 (-0.159715) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.428235 / 0.215209 (0.213026) | 4.262691 / 2.077655 (2.185036) | 2.057348 / 1.504120 (0.553228) | 1.928559 / 1.541195 (0.387364) | 2.120838 / 1.468490 (0.652347) | 0.706300 / 4.584777 (-3.878477) | 3.951828 / 3.745712 (0.206115) | 2.144218 / 5.269862 (-3.125644) | 1.359500 / 4.565676 (-3.206177) | 0.085404 / 0.424275 (-0.338872) | 0.012363 / 0.007607 (0.004756) | 0.529985 / 0.226044 (0.303941) | 5.295831 / 2.268929 (3.026903) | 2.522602 / 55.444624 (-52.922022) | 2.182850 / 6.876477 (-4.693627) | 2.270187 / 2.142072 (0.128114) | 0.841676 / 4.805227 (-3.963551) | 0.168366 / 6.500664 (-6.332298) | 0.065371 / 0.075469 (-0.010098) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.261464 / 1.841788 (-0.580324) | 17.010125 / 8.074308 (8.935817) | 14.304453 / 10.191392 (4.113061) | 0.177782 / 0.680424 (-0.502642) | 0.017762 / 0.534201 (-0.516439) | 0.427283 / 0.579283 (-0.152000) | 0.455176 / 0.434364 (0.020812) | 0.525962 / 0.540337 (-0.014375) | 0.625583 / 1.386936 (-0.761353) |

</details>
</details>

![](https://cml.dev/watermark.png#3b2aba6637dc61f145acda40e4e7b028c3947d72 "CML watermark")