Repository navigation
GH-45010: [C++] Error out in unique/value_counts when unified dictionary overflows the index type - #55253
GH-45010: [C++] Error out in unique/value_counts when unified dictionary overflows the index type#55253ViderStar wants to merge 2 commits into
Conversation
…ictionary overflows the index type
|
|
| // The output keeps the input index type, so the unified dictionary must fit in it | ||
| std::shared_ptr<Array> out_dict; | ||
| RETURN_NOT_OK(dictionary_unifier_->GetResult(&out_dict_type, &out_dict)); | ||
| RETURN_NOT_OK(dictionary_unifier_->GetResultWithIndexType(index_type_, &out_dict)); |
There was a problem hiding this comment.
128 dictionary entries fit in int8: valid indices are 0..127. With two disjoint 64-entry dictionaries, both functions succeed on the parent commit but now return Invalid.
This helper checks the dictionary length instead of the highest index, right?
There was a problem hiding this comment.
Yep, the helper compares the dictionary length rather than the highest index, so exactly 128 entries were wrongly rejected for int8. Fixed in bfc4ac7: the kernel now uses GetResult() again and checks that length - 1 fits in the input index type. The test now covers the exact boundary (64 + 64 = 128 entries succeed for both unique and value_counts, 129 return Invalid).
Side note: GetResultWithIndexType itself has the same off-by-one, and it’s also used by Concatenate and UnifyChunkedArray. I left it out to keep this PR focused - happy to fix it here or in a follow-up, whichever you prefer =)
… type A unified dictionary of 128 entries is addressable by int8 indices (0..127), but DictionaryUnifier::GetResultWithIndexType rejects it because it compares the dictionary length with the index type maximum. Check length - 1 instead and test the exact 128/129 boundary for int8.
|
|
2 similar comments
|
|
|
|
Rationale for this change
uniqueandvalue_countson a chunked dictionary array whose chunks have different dictionaries unify the dictionaries and transpose the indices, but the output keeps the input index type. When the unified dictionary has more entries than that index type can address, the transposed indices wrap around and the result is an invalid array (ValidateFull()fails). The example from the issue still reproduces on pyarrow 26.0.0: twoint8-indexed chunks with 64 and 65 distinct values give indices[63, -128].What changes are included in this PR?
DictionaryHashKernelnow builds the unified dictionary withDictionaryUnifier::GetResultWithIndexType()using the input index type, instead ofGetResult(), which picks a possibly wider index type that the output never uses. If the unified dictionary does not fit, the kernel returns the sameInvaliderror that the unifier already produces in other places ("The unified dictionary requires a larger index type").Widening the output index type would change the declared output type of
unique(FirstType), so I went with an error here.Are these changes tested?
Yes. Added
TestHashKernel.DictionaryUniqueAndValueCountsIndexOverflow, which checks a case that just fits inint8(127 unified entries, result validates) and one that does not (129 entries,Invalid). The test fails without the fix.arrow-compute-vector-testpasses locally (macOS arm64, Debug).Are there any user-facing changes?
unique/value_countsnow return an error instead of an invalid dictionary array in this case.This PR contains a "Critical Fix". It fixes a bug that produced invalid data (out-of-range dictionary indices).
Was AI used for this PR?
In accordance to the AI generation guidelines, please disclose below whether and how AI was used in this PR.
PR code and description written by:
Reviewed before submission by:
pyarrow.compute.uniquegives garbage results with chunked dictionary arrays. #45010