Skip to content

GH-38386: [C++][Python] Support dictionary sort keys in ChunkedArray,… - #54681

Open
robbmcleod wants to merge 1 commit into
apache:mainfrom
robbmcleod:GH-38386-dictionary-sort-keys
Open

robbmcleod wants to merge 1 commit into
apache:mainfrom
robbmcleod:GH-38386-dictionary-sort-keys

Conversation

@robbmcleod

@robbmcleod robbmcleod commented Oct 20, 2026 •

Copy link
Copy Markdown

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:

  • Human
  • AI

Reviewed before submission by:

  • Human
  • AI
  • Not reviewed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant