Skip to content

GH-50695: [C++] Use the full unsigned range for unsigned dictionary index types #50696

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mkzung wants to merge 1 commit into
base: main
Choose a base branch
from
Open

GH-50695: [C++] Use the full unsigned range for unsigned dictionary index types #50696

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 50 additions & 6 deletions cpp/src/arrow/array/array_dict_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<DictionaryBuilder<StringType>&>(*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<const DictionaryType&>(*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<const DictionaryArray&>(*full);
const auto& indices = checked_cast<const UInt8Array&>(*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<const DictionaryType&>(*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<DictionaryBuilder<StringType>&>(*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<const DictionaryType&>(*signed_result->type()).index_type());
}

TEST(TestDictionaryBuilderIndexType, NullValueTypePreservesUnsignedIndexType) {
Expand Down Expand Up @@ -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<DictionaryBuilder<StringType>&>(*boxed_builder);

for (int i = 0; i < 200; ++i) {
ASSERT_OK(builder.Append(std::to_string(i)));
}

std::shared_ptr<Array> 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<const UInt8Array&>(*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
Expand Down
125 changes: 121 additions & 4 deletions cpp/src/arrow/array/builder_adaptive.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,26 @@ std::shared_ptr<DataType> AdaptiveUIntBuilder::type() const {

std::shared_ptr<DataType> 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<const int64_t*>(pending_data_),
Expand All @@ -129,8 +149,9 @@ std::shared_ptr<DataType> 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<ArrayData>* out) {
RETURN_NOT_OK(CommitPendingData());
Expand Down Expand Up @@ -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<const int64_t*>(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<const int64_t*>(pending_data_),
pending_pos_, valid_bytes));
}
pending_has_nulls_ = false;
pending_pos_ = 0;
return Status::OK();
Expand Down Expand Up @@ -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<uint8_t*>(raw_data_) + length_,
chunk_size);
break;
case 2:
internal::DowncastUInts(values, reinterpret_cast<uint16_t*>(raw_data_) + length_,
chunk_size);
break;
case 4:
internal::DowncastUInts(values, reinterpret_cast<uint32_t*>(raw_data_) + length_,
chunk_size);
break;
case 8:
internal::DowncastUInts(values, reinterpret_cast<uint64_t*>(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 <typename new_type>
Status AdaptiveIntBuilder::ExpandUIntSizeN() {
switch (int_size_) {
case 1:
return ExpandIntSizeInternal<new_type, uint8_t>();
case 2:
return ExpandIntSizeInternal<new_type, uint16_t>();
case 4:
return ExpandIntSizeInternal<new_type, uint32_t>();
case 8:
return ExpandIntSizeInternal<new_type, uint64_t>();
default:
DCHECK(false);
}
return Status::OK();
}

Status AdaptiveIntBuilder::ExpandUIntSize(uint8_t new_int_size) {
switch (new_int_size) {
case 1:
return ExpandUIntSizeN<uint8_t>();
case 2:
return ExpandUIntSizeN<uint16_t>();
case 4:
return ExpandUIntSizeN<uint32_t>();
case 8:
return ExpandUIntSizeN<uint64_t>();
default:
DCHECK(false);
}
return Status::OK();
}

Status AdaptiveUIntBuilder::CommitPendingData() {
if (pending_pos_ == 0) {
return Status::OK();
Expand All @@ -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<const uint64_t*>(values), length,
valid_bytes);
}
return AppendValuesInternal(values, length, valid_bytes);
}

Expand Down
16 changes: 15 additions & 1 deletion cpp/src/arrow/array/builder_adaptive.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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 <typename new_type>
Status ExpandIntSizeN();
template <typename new_type>
Status ExpandUIntSizeN();

bool use_unsigned_range_ = false;
};

/// @}
Expand Down
Loading
Loading