diff --git a/cpp/src/arrow/compute/kernels/vector_hash.cc b/cpp/src/arrow/compute/kernels/vector_hash.cc index 9a29dcf333f..57061cf413f 100644 --- a/cpp/src/arrow/compute/kernels/vector_hash.cc +++ b/cpp/src/arrow/compute/kernels/vector_hash.cc @@ -445,9 +445,11 @@ class NullHashKernel : public HashKernel { class DictionaryHashKernel : public HashKernel { public: - explicit DictionaryHashKernel(std::unique_ptr indices_kernel, - std::shared_ptr dictionary_value_type) + DictionaryHashKernel(std::unique_ptr indices_kernel, + std::shared_ptr index_type, + std::shared_ptr dictionary_value_type) : indices_kernel_(std::move(indices_kernel)), + index_type_(std::move(index_type)), dictionary_value_type_(std::move(dictionary_value_type)) {} Status Reset() override { return indices_kernel_->Reset(); } @@ -513,15 +515,25 @@ class DictionaryHashKernel : public HashKernel { return first_dictionary_; } - auto out_dict_type = first_dictionary_->type(); + std::shared_ptr out_dict_type; std::shared_ptr out_dict; RETURN_NOT_OK(dictionary_unifier_->GetResult(&out_dict_type, &out_dict)); + // The output keeps the input index type, so the highest index into the unified + // dictionary (its length minus one) must fit in it + const Int64Scalar max_index(out_dict->length() - 1); + if (out_dict->length() > 0 && + !::arrow::internal::IntegersCanFit(max_index, *index_type_).ok()) { + return Status::Invalid( + "These dictionaries cannot be combined. The unified dictionary requires a " + "larger index type."); + } return out_dict; } private: std::unique_ptr indices_kernel_; std::shared_ptr first_dictionary_; + std::shared_ptr index_type_; std::shared_ptr dictionary_value_type_; std::unique_ptr dictionary_unifier_; }; @@ -597,7 +609,7 @@ Result> DictionaryHashInit(KernelContext* ctx, ARROW_ASSIGN_OR_RAISE(auto indices_hasher, GetHashInit(dict_type.index_type()->id())(ctx, args)); return std::make_unique( - checked_pointer_cast(std::move(indices_hasher)), + checked_pointer_cast(std::move(indices_hasher)), dict_type.index_type(), dict_type.value_type()); } diff --git a/cpp/src/arrow/compute/kernels/vector_hash_test.cc b/cpp/src/arrow/compute/kernels/vector_hash_test.cc index b11e0a722a5..8bd907679f4 100644 --- a/cpp/src/arrow/compute/kernels/vector_hash_test.cc +++ b/cpp/src/arrow/compute/kernels/vector_hash_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -764,6 +765,35 @@ TEST_F(TestHashKernel, DictionaryUniqueAndValueCounts) { } } +TEST_F(TestHashKernel, DictionaryUniqueAndValueCountsIndexOverflow) { + // GH-45010: unifying the chunk dictionaries can produce more entries than the + // index type can address + auto dict_ty = dictionary(int8(), int64()); + auto make_chunk = [&](int64_t first_value, int64_t dict_length) { + std::vector values(dict_length); + std::iota(values.begin(), values.end(), first_value); + std::shared_ptr dict; + ArrayFromVector(values, &dict); + auto indices = ArrayFromJSON(int8(), "[" + std::to_string(dict_length - 1) + "]"); + return std::make_shared(dict_ty, indices, dict); + }; + + // 64 + 64 = 128 unified entries still fit in int8: the highest index is 127 + auto fits = *ChunkedArray::Make({make_chunk(0, 64), make_chunk(100, 64)}, dict_ty); + ASSERT_OK_AND_ASSIGN(auto uniques, Unique(fits)); + ASSERT_OK(uniques->ValidateFull()); + ASSERT_EQ(checked_cast(*uniques).dictionary()->length(), 128); + ASSERT_OK_AND_ASSIGN(auto counts, ValueCounts(fits)); + ASSERT_OK(counts->ValidateFull()); + + // 64 + 65 = 129 unified entries do not + auto overflows = *ChunkedArray::Make({make_chunk(0, 64), make_chunk(100, 65)}, dict_ty); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("larger index type"), + Unique(overflows)); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("larger index type"), + ValueCounts(overflows)); +} + /* TODO(ARROW-4124): Determine if we want to do something that is reproducible with * floats. TEST_F(TestHashKernel, ValueCountsFloat) {