From 25e49e7cdb9d0ad9d37c0ca7d30083a9317dc8c0 Mon Sep 17 00:00:00 2001 From: mkzung <103102868+mkzung@users.noreply.github.com> Date: Sat, 11 Jul 2026 18:17:14 +0300 Subject: [PATCH] GH-37476: [C++][Python] Preserve unsigned dictionary index types when building from values --- cpp/src/arrow/array/array_dict_test.cc | 171 +++++++++++++++++++++++++ cpp/src/arrow/array/builder_dict.h | 69 ++++++++-- cpp/src/arrow/builder.cc | 8 +- python/pyarrow/tests/test_array.py | 45 +++++++ 4 files changed, 277 insertions(+), 16 deletions(-) diff --git a/cpp/src/arrow/array/array_dict_test.cc b/cpp/src/arrow/array/array_dict_test.cc index 5f4335bcbc92..8dd76ddf8541 100644 --- a/cpp/src/arrow/array/array_dict_test.cc +++ b/cpp/src/arrow/array/array_dict_test.cc @@ -1040,6 +1040,177 @@ TYPED_TEST(TestDictionaryBuilderIndexByteWidth, MakeBuilder) { AssertIndexByteWidth(); } +// ---------------------------------------------------------------------- +// GH-37476: the requested index type's signedness must be preserved + +TEST(TestDictionaryBuilderIndexType, TypePreservesUnsignedIndexType) { + for (auto index_type : {uint8(), uint16(), uint32(), uint64()}) { + ARROW_SCOPED_TRACE("index_type = ", index_type->ToString()); + auto dict_type = dictionary(index_type, utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + + auto builder_type = boxed_builder->type(); + const auto& dict_builder_type = checked_cast(*builder_type); + AssertTypeEqual(*index_type, *dict_builder_type.index_type()); + } +} + +TEST(TestDictionaryBuilderIndexType, TypePreservesSignedIndexType) { + for (auto index_type : {int8(), int16(), int32(), int64()}) { + ARROW_SCOPED_TRACE("index_type = ", index_type->ToString()); + auto dict_type = dictionary(index_type, utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + + auto builder_type = boxed_builder->type(); + const auto& dict_builder_type = checked_cast(*builder_type); + AssertTypeEqual(*index_type, *dict_builder_type.index_type()); + } +} + +TEST(TestDictionaryBuilderIndexType, FinishPreservesUnsignedIndexType) { + for (auto index_type : {uint8(), uint16(), uint32(), uint64()}) { + ARROW_SCOPED_TRACE("index_type = ", index_type->ToString()); + auto dict_type = dictionary(index_type, utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + ASSERT_OK(builder.Append("a")); + ASSERT_OK(builder.Append("b")); + ASSERT_OK(builder.AppendNull()); + ASSERT_OK(builder.Append("a")); + + std::shared_ptr result; + ASSERT_OK(builder.Finish(&result)); + ASSERT_OK(result->ValidateFull()); + + auto ex_dict = ArrayFromJSON(utf8(), R"(["a", "b"])"); + auto ex_indices = ArrayFromJSON(index_type, "[0, 1, null, 0]"); + DictionaryArray expected(dict_type, ex_indices, ex_dict); + AssertTypeEqual(*dict_type, *result->type()); + AssertArraysEqual(expected, *result); + } +} + +TEST(TestDictionaryBuilderIndexType, UnsignedIndexWidthStillAdapts) { + // The index width remains adaptive, exactly as it is for signed index types: + // only the signedness of the requested type is preserved. + auto dict_type = dictionary(uint8(), utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + // More distinct values than a uint8 index start width accommodates. + for (int i = 0; i < 200; ++i) { + ASSERT_OK(builder.Append(std::to_string(i))); + } + + std::shared_ptr result; + ASSERT_OK(builder.Finish(&result)); + ASSERT_OK(result->ValidateFull()); + + const auto& result_type = checked_cast(*result->type()); + // Widened, but still unsigned rather than falling back to a signed type. + AssertTypeEqual(*uint16(), *result_type.index_type()); +} + +TEST(TestDictionaryBuilderIndexType, NullValueTypeUnsignedIndexType) { + auto dict_type = dictionary(uint32(), null()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + ASSERT_OK(builder.AppendNull()); + + std::shared_ptr result; + ASSERT_OK(builder.Finish(&result)); + ASSERT_OK(result->ValidateFull()); + + const auto& result_type = checked_cast(*result->type()); + AssertTypeEqual(*uint32(), *result_type.index_type()); +} + +TEST(TestDictionaryBuilderIndexType, FinishDeltaPreservesUnsignedIndexType) { + auto dict_type = dictionary(uint32(), utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + ASSERT_OK(builder.Append("a")); + ASSERT_OK(builder.Append("b")); + + std::shared_ptr result_indices, result_delta; + ASSERT_OK(builder.FinishDelta(&result_indices, &result_delta)); + ASSERT_OK(result_indices->ValidateFull()); + + // The delta indices must carry the requested index type too, not the signed + // type that the adaptive indices builder produces internally. + AssertTypeEqual(*uint32(), *result_indices->type()); + AssertArraysEqual(*ArrayFromJSON(uint32(), "[0, 1]"), *result_indices); +} + +TEST(TestDictionaryBuilderIndexType, SuppliedDictionaryPreservesUnsignedIndexType) { + // The supplied-dictionary constructor starts the adaptive indices builder at its + // default width rather than at the requested one, so the requested *width* is not + // honoured on this path. That predates this change: a requested int32 reports int8 + // too. What must hold is that the *signedness* survives, i.e. uint8 and not int8. + auto dict_values = ArrayFromJSON(utf8(), R"(["a", "b"])"); + + ASSERT_OK_AND_ASSIGN(auto unsigned_builder, + MakeDictionaryBuilder(dictionary(uint32(), utf8()), dict_values)); + auto unsigned_type = unsigned_builder->type(); + AssertTypeEqual(*uint8(), + *checked_cast(*unsigned_type).index_type()); + + ASSERT_OK_AND_ASSIGN(auto signed_builder, + MakeDictionaryBuilder(dictionary(int32(), utf8()), dict_values)); + auto signed_type = signed_builder->type(); + AssertTypeEqual(*int8(), + *checked_cast(*signed_type).index_type()); +} + +TEST(TestDictionaryBuilderIndexType, FixedSizeBinaryUnsignedIndexType) { + auto dict_type = dictionary(uint16(), fixed_size_binary(2)); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + ASSERT_OK(builder.Append("ab")); + ASSERT_OK(builder.Append("cd")); + ASSERT_OK(builder.Append("ab")); + + std::shared_ptr result; + ASSERT_OK(builder.Finish(&result)); + ASSERT_OK(result->ValidateFull()); + + const auto& result_type = checked_cast(*result->type()); + AssertTypeEqual(*uint16(), *result_type.index_type()); +} + +TEST(TestDictionaryBuilderIndexType, OrderedAndUnsignedIndexType) { + // The ordered flag (GH-49689) and the index type must both survive. + auto dict_type = dictionary(uint32(), utf8(), /*ordered=*/true); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + ASSERT_OK(builder.Append("a")); + ASSERT_OK(builder.Append("b")); + + std::shared_ptr result; + ASSERT_OK(builder.Finish(&result)); + ASSERT_OK(result->ValidateFull()); + + const auto& result_type = checked_cast(*result->type()); + AssertTypeEqual(*uint32(), *result_type.index_type()); + ASSERT_TRUE(result_type.ordered()); +} + +TEST(TestDictionaryBuilderIndexType, ExactIndexTypeStillHonoursUnsigned) { + // The exact-index builder already honoured unsigned index types; guard it. + auto dict_type = dictionary(uint16(), utf8()); + std::unique_ptr boxed_builder; + ASSERT_OK(MakeBuilderExactIndex(default_memory_pool(), dict_type, &boxed_builder)); + + auto builder_type = boxed_builder->type(); + const auto& dict_builder_type = checked_cast(*builder_type); + AssertTypeEqual(*uint16(), *dict_builder_type.index_type()); +} + // ---------------------------------------------------------------------- // DictionaryArray tests diff --git a/cpp/src/arrow/array/builder_dict.h b/cpp/src/arrow/array/builder_dict.h index 269cdeee646d..dd178eef826e 100644 --- a/cpp/src/arrow/array/builder_dict.h +++ b/cpp/src/arrow/array/builder_dict.h @@ -134,6 +134,33 @@ class ARROW_EXPORT DictionaryMemoTable { namespace internal { +/// \brief Return the unsigned integer type of the same width when an unsigned +/// dictionary index type was requested. +/// +/// The adaptive indices builder only ever produces signed integer types. Dictionary +/// indices are non-negative, so the signed and unsigned integer types of a given +/// width have identical memory layout and reporting one as the other is +/// value-preserving (GH-37476). The width itself stays adaptive, exactly as it is +/// for signed index types; only the signedness of the requested type is preserved. +inline std::shared_ptr MaybeUnsignedIndexType( + const std::shared_ptr& index_type, bool unsigned_index) { + if (!unsigned_index) { + return index_type; + } + switch (index_type->id()) { + case Type::INT8: + return ::arrow::uint8(); + case Type::INT16: + return ::arrow::uint16(); + case Type::INT32: + return ::arrow::uint32(); + case Type::INT64: + return ::arrow::uint64(); + default: + return index_type; + } +} + /// \brief Array builder for created encoded DictionaryArray from /// dense array /// @@ -154,14 +181,16 @@ class DictionaryBuilderBase : public ArrayBuilder { const std::shared_ptr&> value_type, MemoryPool* pool = default_memory_pool(), - int64_t alignment = kDefaultBufferAlignment, bool ordered = false) + int64_t alignment = kDefaultBufferAlignment, bool ordered = false, + bool unsigned_index = false) : ArrayBuilder(pool, alignment), memo_table_(new internal::DictionaryMemoTable(pool, value_type)), delta_offset_(0), byte_width_(-1), indices_builder_(start_int_size, pool, alignment), value_type_(value_type), - ordered_(ordered) {} + ordered_(ordered), + unsigned_index_(unsigned_index) {} template explicit DictionaryBuilderBase( @@ -199,14 +228,16 @@ class DictionaryBuilderBase : public ArrayBuilder { const std::shared_ptr&> value_type, MemoryPool* pool = default_memory_pool(), - int64_t alignment = kDefaultBufferAlignment, bool ordered = false) + int64_t alignment = kDefaultBufferAlignment, bool ordered = false, + bool unsigned_index = false) : ArrayBuilder(pool, alignment), memo_table_(new internal::DictionaryMemoTable(pool, value_type)), delta_offset_(0), byte_width_(static_cast(*value_type).byte_width()), indices_builder_(start_int_size, pool, alignment), value_type_(value_type), - ordered_(ordered) {} + ordered_(ordered), + unsigned_index_(unsigned_index) {} template explicit DictionaryBuilderBase( @@ -244,14 +275,15 @@ class DictionaryBuilderBase : public ArrayBuilder { explicit DictionaryBuilderBase(const std::shared_ptr& dictionary, MemoryPool* pool = default_memory_pool(), int64_t alignment = kDefaultBufferAlignment, - bool ordered = false) + bool ordered = false, bool unsigned_index = false) : ArrayBuilder(pool, alignment), memo_table_(new internal::DictionaryMemoTable(pool, dictionary)), delta_offset_(0), byte_width_(-1), indices_builder_(pool, alignment), value_type_(dictionary->type()), - ordered_(ordered) {} + ordered_(ordered), + unsigned_index_(unsigned_index) {} ~DictionaryBuilderBase() override = default; @@ -486,6 +518,7 @@ class DictionaryBuilderBase : public ArrayBuilder { std::shared_ptr indices_data; std::shared_ptr delta_data; ARROW_RETURN_NOT_OK(FinishWithDictOffset(delta_offset_, &indices_data, &delta_data)); + indices_data->type = MaybeUnsignedIndexType(indices_data->type, unsigned_index_); *out_indices = MakeArray(indices_data); *out_delta = MakeArray(delta_data); return Status::OK(); @@ -498,7 +531,9 @@ class DictionaryBuilderBase : public ArrayBuilder { Status Finish(std::shared_ptr* out) { return FinishTyped(out); } std::shared_ptr type() const override { - return ::arrow::dictionary(indices_builder_.type(), value_type_, ordered_); + return ::arrow::dictionary( + MaybeUnsignedIndexType(indices_builder_.type(), unsigned_index_), value_type_, + ordered_); } protected: @@ -570,6 +605,7 @@ class DictionaryBuilderBase : public ArrayBuilder { BuilderType indices_builder_; std::shared_ptr value_type_; bool ordered_ = false; + bool unsigned_index_ = false; }; template @@ -581,10 +617,12 @@ class DictionaryBuilderBase : public ArrayBuilder { start_int_size, const std::shared_ptr& value_type, MemoryPool* pool = default_memory_pool(), - int64_t alignment = kDefaultBufferAlignment, bool ordered = false) + int64_t alignment = kDefaultBufferAlignment, bool ordered = false, + bool unsigned_index = false) : ArrayBuilder(pool, alignment), indices_builder_(start_int_size, pool, alignment), - ordered_(ordered) {} + ordered_(ordered), + unsigned_index_(unsigned_index) {} explicit DictionaryBuilderBase(const std::shared_ptr& value_type, MemoryPool* pool = default_memory_pool(), @@ -623,10 +661,11 @@ class DictionaryBuilderBase : public ArrayBuilder { explicit DictionaryBuilderBase(const std::shared_ptr& dictionary, MemoryPool* pool = default_memory_pool(), int64_t alignment = kDefaultBufferAlignment, - bool ordered = false) + bool ordered = false, bool unsigned_index = false) : ArrayBuilder(pool, alignment), indices_builder_(pool, alignment), - ordered_(ordered) {} + ordered_(ordered), + unsigned_index_(unsigned_index) {} /// \brief Append a scalar null value Status AppendNull() final { @@ -678,7 +717,8 @@ class DictionaryBuilderBase : public ArrayBuilder { Status FinishInternal(std::shared_ptr* out) override { ARROW_RETURN_NOT_OK(indices_builder_.FinishInternal(out)); - (*out)->type = dictionary((*out)->type, null(), ordered_); + (*out)->type = dictionary(MaybeUnsignedIndexType((*out)->type, unsigned_index_), + null(), ordered_); (*out)->dictionary = NullArray(0).data(); return Status::OK(); } @@ -690,12 +730,15 @@ class DictionaryBuilderBase : public ArrayBuilder { Status Finish(std::shared_ptr* out) { return FinishTyped(out); } std::shared_ptr type() const override { - return ::arrow::dictionary(indices_builder_.type(), null(), ordered_); + return ::arrow::dictionary( + MaybeUnsignedIndexType(indices_builder_.type(), unsigned_index_), null(), + ordered_); } protected: BuilderType indices_builder_; bool ordered_ = false; + bool unsigned_index_ = false; }; } // namespace internal diff --git a/cpp/src/arrow/builder.cc b/cpp/src/arrow/builder.cc index 2190f1ce04a9..d4269f57d959 100644 --- a/cpp/src/arrow/builder.cc +++ b/cpp/src/arrow/builder.cc @@ -170,9 +170,10 @@ struct DictionaryBuilderCase { using AdaptiveBuilderType = DictionaryBuilder; using ExactBuilderType = internal::DictionaryBuilderBase; + const bool unsigned_index = is_unsigned_integer(index_type->id()); if (dictionary != nullptr) { - out->reset( - new AdaptiveBuilderType(dictionary, pool, kDefaultBufferAlignment, ordered)); + out->reset(new AdaptiveBuilderType(dictionary, pool, kDefaultBufferAlignment, + ordered, unsigned_index)); } else if (exact_index_type) { if (!is_integer(index_type->id())) { return Status::TypeError("MakeBuilder: invalid index type ", *index_type); @@ -182,7 +183,8 @@ struct DictionaryBuilderCase { } else { auto start_int_size = index_type->byte_width(); out->reset(new AdaptiveBuilderType(start_int_size, value_type, pool, - kDefaultBufferAlignment, ordered)); + kDefaultBufferAlignment, ordered, + unsigned_index)); } return Status::OK(); } diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index adc3e097b54a..fe3f6c39cec0 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -4527,3 +4527,48 @@ def test_dunders_checked_overflow(): arr ** pa.scalar(2, type=pa.int8()) with pytest.raises(pa.ArrowInvalid, match=error_match): arr / (-arr) + + +@pytest.mark.parametrize("index_type", [pa.uint8(), pa.uint16(), + pa.uint32(), pa.uint64()]) +def test_dictionary_array_preserves_unsigned_index_type(index_type): + # GH-37476: an unsigned dictionary index type used to be silently replaced + # by the signed integer type of the same width. + dict_type = pa.dictionary(index_type, pa.string()) + + arr = pa.array(["a", "b", None, "a"], type=dict_type) + assert arr.type == dict_type + assert arr.to_pylist() == ["a", "b", None, "a"] + arr.validate(full=True) + + chunked = pa.chunked_array([["a", "b", "a"]], dict_type) + assert chunked.type == dict_type + + +@pytest.mark.parametrize("index_type", [pa.int8(), pa.int16(), + pa.int32(), pa.int64()]) +def test_dictionary_array_signed_index_type_unchanged(index_type): + dict_type = pa.dictionary(index_type, pa.string()) + arr = pa.array(["a", "b", "a"], type=dict_type) + assert arr.type == dict_type + + +def test_dictionary_array_index_width_still_adapts(): + # The index width remains adaptive, as it is for signed index types; only + # the signedness of the requested index type is preserved. + values = [str(i) for i in range(200)] + arr = pa.array(values, type=pa.dictionary(pa.uint8(), pa.string())) + assert arr.type == pa.dictionary(pa.uint16(), pa.string()) + assert arr.to_pylist() == values + + +@pytest.mark.pandas +def test_dictionary_uint64_index_to_pandas_not_supported(): + # GH-37476: uint64 dictionary indices are now preserved instead of being + # silently downgraded to int64. Arrow deliberately does not support converting + # them to pandas (pandas categorical codes are signed), so that pre-existing + # limitation is now reachable, where the silent downgrade used to mask it. + # uint8/uint16/uint32 dictionary indices convert to pandas as before. + arr = pa.array(["a", "b"], type=pa.dictionary(pa.uint64(), pa.string())) + with pytest.raises(pa.ArrowTypeError, match="UInt64 dictionary indices"): + arr.to_pandas()