From 787d0cdedcb422f1f1c89fb02173bb99ba8b4a54 Mon Sep 17 00:00:00 2001 From: mkzung <103102868+mkzung@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:59:13 +0500 Subject: [PATCH] GH-50695: [C++] Use the full unsigned range for unsigned dictionary index types An unsigned dictionary index previously widened on the signed thresholds because the dictionary builder reuses the signed AdaptiveIntBuilder: a uint8 index widened after 128 distinct values rather than 256. Follow-up to GH-37476 suggested by the review there. Give AdaptiveIntBuilder an opt-in unsigned mode: with use_unsigned_range the appended values are taken as non-negative, width detection and lane writes go through the existing unsigned helpers (DetectUIntWidth, DowncastUInts), and the reported type is unsigned. The dictionary builder passes its use_unsigned_index flag into the indices builder and no longer needs the MaybeUnsignedIndexType mapping, which is removed: the builder now reports the index type natively on every path (type, FinishInternal, FinishDelta, the NullType specialization). The index-type tests assert the new thresholds: 256 distinct values keep a uint8 index, with the indices round-tripping intact through the one-byte lanes, and the 257th widens to uint16; the signed thresholds are unchanged. --- cpp/src/arrow/array/array_dict_test.cc | 56 +++++++++-- cpp/src/arrow/array/builder_adaptive.cc | 125 +++++++++++++++++++++++- cpp/src/arrow/array/builder_adaptive.h | 16 ++- cpp/src/arrow/array/builder_dict.h | 45 +++------ cpp/src/arrow/builder.cc | 22 ----- python/pyarrow/tests/test_array.py | 17 ++-- 6 files changed, 209 insertions(+), 72 deletions(-) diff --git a/cpp/src/arrow/array/array_dict_test.cc b/cpp/src/arrow/array/array_dict_test.cc index 23335ebb008a..83e33c95ab23 100644 --- a/cpp/src/arrow/array/array_dict_test.cc +++ b/cpp/src/arrow/array/array_dict_test.cc @@ -1070,22 +1070,47 @@ TEST(TestDictionaryBuilderIndexType, PreservesRequestedIndexType) { } } -TEST(TestDictionaryBuilderIndexType, WidthAdaptsWhenUnsigned) { - // The width stays adaptive, as it does for signed indices. The underlying builder is - // signed, so it widens after 128 distinct values rather than the 256 a uint8 could - // hold, but the widened type stays unsigned rather than falling back to a signed type. +TEST(TestDictionaryBuilderIndexType, WidthAdaptsOnUnsignedThresholds) { + // An unsigned index uses the full unsigned range: 256 distinct values still fit + // a uint8 index, and the 257th widens it to uint16. auto dict_type = dictionary(uint8(), utf8()); ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); auto& builder = checked_cast&>(*boxed_builder); - for (int i = 0; i < 200; ++i) { + for (int i = 0; i < 256; ++i) { ASSERT_OK(builder.Append(std::to_string(i))); } - + AssertTypeEqual(*uint8(), + *checked_cast(*builder.type()).index_type()); + { + // Indices above 127 must round-trip through the one-byte lanes intact. + ASSERT_OK_AND_ASSIGN(auto full, builder.Finish()); + ASSERT_OK(full->ValidateFull()); + const auto& dict_array = checked_cast(*full); + const auto& indices = checked_cast(*dict_array.indices()); + ASSERT_EQ(indices.length(), 256); + ASSERT_EQ(indices.Value(255), 255); + ASSERT_EQ(dict_array.dictionary()->length(), 256); + } + for (int i = 0; i < 256; ++i) { + ASSERT_OK(builder.Append(std::to_string(i))); + } + ASSERT_OK(builder.Append("one more")); ASSERT_OK_AND_ASSIGN(auto result, builder.Finish()); ASSERT_OK(result->ValidateFull()); AssertTypeEqual(*uint16(), *checked_cast(*result->type()).index_type()); + + // The signed counterpart still widens on the signed threshold. + ASSERT_OK_AND_ASSIGN(auto signed_builder, MakeBuilder(dictionary(int8(), utf8()))); + auto& sbuilder = checked_cast&>(*signed_builder); + for (int i = 0; i < 200; ++i) { + ASSERT_OK(sbuilder.Append(std::to_string(i))); + } + ASSERT_OK_AND_ASSIGN(auto signed_result, sbuilder.Finish()); + AssertTypeEqual( + *int16(), + *checked_cast(*signed_result->type()).index_type()); } TEST(TestDictionaryBuilderIndexType, NullValueTypePreservesUnsignedIndexType) { @@ -1121,6 +1146,25 @@ TEST(TestDictionaryBuilderIndexType, FinishDeltaPreservesUnsignedIndexType) { AssertArraysEqual(*ArrayFromJSON(uint32(), "[0, 1]"), *result_indices); } +TEST(TestDictionaryBuilderIndexType, FinishDeltaUsesUnsignedRange) { + // The delta path shares the unsigned width logic: 200 distinct values keep a + // uint8 index and the values above 127 come through intact. + auto dict_type = dictionary(uint8(), utf8()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + + for (int i = 0; i < 200; ++i) { + ASSERT_OK(builder.Append(std::to_string(i))); + } + + std::shared_ptr result_indices, result_delta; + ASSERT_OK(builder.FinishDelta(&result_indices, &result_delta)); + ASSERT_OK(result_indices->ValidateFull()); + AssertTypeEqual(*uint8(), *result_indices->type()); + ASSERT_EQ(checked_cast(*result_indices).Value(199), 199); + ASSERT_EQ(result_delta->length(), 200); +} + TEST(TestDictionaryBuilderIndexType, SuppliedDictionaryPreservesUnsignedIndexType) { // The supplied-dictionary constructor starts the adaptive builder at its default width // rather than the requested one, so a requested uint32 reports uint8. The width is not diff --git a/cpp/src/arrow/array/builder_adaptive.cc b/cpp/src/arrow/array/builder_adaptive.cc index 3cd5a46321f4..204cdce06e85 100644 --- a/cpp/src/arrow/array/builder_adaptive.cc +++ b/cpp/src/arrow/array/builder_adaptive.cc @@ -108,6 +108,26 @@ std::shared_ptr AdaptiveUIntBuilder::type() const { std::shared_ptr AdaptiveIntBuilder::type() const { auto int_size = int_size_; + if (use_unsigned_range_) { + if (pending_pos_ != 0) { + const uint8_t* valid_bytes = pending_has_nulls_ ? pending_valid_ : nullptr; + int_size = + internal::DetectUIntWidth(pending_data_, valid_bytes, pending_pos_, int_size_); + } + switch (int_size) { + case 1: + return uint8(); + case 2: + return uint16(); + case 4: + return uint32(); + case 8: + return uint64(); + default: + DCHECK(false); + } + return nullptr; + } if (pending_pos_ != 0) { const uint8_t* valid_bytes = pending_has_nulls_ ? pending_valid_ : nullptr; int_size = internal::DetectIntWidth(reinterpret_cast(pending_data_), @@ -129,8 +149,9 @@ std::shared_ptr AdaptiveIntBuilder::type() const { } AdaptiveIntBuilder::AdaptiveIntBuilder(uint8_t start_int_size, MemoryPool* pool, - int64_t alignment) - : AdaptiveIntBuilderBase(start_int_size, pool, alignment) {} + int64_t alignment, bool use_unsigned_range) + : AdaptiveIntBuilderBase(start_int_size, pool, alignment), + use_unsigned_range_(use_unsigned_range) {} Status AdaptiveIntBuilder::FinishInternal(std::shared_ptr* out) { RETURN_NOT_OK(CommitPendingData()); @@ -158,8 +179,12 @@ Status AdaptiveIntBuilder::CommitPendingData() { } RETURN_NOT_OK(Reserve(pending_pos_)); const uint8_t* valid_bytes = pending_has_nulls_ ? pending_valid_ : nullptr; - RETURN_NOT_OK(AppendValuesInternal(reinterpret_cast(pending_data_), - pending_pos_, valid_bytes)); + if (use_unsigned_range_) { + RETURN_NOT_OK(AppendValuesUnsignedInternal(pending_data_, pending_pos_, valid_bytes)); + } else { + RETURN_NOT_OK(AppendValuesInternal(reinterpret_cast(pending_data_), + pending_pos_, valid_bytes)); + } pending_has_nulls_ = false; pending_pos_ = 0; return Status::OK(); @@ -223,6 +248,94 @@ Status AdaptiveIntBuilder::AppendValuesInternal(const int64_t* values, int64_t l return Status::OK(); } +Status AdaptiveIntBuilder::AppendValuesUnsignedInternal(const uint64_t* values, + int64_t length, + const uint8_t* valid_bytes) { + if (pending_pos_ > 0) { + // UnsafeAppendToBitmap expects length_ to be the pre-update value, satisfy it + DCHECK_EQ(length, pending_pos_) << "AppendValuesInternal called while data pending"; + length_ -= pending_pos_; + } + + while (length > 0) { + // See AdaptiveIntBuilder::AppendValuesInternal + const int64_t chunk_size = std::min(length, kAdaptiveIntChunkSize); + + uint8_t new_int_size; + new_int_size = internal::DetectUIntWidth(values, valid_bytes, chunk_size, int_size_); + + DCHECK_GE(new_int_size, int_size_); + if (new_int_size > int_size_) { + // This updates int_size_ + RETURN_NOT_OK(ExpandUIntSize(new_int_size)); + } + + switch (int_size_) { + case 1: + internal::DowncastUInts(values, reinterpret_cast(raw_data_) + length_, + chunk_size); + break; + case 2: + internal::DowncastUInts(values, reinterpret_cast(raw_data_) + length_, + chunk_size); + break; + case 4: + internal::DowncastUInts(values, reinterpret_cast(raw_data_) + length_, + chunk_size); + break; + case 8: + internal::DowncastUInts(values, reinterpret_cast(raw_data_) + length_, + chunk_size); + break; + default: + DCHECK(false); + } + + // UnsafeAppendToBitmap increments length_ by chunk_size + ArrayBuilder::UnsafeAppendToBitmap(valid_bytes, chunk_size); + values += chunk_size; + if (valid_bytes != nullptr) { + valid_bytes += chunk_size; + } + length -= chunk_size; + } + + return Status::OK(); +} + +template +Status AdaptiveIntBuilder::ExpandUIntSizeN() { + switch (int_size_) { + case 1: + return ExpandIntSizeInternal(); + case 2: + return ExpandIntSizeInternal(); + case 4: + return ExpandIntSizeInternal(); + case 8: + return ExpandIntSizeInternal(); + default: + DCHECK(false); + } + return Status::OK(); +} + +Status AdaptiveIntBuilder::ExpandUIntSize(uint8_t new_int_size) { + switch (new_int_size) { + case 1: + return ExpandUIntSizeN(); + case 2: + return ExpandUIntSizeN(); + case 4: + return ExpandUIntSizeN(); + case 8: + return ExpandUIntSizeN(); + default: + DCHECK(false); + } + return Status::OK(); +} + Status AdaptiveUIntBuilder::CommitPendingData() { if (pending_pos_ == 0) { return Status::OK(); @@ -240,6 +353,10 @@ Status AdaptiveIntBuilder::AppendValues(const int64_t* values, int64_t length, RETURN_NOT_OK(CommitPendingData()); RETURN_NOT_OK(Reserve(length)); + if (use_unsigned_range_) { + return AppendValuesUnsignedInternal(reinterpret_cast(values), length, + valid_bytes); + } return AppendValuesInternal(values, length, valid_bytes); } diff --git a/cpp/src/arrow/array/builder_adaptive.h b/cpp/src/arrow/array/builder_adaptive.h index 0cea571be3e3..a703ee45c882 100644 --- a/cpp/src/arrow/array/builder_adaptive.h +++ b/cpp/src/arrow/array/builder_adaptive.h @@ -173,9 +173,16 @@ class ARROW_EXPORT AdaptiveUIntBuilder : public internal::AdaptiveIntBuilderBase class ARROW_EXPORT AdaptiveIntBuilder : public internal::AdaptiveIntBuilderBase { public: + /// \brief Create a builder starting at `start_int_size` bytes wide. + /// + /// With `use_unsigned_range` the appended values are taken to be non-negative + /// and the width adapts on unsigned thresholds, so one byte holds 256 + /// distinct values rather than 128, and the reported type is unsigned. + /// The dictionary builder uses this for unsigned index types. explicit AdaptiveIntBuilder(uint8_t start_int_size, MemoryPool* pool = default_memory_pool(), - int64_t alignment = kDefaultBufferAlignment); + int64_t alignment = kDefaultBufferAlignment, + bool use_unsigned_range = false); explicit AdaptiveIntBuilder(MemoryPool* pool = default_memory_pool(), int64_t alignment = kDefaultBufferAlignment) @@ -205,9 +212,16 @@ class ARROW_EXPORT AdaptiveIntBuilder : public internal::AdaptiveIntBuilderBase Status AppendValuesInternal(const int64_t* values, int64_t length, const uint8_t* valid_bytes); + Status AppendValuesUnsignedInternal(const uint64_t* values, int64_t length, + const uint8_t* valid_bytes); + Status ExpandUIntSize(uint8_t new_int_size); template Status ExpandIntSizeN(); + template + Status ExpandUIntSizeN(); + + bool use_unsigned_range_ = false; }; /// @} diff --git a/cpp/src/arrow/array/builder_dict.h b/cpp/src/arrow/array/builder_dict.h index 4dd4e6ea44fd..bc777ca675e3 100644 --- a/cpp/src/arrow/array/builder_dict.h +++ b/cpp/src/arrow/array/builder_dict.h @@ -134,12 +134,6 @@ class ARROW_EXPORT DictionaryMemoTable { namespace internal { -/// \brief Return the unsigned integer type of the same width when an unsigned -/// dictionary index type was requested. Defined in builder.cc; see there for why -/// this is value-preserving and why the width widens on the signed threshold. -ARROW_EXPORT std::shared_ptr MaybeUnsignedIndexType( - const std::shared_ptr& index_type, bool use_unsigned_index); - /// \brief Array builder for created encoded DictionaryArray from /// dense array /// @@ -166,10 +160,9 @@ class DictionaryBuilderBase : public ArrayBuilder { memo_table_(new internal::DictionaryMemoTable(pool, value_type)), delta_offset_(0), byte_width_(-1), - indices_builder_(start_int_size, pool, alignment), + indices_builder_(start_int_size, pool, alignment, use_unsigned_index), value_type_(value_type), - ordered_(ordered), - use_unsigned_index_(use_unsigned_index) {} + ordered_(ordered) {} template explicit DictionaryBuilderBase( @@ -213,10 +206,9 @@ class DictionaryBuilderBase : public ArrayBuilder { 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), + indices_builder_(start_int_size, pool, alignment, use_unsigned_index), value_type_(value_type), - ordered_(ordered), - use_unsigned_index_(use_unsigned_index) {} + ordered_(ordered) {} template explicit DictionaryBuilderBase( @@ -259,10 +251,9 @@ class DictionaryBuilderBase : public ArrayBuilder { memo_table_(new internal::DictionaryMemoTable(pool, dictionary)), delta_offset_(0), byte_width_(-1), - indices_builder_(pool, alignment), + indices_builder_(sizeof(uint8_t), pool, alignment, use_unsigned_index), value_type_(dictionary->type()), - ordered_(ordered), - use_unsigned_index_(use_unsigned_index) {} + ordered_(ordered) {} ~DictionaryBuilderBase() override = default; @@ -497,7 +488,6 @@ 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, use_unsigned_index_); *out_indices = MakeArray(indices_data); *out_delta = MakeArray(delta_data); return Status::OK(); @@ -510,9 +500,7 @@ class DictionaryBuilderBase : public ArrayBuilder { Status Finish(std::shared_ptr* out) { return FinishTyped(out); } std::shared_ptr type() const override { - return ::arrow::dictionary( - MaybeUnsignedIndexType(indices_builder_.type(), use_unsigned_index_), value_type_, - ordered_); + return ::arrow::dictionary(indices_builder_.type(), value_type_, ordered_); } protected: @@ -584,7 +572,6 @@ class DictionaryBuilderBase : public ArrayBuilder { BuilderType indices_builder_; std::shared_ptr value_type_; bool ordered_ = false; - bool use_unsigned_index_ = false; }; template @@ -599,9 +586,8 @@ class DictionaryBuilderBase : public ArrayBuilder { int64_t alignment = kDefaultBufferAlignment, bool ordered = false, bool use_unsigned_index = false) : ArrayBuilder(pool, alignment), - indices_builder_(start_int_size, pool, alignment), - ordered_(ordered), - use_unsigned_index_(use_unsigned_index) {} + indices_builder_(start_int_size, pool, alignment, use_unsigned_index), + ordered_(ordered) {} explicit DictionaryBuilderBase(const std::shared_ptr& value_type, MemoryPool* pool = default_memory_pool(), @@ -642,9 +628,8 @@ class DictionaryBuilderBase : public ArrayBuilder { int64_t alignment = kDefaultBufferAlignment, bool ordered = false, bool use_unsigned_index = false) : ArrayBuilder(pool, alignment), - indices_builder_(pool, alignment), - ordered_(ordered), - use_unsigned_index_(use_unsigned_index) {} + indices_builder_(sizeof(uint8_t), pool, alignment, use_unsigned_index), + ordered_(ordered) {} /// \brief Append a scalar null value Status AppendNull() final { @@ -696,8 +681,7 @@ class DictionaryBuilderBase : public ArrayBuilder { Status FinishInternal(std::shared_ptr* out) override { ARROW_RETURN_NOT_OK(indices_builder_.FinishInternal(out)); - (*out)->type = dictionary(MaybeUnsignedIndexType((*out)->type, use_unsigned_index_), - null(), ordered_); + (*out)->type = dictionary((*out)->type, null(), ordered_); (*out)->dictionary = NullArray(0).data(); return Status::OK(); } @@ -709,15 +693,12 @@ class DictionaryBuilderBase : public ArrayBuilder { Status Finish(std::shared_ptr* out) { return FinishTyped(out); } std::shared_ptr type() const override { - return ::arrow::dictionary( - MaybeUnsignedIndexType(indices_builder_.type(), use_unsigned_index_), null(), - ordered_); + return ::arrow::dictionary(indices_builder_.type(), null(), ordered_); } protected: BuilderType indices_builder_; bool ordered_ = false; - bool use_unsigned_index_ = false; }; } // namespace internal diff --git a/cpp/src/arrow/builder.cc b/cpp/src/arrow/builder.cc index 1d0a42099b3e..5bfbbd98f533 100644 --- a/cpp/src/arrow/builder.cc +++ b/cpp/src/arrow/builder.cc @@ -27,7 +27,6 @@ #include "arrow/util/checked_cast.h" #include "arrow/util/hashing.h" #include "arrow/util/logging_internal.h" -#include "arrow/util/unreachable.h" #include "arrow/visit_type_inline.h" namespace arrow { @@ -50,27 +49,6 @@ namespace internal { /// width stays adaptive, as it is for signed index types, and it widens on the signed /// threshold: a uint8 index widens after 128 distinct values rather than the 256 a real /// uint8 could hold, so the extra bit does not delay widening. -std::shared_ptr MaybeUnsignedIndexType( - const std::shared_ptr& index_type, bool use_unsigned_index) { - if (!use_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: - // The adaptive index builder only ever produces signed int8/16/32/64, so no - // other type reaches this point when an unsigned index was requested. - Unreachable("MaybeUnsignedIndexType: adaptive dictionary index type is not signed"); - } -} - } // namespace internal // Generic int builder that delegates to the builder for a specific diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index bc4e521dca82..6670560db70c 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -4602,14 +4602,17 @@ def test_dictionary_array_preserves_index_type(index_type): assert chunked.type == dict_type -@pytest.mark.parametrize("start_type, widened_type", [(pa.int8(), pa.int16()), - (pa.uint8(), pa.uint16())]) -def test_dictionary_array_index_width_adapts(start_type, widened_type): - # The index width adapts to the number of distinct values, as it does for signed - # indices; only the signedness of the requested type is preserved. - values = [str(i) for i in range(200)] +@pytest.mark.parametrize( + "start_type, distinct, expected_type", + [(pa.int8(), 200, pa.int16()), # signed widens past 127 + (pa.uint8(), 200, pa.uint8()), # unsigned uses the full byte + (pa.uint8(), 257, pa.uint16())]) # and widens past 255 +def test_dictionary_array_index_width_adapts(start_type, distinct, expected_type): + # The index width adapts to the number of distinct values, on the thresholds + # of the requested signedness. + values = [str(i) for i in range(distinct)] arr = pa.array(values, type=pa.dictionary(start_type, pa.string())) - assert arr.type == pa.dictionary(widened_type, pa.string()) + assert arr.type == pa.dictionary(expected_type, pa.string()) assert arr.to_pylist() == values