robbmcleod opened a new pull request, #54681:
URL: https://github.com/apache/arrow/pull/54681
GH-38386: [C++][Python] Support dictionary sort keys in ChunkedArray,
RecordBatch and Table sorting
### Rationale for this change
GH-29887 made `array_sort_indices` accept a dictionary `Array`, but the
other `sort_indices` paths still reject dictionary columns.
The existing `Array` path also orders NaN wrongly in floating-point
dictionaries, which may or may not be a problem as it would be an unusual
use-case (from my point of view). It ranks NaN as the greatest value, so the
result differs from sorting the decoded values for ascending/`at_start` and
descending/`at_end`.
### What changes are included in this PR?
Every sort path now converts dictionary-encoded sort keys with one internal
helper, `DictionaryToSortable`. It extends the GH-29887 approach from one array
to many chunks.
Algorithm description:
1. One dictionary array: rank the `k` dictionary values, so equal values
share a rank and null values get a null rank. Then take those ranks through the
indices, giving one integer per row. Sorting those integers gives the same
order as sorting the strings, and string comparisons only happen while ranking
the `k` dictionary values.
a. If someone was to have a use case where there were far more keys than
rows in the container then this is inefficient, but I thought this would be
rare enough to not justify a branch.
2. Float dictionaries are decoded instead, because a rank can’t represent
NaN.
3. Chunked data: chunks can carry different dictionaries, so ranks have to
agree across them.
a. With a shared dictionary this is a non-issue.
b. Otherwise the distinct dictionaries are concatenated, deduplicated in
one hash pass, and only the unique values are ranked. Each chunk then takes its
slice of the ranks through its own indices.
4. Per container:
a. ChunkedArray: joins the rank chunks into one array, so it does one
counting sort instead of sorting each chunk and merging.
b. RecordBatch: converts dictionary columns while it resolves the sort
keys.
c. Table: converts each key across all batches at once, and every batch
reuses those table-wide ranks, so the batch merge compares consistent numbers.
`rank` and `select_k` still reject dictionary input and are left for a
follow-up. I can tack those onto this PR if desired using a similar approach.
### Are these changes tested?
- **C++:** new `TestDictionarySortIndices` cases in `vector_sort_test.cc`
compare every result against the same sort on decoded data, across all orders
and null placements. Tests include:
- chunks with shared, differing and many overlapping dictionaries
- null indices and null dictionary values
- empty and all-null chunks
- float dictionaries containing NaN
- other value types: timestamp (with and without time zone), date32,
time64, duration, int64, decimal128, binary, large_string and
fixed_size_binary(16)
- 2 and 9 sort keys, so both RecordBatch sorters run
- tables with mismatched chunking
- sort keys on nested struct fields
- **Python:** `test_sort_indices_dictionary`,
`test_sort_indices_dictionary_nan` and `test_table_sort_by_dictionary`.
I have approached this from the perspective of using `dictionary<int, str>`
internally so I may have blind spots for other patterns. I have run tests on
using `timestamp`, `float`, and `fixed_size_binary[16]` for UUIDs. `arrow.uuid`
doesn't work but I know from prior work that the extension types don't work
well in nested structures like `dictionary`, so I haven't pursued that.
**Benchmark.** Median of 5 runs over 1M rows with 5 % nulls, using
'cardinality' unique valued strings and an `int64` key with 10 values. "16
chunks, per-chunk" means each chunk was dictionary-encoded separately.
| Input | Keys | Cardinality | `string` | dictionary (this PR) | Speedup |
Python workaround |
|---|---|---|---|---|---|---|
| ChunkedArray, 1 chunk | 1 | 1,000 | 225 ms | 6.1 ms | 37x | 6.7 ms |
| ChunkedArray, 16 chunks, shared dictionary | 1 | 24 | 163 ms | 5.9 ms |
27x | 6.7 ms |
| ChunkedArray, 16 chunks, per-chunk dictionaries | 1 | 100,000 | 280 ms |
124 ms | 2.3x | 124 ms |
| RecordBatch | 2 | 1,000 | 270 ms | 95 ms | 2.8x | 90 ms |
| Table, 16 chunks, shared dictionary | 1 | 1,000 | 225 ms | 6.5 ms | 35x |
6.7 ms |
| Table, 16 chunks, shared dictionary | 2 | 100,000 | 359 ms | 159 ms | 2.3x
| 179 ms |
| Table, 16 chunks, per-chunk dictionaries | 2 | 100,000 | 392 ms | 213 ms |
1.8x | 200 ms |
The last column is a Python (3.14) workaround we use internally: rank the
dictionary with `pc.rank`, `take` through the indices, and sort the result.
This PR is essentially that algorithm implemented in C++ arrow.
### Are there any user-facing changes?
- `sort_indices` on a `ChunkedArray`, `RecordBatch` or `Table`,
`array_sort_indices` on a `ChunkedArray`, and `Table.sort_by` /
`RecordBatch.sort_by` now accept dictionary-encoded sort keys.
- Float dictionaries with NaN now sort the same as their decoded values.
As with `Array`, dictionaries are ordered by their values, and the `ordered`
flag is ignored.
### Was AI used for this PR?
The implementation, tests and this description were drafted with Claude
Claude Opus 5.5. That covers the helper, the sort-path changes, the C++ and
Python tests, and verification against decoded data. Code was reviewed
line-by-line and comments were edited by hand. The PR description was heavily
edited.
**PR code and description written by:**
- [x] Human
- [x] AI
**Reviewed before submission by:**
- [X] Human
- [x] AI
- [ ] Not reviewed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
