diff --git a/cpp/src/arrow/array/array_dict_test.cc b/cpp/src/arrow/array/array_dict_test.cc index 5f4335bcbc92..23335ebb008a 100644 --- a/cpp/src/arrow/array/array_dict_test.cc +++ b/cpp/src/arrow/array/array_dict_test.cc @@ -1040,6 +1040,123 @@ TYPED_TEST(TestDictionaryBuilderIndexByteWidth, MakeBuilder) { AssertIndexByteWidth(); } +// ---------------------------------------------------------------------- +// GH-37476: the requested dictionary index type's signedness must be preserved. + +TEST(TestDictionaryBuilderIndexType, PreservesRequestedIndexType) { + // Both the builder's reported type and the finished array must carry the requested + // index type, for signed and unsigned widths alike. + for (auto index_type : + {int8(), int16(), int32(), int64(), uint8(), uint16(), uint32(), uint64()}) { + ARROW_SCOPED_TRACE("index_type = ", index_type->ToString()); + auto dict_type = dictionary(index_type, utf8()); + ASSERT_OK_AND_ASSIGN(auto builder, MakeBuilder(dict_type)); + AssertTypeEqual(*index_type, + *checked_cast(*builder->type()).index_type()); + + auto& dict_builder = checked_cast&>(*builder); + ASSERT_OK(dict_builder.Append("a")); + ASSERT_OK(dict_builder.Append("b")); + ASSERT_OK(dict_builder.AppendNull()); + ASSERT_OK(dict_builder.Append("a")); + ASSERT_OK_AND_ASSIGN(auto result, dict_builder.Finish()); + 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, 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. + 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))); + } + + ASSERT_OK_AND_ASSIGN(auto result, builder.Finish()); + ASSERT_OK(result->ValidateFull()); + AssertTypeEqual(*uint16(), + *checked_cast(*result->type()).index_type()); +} + +TEST(TestDictionaryBuilderIndexType, NullValueTypePreservesUnsignedIndexType) { + // The NullType value builder is a separate specialization with its own type() and + // FinishInternal, so it needs its own guard. + auto dict_type = dictionary(uint32(), null()); + ASSERT_OK_AND_ASSIGN(auto boxed_builder, MakeBuilder(dict_type)); + auto& builder = checked_cast&>(*boxed_builder); + AssertTypeEqual(*uint32(), + *checked_cast(*builder.type()).index_type()); + + ASSERT_OK(builder.AppendNull()); + ASSERT_OK_AND_ASSIGN(auto result, builder.Finish()); + ASSERT_OK(result->ValidateFull()); + AssertTypeEqual(*uint32(), + *checked_cast(*result->type()).index_type()); +} + +TEST(TestDictionaryBuilderIndexType, FinishDeltaPreservesUnsignedIndexType) { + // FinishDelta is a distinct path from Finish and must carry the requested index type + // too, not the signed type the adaptive builder produces internally. + 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()); + AssertTypeEqual(*uint32(), *result_indices->type()); + AssertArraysEqual(*ArrayFromJSON(uint32(), "[0, 1]"), *result_indices); +} + +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 + // honoured on this path (a signed request behaves the same way), but the signedness + // must survive regardless: uint8, not int8. + auto dict_values = ArrayFromJSON(utf8(), R"(["a", "b"])"); + ASSERT_OK_AND_ASSIGN(auto builder, + MakeDictionaryBuilder(dictionary(uint32(), utf8()), dict_values)); + AssertTypeEqual(*uint8(), + *checked_cast(*builder->type()).index_type()); + + auto& dict_builder = checked_cast&>(*builder); + ASSERT_OK(dict_builder.Append("a")); + ASSERT_OK(dict_builder.Append("b")); + ASSERT_OK_AND_ASSIGN(auto result, dict_builder.Finish()); + ASSERT_OK(result->ValidateFull()); + AssertTypeEqual(*uint8(), + *checked_cast(*result->type()).index_type()); +} + +TEST(TestDictionaryBuilderIndexType, ExactIndexBuilderPreservesUnsignedIndexType) { + // MakeBuilderExactIndex is a separate builder that honours unsigned index types on its + // own; guard that it keeps doing so. + auto dict_type = dictionary(uint16(), utf8()); + std::unique_ptr boxed_builder; + ASSERT_OK(MakeBuilderExactIndex(default_memory_pool(), dict_type, &boxed_builder)); + AssertTypeEqual( + *uint16(), + *checked_cast(*boxed_builder->type()).index_type()); + + ASSERT_OK_AND_ASSIGN(auto result, boxed_builder->Finish()); + ASSERT_OK(result->ValidateFull()); + AssertTypeEqual(*uint16(), + *checked_cast(*result->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..4dd4e6ea44fd 100644 --- a/cpp/src/arrow/array/builder_dict.h +++ b/cpp/src/arrow/array/builder_dict.h @@ -134,6 +134,12 @@ 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 /// @@ -154,14 +160,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 use_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), + use_unsigned_index_(use_unsigned_index) {} template explicit DictionaryBuilderBase( @@ -199,14 +207,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 use_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), + use_unsigned_index_(use_unsigned_index) {} template explicit DictionaryBuilderBase( @@ -244,14 +254,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 use_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), + use_unsigned_index_(use_unsigned_index) {} ~DictionaryBuilderBase() override = default; @@ -486,6 +497,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, use_unsigned_index_); *out_indices = MakeArray(indices_data); *out_delta = MakeArray(delta_data); return Status::OK(); @@ -498,7 +510,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(), use_unsigned_index_), value_type_, + ordered_); } protected: @@ -570,6 +584,7 @@ class DictionaryBuilderBase : public ArrayBuilder { BuilderType indices_builder_; std::shared_ptr value_type_; bool ordered_ = false; + bool use_unsigned_index_ = false; }; template @@ -581,10 +596,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 use_unsigned_index = false) : ArrayBuilder(pool, alignment), indices_builder_(start_int_size, pool, alignment), - ordered_(ordered) {} + ordered_(ordered), + use_unsigned_index_(use_unsigned_index) {} explicit DictionaryBuilderBase(const std::shared_ptr& value_type, MemoryPool* pool = default_memory_pool(), @@ -623,10 +640,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 use_unsigned_index = false) : ArrayBuilder(pool, alignment), indices_builder_(pool, alignment), - ordered_(ordered) {} + ordered_(ordered), + use_unsigned_index_(use_unsigned_index) {} /// \brief Append a scalar null value Status AppendNull() final { @@ -678,7 +696,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, use_unsigned_index_), + null(), ordered_); (*out)->dictionary = NullArray(0).data(); return Status::OK(); } @@ -690,12 +709,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(), use_unsigned_index_), 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 2190f1ce04a9..1d0a42099b3e 100644 --- a/cpp/src/arrow/builder.cc +++ b/cpp/src/arrow/builder.cc @@ -27,6 +27,7 @@ #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 { @@ -38,6 +39,40 @@ class MemoryPool; using arrow::internal::checked_cast; +namespace internal { + +/// Return the unsigned integer type of the same width when an unsigned dictionary +/// index type was requested (GH-37476). +/// +/// 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. The +/// 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 // type. Used to reduce the number of template instantiations in the // exact_index_type case below, to reduce build time and memory usage. @@ -170,9 +205,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 +218,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/src/arrow/python/arrow_to_pandas.cc b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc index 348d352a0482..ebf151ed6086 100644 --- a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc +++ b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc @@ -1863,13 +1863,17 @@ class CategoricalWriter } Status WriteIndicesUniform(const ChunkedArray& data) { - // For unsigned types, upcast to signed since pandas uses -1 for nulls - // uint8 to int16, uint16 to int32, uint32 to int64, signed types unchanged + // For unsigned types, use a signed output since pandas uses -1 for nulls: + // uint8 to int16, uint16 to int32, uint32 to int64. uint64 also maps to int64, + // which is safe because the indices are bounds-checked below the dictionary length + // and so never reach the int64 range limit. Signed types are unchanged. using OutputType = std::conditional_t< std::is_same::value, int16_t, std::conditional_t< std::is_same::value, int32_t, - std::conditional_t::value, int64_t, T>>>; + std::conditional_t< + std::is_same::value, int64_t, + std::conditional_t::value, int64_t, T>>>>; const int npy_output_type = std::is_same::value ? NPY_INT16 : std::is_same::value ? NPY_INT32 : std::is_same::value @@ -2044,9 +2048,7 @@ Status MakeWriter(const PandasOptions& options, PandasWriter::type writer_type, CATEGORICAL_CASE(UInt8Type); CATEGORICAL_CASE(UInt16Type); CATEGORICAL_CASE(UInt32Type); - case Type::UINT64: - return Status::TypeError( - "Converting UInt64 dictionary indices to pandas is not supported."); + CATEGORICAL_CASE(UInt64Type); default: // Unreachable ARROW_DCHECK(false); diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index adc3e097b54a..6caf9d2b5550 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -4527,3 +4527,42 @@ 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.int8(), pa.int16(), pa.int32(), pa.int64(), + pa.uint8(), pa.uint16(), pa.uint32(), + pa.uint64()]) +def test_dictionary_array_preserves_index_type(index_type): + # GH-37476: an unsigned dictionary index type must be preserved, not replaced by the + # signed integer type of the same width. Signed index types are kept as-is. + 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("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)] + arr = pa.array(values, type=pa.dictionary(start_type, pa.string())) + assert arr.type == pa.dictionary(widened_type, pa.string()) + assert arr.to_pylist() == values + + +@pytest.mark.pandas +def test_dictionary_uint64_index_to_pandas(): + # GH-37476: uint64 dictionary indices are preserved, and converting them to pandas + # maps the indices to int64 categorical codes, which is safe because the indices are + # bounded by the dictionary length. + arr = pa.array(["a", "b", None, "a"], type=pa.dictionary(pa.uint64(), pa.string())) + result = arr.to_pandas() + assert list(result.cat.categories) == ["a", "b"] + assert result.cat.codes.tolist() == [0, 1, -1, 0] diff --git a/python/pyarrow/tests/test_pandas.py b/python/pyarrow/tests/test_pandas.py index 2a782de1648c..d2ad7ac7f3e4 100644 --- a/python/pyarrow/tests/test_pandas.py +++ b/python/pyarrow/tests/test_pandas.py @@ -4142,21 +4142,14 @@ def test_dictionary_with_pandas(): d1 = pa.DictionaryArray.from_arrays(indices, dictionary) d2 = pa.DictionaryArray.from_arrays(indices, dictionary, mask=mask) - if index_type == 'uint64': - # uint64 is not supported due to overflow risk (values > 2^63-1) - with pytest.raises(TypeError, - match="UInt64 dictionary indices"): - d1.to_pandas() - continue - pandas1 = d1.to_pandas() # Pandas Categorical uses signed int codes. Arrow converts: - # uint8 to int16, uint16 to int32, uint32 to int64, signed types unchanged + # uint8 to int16, uint16 to int32, uint32 and uint64 to int64, signed unchanged if index_type == 'uint8': compare_indices = indices.astype('int16') elif index_type == 'uint16': compare_indices = indices.astype('int32') - elif index_type == 'uint32': + elif index_type in ('uint32', 'uint64'): compare_indices = indices.astype('int64') else: compare_indices = indices @@ -4172,7 +4165,7 @@ def test_dictionary_with_pandas(): signed_indices = indices.astype('int16') elif index_type == 'uint16': signed_indices = indices.astype('int32') - elif index_type == 'uint32': + elif index_type in ('uint32', 'uint64'): signed_indices = indices.astype('int64') else: signed_indices = indices