Repository navigation
GH-38386: [C++][Python] Support dictionary sort keys in ChunkedArray,… - #54681
Open
robbmcleod wants to merge 1 commit into
Open
robbmcleod wants to merge 1 commit into
robbmcleod wants to merge 1 commit into
Conversation
…Array, RecordBatch and Table sorting
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GH-38386: [C++][Python] Support dictionary sort keys in ChunkedArray, RecordBatch and Table sorting
Rationale for this change
GH-29887 made
array_sort_indicesaccept a dictionaryArray, but the othersort_indicespaths still reject dictionary columns.The existing
Arraypath 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_startand 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:
kdictionary 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 thekdictionary 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.
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.
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.
rankandselect_kstill 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?
TestDictionarySortIndicescases invector_sort_test.cccompare every result against the same sort on decoded data, across all orders and null placements. Tests include:test_sort_indices_dictionary,test_sort_indices_dictionary_nanandtest_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 usingtimestamp,float, andfixed_size_binary[16]for UUIDs.arrow.uuiddoesn't work but I know from prior work that the extension types don't work well in nested structures likedictionary, so I haven't pursued that.Benchmark. Median of 5 runs over 1M rows with 5 % nulls, using 'cardinality' unique valued strings and an
int64key with 10 values. "16 chunks, per-chunk" means each chunk was dictionary-encoded separately.stringThe last column is a Python (3.14) workaround we use internally: rank the dictionary with
pc.rank,takethrough the indices, and sort the result. This PR is essentially that algorithm implemented in C++ arrow.Are there any user-facing changes?
sort_indiceson aChunkedArray,RecordBatchorTable,array_sort_indiceson aChunkedArray, andTable.sort_by/RecordBatch.sort_bynow accept dictionary-encoded sort keys.As with
Array, dictionaries are ordered by their values, and theorderedflag 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:
Reviewed before submission by: