Fix `sort` with indices mapping
State: closed | Author: mariosasko (COLLABORATOR)

Fixes the `key` range in the `query_table` call in `sort` to account for an indices mapping

Fix #5586 

Comments: <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.008740 / 0.011353 (-0.002613) | 0.004501 / 0.011008 (-0.006507) | 0.100045 / 0.038508 (0.061537) | 0.029999 / 0.023109 (0.006890) | 0.303556 / 0.275898 (0.027658) | 0.335342 / 0.323480 (0.011863) | 0.006996 / 0.007986 (-0.000989) | 0.004183 / 0.004328 (-0.000145) | 0.076434 / 0.004250 (0.072183) | 0.033899 / 0.037052 (-0.003153) | 0.301312 / 0.258489 (0.042823) | 0.343136 / 0.293841 (0.049295) | 0.034062 / 0.128546 (-0.094484) | 0.011465 / 0.075646 (-0.064181) | 0.323134 / 0.419271 (-0.096137) | 0.040820 / 0.043533 (-0.002713) | 0.301708 / 0.255139 (0.046569) | 0.329528 / 0.283200 (0.046328) | 0.088393 / 0.141683 (-0.053290) | 1.460996 / 1.452155 (0.008842) | 1.531145 / 1.492716 (0.038429) |

### 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.191918 / 0.018006 (0.173912) | 0.414099 / 0.000490 (0.413610) | 0.000411 / 0.000200 (0.000211) | 0.000060 / 0.000054 (0.000005) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022707 / 0.037411 (-0.014704) | 0.096991 / 0.014526 (0.082465) | 0.106070 / 0.176557 (-0.070487) | 0.151275 / 0.737135 (-0.585860) | 0.108909 / 0.296338 (-0.187430) |

### 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.422499 / 0.215209 (0.207289) | 4.205551 / 2.077655 (2.127896) | 1.918960 / 1.504120 (0.414841) | 1.715421 / 1.541195 (0.174227) | 1.768969 / 1.468490 (0.300479) | 0.692243 / 4.584777 (-3.892534) | 3.382452 / 3.745712 (-0.363260) | 1.943695 / 5.269862 (-3.326166) | 1.250482 / 4.565676 (-3.315195) | 0.082084 / 0.424275 (-0.342191) | 0.012446 / 0.007607 (0.004839) | 0.525584 / 0.226044 (0.299539) | 5.275530 / 2.268929 (3.006602) | 2.386207 / 55.444624 (-53.058418) | 2.043920 / 6.876477 (-4.832557) | 2.030932 / 2.142072 (-0.111140) | 0.810233 / 4.805227 (-3.994994) | 0.148139 / 6.500664 (-6.352525) | 0.064617 / 0.075469 (-0.010852) |

### 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.227352 / 1.841788 (-0.614436) | 13.527623 / 8.074308 (5.453315) | 14.018551 / 10.191392 (3.827159) | 0.140333 / 0.680424 (-0.540091) | 0.028349 / 0.534201 (-0.505852) | 0.394904 / 0.579283 (-0.184379) | 0.406532 / 0.434364 (-0.027831) | 0.471714 / 0.540337 (-0.068624) | 0.568517 / 1.386936 (-0.818419) |

</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.006623 / 0.011353 (-0.004730) | 0.004464 / 0.011008 (-0.006544) | 0.076342 / 0.038508 (0.037834) | 0.027451 / 0.023109 (0.004341) | 0.343851 / 0.275898 (0.067953) | 0.385723 / 0.323480 (0.062243) | 0.005624 / 0.007986 (-0.002362) | 0.004685 / 0.004328 (0.000356) | 0.075669 / 0.004250 (0.071419) | 0.037297 / 0.037052 (0.000244) | 0.343363 / 0.258489 (0.084874) | 0.396115 / 0.293841 (0.102274) | 0.031577 / 0.128546 (-0.096970) | 0.011557 / 0.075646 (-0.064090) | 0.085626 / 0.419271 (-0.333645) | 0.041699 / 0.043533 (-0.001834) | 0.340826 / 0.255139 (0.085687) | 0.377167 / 0.283200 (0.093967) | 0.088632 / 0.141683 (-0.053051) | 1.464500 / 1.452155 (0.012345) | 1.556686 / 1.492716 (0.063969) |

### 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.231136 / 0.018006 (0.213130) | 0.402687 / 0.000490 (0.402197) | 0.000590 / 0.000200 (0.000390) | 0.000059 / 0.000054 (0.000004) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.024926 / 0.037411 (-0.012485) | 0.101062 / 0.014526 (0.086536) | 0.106481 / 0.176557 (-0.070075) | 0.159167 / 0.737135 (-0.577968) | 0.110948 / 0.296338 (-0.185390) |

### 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.441813 / 0.215209 (0.226603) | 4.416332 / 2.077655 (2.338677) | 2.080621 / 1.504120 (0.576501) | 1.877832 / 1.541195 (0.336637) | 1.944778 / 1.468490 (0.476288) | 0.704634 / 4.584777 (-3.880143) | 3.433955 / 3.745712 (-0.311758) | 1.863493 / 5.269862 (-3.406368) | 1.168869 / 4.565676 (-3.396807) | 0.084095 / 0.424275 (-0.340180) | 0.012440 / 0.007607 (0.004833) | 0.545122 / 0.226044 (0.319077) | 5.472214 / 2.268929 (3.203285) | 2.514580 / 55.444624 (-52.930044) | 2.164570 / 6.876477 (-4.711907) | 2.193467 / 2.142072 (0.051395) | 0.809056 / 4.805227 (-3.996171) | 0.152343 / 6.500664 (-6.348321) | 0.067610 / 0.075469 (-0.007859) |

### 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.280968 / 1.841788 (-0.560820) | 13.887674 / 8.074308 (5.813366) | 13.160405 / 10.191392 (2.969013) | 0.128601 / 0.680424 (-0.551823) | 0.016420 / 0.534201 (-0.517780) | 0.382810 / 0.579283 (-0.196473) | 0.394386 / 0.434364 (-0.039978) | 0.470254 / 0.540337 (-0.070083) | 0.566907 / 1.386936 (-0.820029) |

</details>
</details>

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


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

<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.008673 / 0.011353 (-0.002679) | 0.004475 / 0.011008 (-0.006533) | 0.102060 / 0.038508 (0.063552) | 0.029438 / 0.023109 (0.006329) | 0.351785 / 0.275898 (0.075887) | 0.388199 / 0.323480 (0.064719) | 0.007011 / 0.007986 (-0.000974) | 0.003317 / 0.004328 (-0.001012) | 0.080931 / 0.004250 (0.076681) | 0.033449 / 0.037052 (-0.003603) | 0.360329 / 0.258489 (0.101840) | 0.400069 / 0.293841 (0.106228) | 0.033628 / 0.128546 (-0.094918) | 0.011462 / 0.075646 (-0.064184) | 0.323781 / 0.419271 (-0.095490) | 0.040686 / 0.043533 (-0.002847) | 0.332715 / 0.255139 (0.077576) | 0.370339 / 0.283200 (0.087139) | 0.084633 / 0.141683 (-0.057050) | 1.459452 / 1.452155 (0.007297) | 1.547719 / 1.492716 (0.055003) |

### 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.187051 / 0.018006 (0.169045) | 0.402625 / 0.000490 (0.402135) | 0.002218 / 0.000200 (0.002018) | 0.000070 / 0.000054 (0.000016) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.025240 / 0.037411 (-0.012171) | 0.102201 / 0.014526 (0.087675) | 0.108629 / 0.176557 (-0.067927) | 0.156686 / 0.737135 (-0.580449) | 0.111383 / 0.296338 (-0.184955) |

### 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.418099 / 0.215209 (0.202890) | 4.163345 / 2.077655 (2.085690) | 1.868419 / 1.504120 (0.364300) | 1.662066 / 1.541195 (0.120871) | 1.705912 / 1.468490 (0.237422) | 0.696391 / 4.584777 (-3.888386) | 3.338307 / 3.745712 (-0.407405) | 1.923255 / 5.269862 (-3.346607) | 1.249220 / 4.565676 (-3.316457) | 0.082037 / 0.424275 (-0.342238) | 0.012232 / 0.007607 (0.004624) | 0.523913 / 0.226044 (0.297869) | 5.290036 / 2.268929 (3.021107) | 2.319729 / 55.444624 (-53.124896) | 1.987345 / 6.876477 (-4.889132) | 2.044516 / 2.142072 (-0.097556) | 0.812098 / 4.805227 (-3.993129) | 0.147327 / 6.500664 (-6.353337) | 0.063838 / 0.075469 (-0.011631) |

### 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.219652 / 1.841788 (-0.622136) | 13.271513 / 8.074308 (5.197205) | 13.799982 / 10.191392 (3.608590) | 0.150055 / 0.680424 (-0.530369) | 0.028804 / 0.534201 (-0.505397) | 0.395452 / 0.579283 (-0.183831) | 0.398758 / 0.434364 (-0.035606) | 0.468575 / 0.540337 (-0.071763) | 0.553324 / 1.386936 (-0.833612) |

</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.006498 / 0.011353 (-0.004855) | 0.004439 / 0.011008 (-0.006569) | 0.076525 / 0.038508 (0.038017) | 0.027184 / 0.023109 (0.004074) | 0.364705 / 0.275898 (0.088807) | 0.409481 / 0.323480 (0.086001) | 0.004831 / 0.007986 (-0.003154) | 0.004524 / 0.004328 (0.000196) | 0.075403 / 0.004250 (0.071153) | 0.039013 / 0.037052 (0.001960) | 0.364042 / 0.258489 (0.105553) | 0.413090 / 0.293841 (0.119249) | 0.032052 / 0.128546 (-0.096495) | 0.011514 / 0.075646 (-0.064132) | 0.085219 / 0.419271 (-0.334053) | 0.041448 / 0.043533 (-0.002085) | 0.350371 / 0.255139 (0.095232) | 0.386670 / 0.283200 (0.103470) | 0.089824 / 0.141683 (-0.051859) | 1.487392 / 1.452155 (0.035238) | 1.537201 / 1.492716 (0.044485) |

### 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.231555 / 0.018006 (0.213549) | 0.407505 / 0.000490 (0.407016) | 0.000382 / 0.000200 (0.000182) | 0.000060 / 0.000054 (0.000006) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.026665 / 0.037411 (-0.010747) | 0.105852 / 0.014526 (0.091326) | 0.108228 / 0.176557 (-0.068328) | 0.164164 / 0.737135 (-0.572972) | 0.114284 / 0.296338 (-0.182054) |

### 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.448957 / 0.215209 (0.233748) | 4.500058 / 2.077655 (2.422403) | 2.331660 / 1.504120 (0.827541) | 2.119904 / 1.541195 (0.578710) | 2.101489 / 1.468490 (0.632999) | 0.696580 / 4.584777 (-3.888197) | 3.364206 / 3.745712 (-0.381506) | 2.550157 / 5.269862 (-2.719704) | 1.496455 / 4.565676 (-3.069222) | 0.083289 / 0.424275 (-0.340986) | 0.012283 / 0.007607 (0.004676) | 0.555581 / 0.226044 (0.329537) | 5.556284 / 2.268929 (3.287355) | 2.595261 / 55.444624 (-52.849363) | 2.234793 / 6.876477 (-4.641683) | 2.280150 / 2.142072 (0.138078) | 0.817885 / 4.805227 (-3.987343) | 0.151481 / 6.500664 (-6.349183) | 0.066764 / 0.075469 (-0.008705) |

### 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.318875 / 1.841788 (-0.522913) | 14.220380 / 8.074308 (6.146072) | 13.922773 / 10.191392 (3.731381) | 0.154608 / 0.680424 (-0.525816) | 0.016343 / 0.534201 (-0.517858) | 0.380758 / 0.579283 (-0.198525) | 0.392595 / 0.434364 (-0.041769) | 0.468844 / 0.540337 (-0.071493) | 0.561047 / 1.386936 (-0.825889) |

</details>
</details>

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