Skip to content

Commit 91f599c

Browse files
zhangweilstXeon Liu
zhangweilst
authored and
Xeon Liu
committed
feat: support dictionary for binary type
1 parent b7ed2e1 commit 91f599c

18 files changed

Lines changed: 377 additions & 114 deletions

‎docs/source/user_guide/compaction.rst‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,9 +122,10 @@ one is read as ordinary values, and the writer takes both. A high-cardinality
122122
column that started dictionary-encoded and fell back to plain encoding therefore
123123
does not qualify, even though it still carries a dictionary page. ``BINARY`` is
124124
not forwarded although Parquet stores it in the same physical type and
125-
dictionary-encodes it the same way, because the value accessors cannot read a
126-
``BINARY`` dictionary. The rewrite also overrides the option back to ``false``
127-
when the table writes a format other than Parquet, when
125+
dictionary-encodes it the same way; this option only selects ``STRING`` columns.
126+
A file with a serialized Arrow dictionary schema can independently produce a
127+
``BINARY`` dictionary, which the value accessors also support. The rewrite also
128+
overrides the option back to ``false`` when the table writes a format other than Parquet, when
128129
``parquet.enable-dictionary`` is ``false`` because the writer would only expand
129130
the values again, or when variant/map shredding is configured because those
130131
writers reshape each batch against a fixed physical schema. Setting the option

‎src/paimon/common/predicate/literal_converter.cpp‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -214,17 +214,26 @@ Result<std::vector<Literal>> LiteralConverter::ConvertLiteralsFromArray(const ar
214214
auto* dict_type = checked_cast<arrow::DictionaryType*>(dict_array.type().get());
215215
auto value_type_id = dict_type->value_type()->id();
216216
auto index_type_id = dict_type->index_type()->id();
217-
if (value_type_id == arrow::Type::type::STRING &&
217+
if ((value_type_id == arrow::Type::type::STRING ||
218+
value_type_id == arrow::Type::type::BINARY) &&
218219
index_type_id == arrow::Type::type::INT32) {
219-
return GetLiteralFromDictionaryArray<arrow::StringArray, arrow::Int32Array>(
220-
dict_array, FieldType::STRING, own_data);
221-
} else if (value_type_id == arrow::Type::type::LARGE_STRING &&
220+
FieldType literal_type = value_type_id == arrow::Type::type::STRING
221+
? FieldType::STRING
222+
: FieldType::BINARY;
223+
return GetLiteralFromDictionaryArray<arrow::BinaryArray, arrow::Int32Array>(
224+
dict_array, literal_type, own_data);
225+
} else if ((value_type_id == arrow::Type::type::LARGE_STRING ||
226+
value_type_id == arrow::Type::type::LARGE_BINARY) &&
222227
index_type_id == arrow::Type::type::INT64) {
223-
return GetLiteralFromDictionaryArray<arrow::LargeStringArray, arrow::Int64Array>(
224-
dict_array, FieldType::STRING, own_data);
228+
FieldType literal_type = value_type_id == arrow::Type::type::LARGE_STRING
229+
? FieldType::STRING
230+
: FieldType::BINARY;
231+
return GetLiteralFromDictionaryArray<arrow::LargeBinaryArray, arrow::Int64Array>(
232+
dict_array, literal_type, own_data);
225233
} else {
226234
return Status::Invalid(
227-
"only support [STRING, INT32] or [LARGE_STRING, INT64] for DictionaryArray");
235+
"only support [STRING|BINARY, INT32] or "
236+
"[LARGE_STRING|LARGE_BINARY, INT64] for DictionaryArray");
228237
}
229238
}
230239
default:

‎src/paimon/common/predicate/literal_converter.h‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -134,17 +134,12 @@ class PAIMON_EXPORT LiteralConverter {
134134
literals.emplace_back(literal_type);
135135
} else {
136136
int64_t dict_index = indices->Value(i);
137-
if constexpr (std::is_same_v<DictArrayType, arrow::StringArray>) {
138-
int32_t length = 0;
139-
const uint8_t* value = dictionary->GetValue(dict_index, &length);
140-
literals.emplace_back(literal_type, reinterpret_cast<const char*>(value),
141-
length, own_data);
142-
} else {
143-
int64_t length = 0;
144-
const uint8_t* value = dictionary->GetValue(dict_index, &length);
145-
literals.emplace_back(literal_type, reinterpret_cast<const char*>(value),
146-
length, own_data);
137+
if (dictionary->IsNull(dict_index)) {
138+
literals.emplace_back(literal_type);
139+
continue;
147140
}
141+
auto value = dictionary->GetView(dict_index);
142+
literals.emplace_back(literal_type, value.data(), value.size(), own_data);
148143
}
149144
}
150145
return literals;

‎src/paimon/common/predicate/literal_converter_test.cpp‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,59 @@ TEST_F(LiteralConverterTest, TestDictType) {
449449
Literal(FieldType::STRING, "foo", 3), Literal(FieldType::STRING)}));
450450
}
451451

452+
TEST_F(LiteralConverterTest, TestDictionaryValueAndIndexTypes) {
453+
const std::vector<std::pair<std::shared_ptr<arrow::DataType>, std::shared_ptr<arrow::DataType>>>
454+
types = {{arrow::utf8(), arrow::int32()},
455+
{arrow::binary(), arrow::int32()},
456+
{arrow::large_utf8(), arrow::int64()},
457+
{arrow::large_binary(), arrow::int64()}};
458+
for (const auto& [value_type, index_type] : types) {
459+
SCOPED_TRACE(value_type->ToString());
460+
const FieldType literal_type =
461+
arrow::is_string(value_type->id()) ? FieldType::STRING : FieldType::BINARY;
462+
auto dictionary = arrow::ipc::internal::json::ArrayFromJSON(
463+
value_type, R"(["unused", "a\u0000b", "", null, "tail"])")
464+
.ValueOrDie()
465+
->Slice(1);
466+
const std::vector<Literal> expected = {
467+
Literal(literal_type, "a\0b", 3), Literal(literal_type),
468+
Literal(literal_type, "", 0), Literal(literal_type),
469+
Literal(literal_type, "tail", 4), Literal(literal_type, "a\0b", 3)};
470+
SCOPED_TRACE(index_type->ToString());
471+
auto indices =
472+
arrow::ipc::internal::json::ArrayFromJSON(index_type, "[3, 0, null, 1, 2, 3, 0]")
473+
.ValueOrDie();
474+
auto array = arrow::DictionaryArray::FromArrays(indices, dictionary).ValueOrDie();
475+
ASSERT_TRUE(array->ValidateFull().ok());
476+
auto sliced = array->Slice(1);
477+
for (bool own_data : {false, true}) {
478+
SCOPED_TRACE(own_data);
479+
ASSERT_OK_AND_ASSIGN(std::vector<Literal> actual,
480+
LiteralConverter::ConvertLiteralsFromArray(*sliced, own_data));
481+
ASSERT_EQ(expected, actual);
482+
}
483+
}
484+
}
485+
486+
TEST_F(LiteralConverterTest, TestBinaryDictionaryPreservesBytes) {
487+
const std::string bytes("\x00\xff\x80", 3);
488+
arrow::BinaryBuilder builder;
489+
ASSERT_TRUE(builder.Append(bytes).ok());
490+
std::shared_ptr<arrow::Array> dictionary;
491+
ASSERT_TRUE(builder.Finish(&dictionary).ok());
492+
auto indices =
493+
arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, null, 0]").ValueOrDie();
494+
auto array = arrow::DictionaryArray::FromArrays(indices, dictionary).ValueOrDie();
495+
ASSERT_OK_AND_ASSIGN(std::vector<Literal> actual,
496+
LiteralConverter::ConvertLiteralsFromArray(*array, /*own_data=*/true));
497+
array.reset();
498+
dictionary.reset();
499+
ASSERT_EQ(actual,
500+
std::vector<Literal>({Literal(FieldType::BINARY, bytes.data(), bytes.size()),
501+
Literal(FieldType::BINARY),
502+
Literal(FieldType::BINARY, bytes.data(), bytes.size())}));
503+
}
504+
452505
TEST_F(LiteralConverterTest, TestLiteralsToArray) {
453506
// Every writable field type converts to an array and back, the null literal included, so each
454507
// case asserts `ConvertLiteralsToArray` and `ConvertLiteralsFromArray` agree on every type.

‎src/paimon/common/predicate/null_false_leaf_binary_function.cpp‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,15 @@ std::optional<FieldType> ArrayFieldType(const arrow::Array& array) {
9898
const std::shared_ptr<arrow::DataType>& type = array.type();
9999
if (type->id() == arrow::Type::DICTIONARY) {
100100
const auto& dict_type = checked_cast<const arrow::DictionaryType&>(*type);
101-
const bool is_string = dict_type.value_type()->id() == arrow::Type::STRING &&
102-
dict_type.index_type()->id() == arrow::Type::INT32;
103-
const bool is_large_string = dict_type.value_type()->id() == arrow::Type::LARGE_STRING &&
104-
dict_type.index_type()->id() == arrow::Type::INT64;
105-
if (is_string || is_large_string) {
106-
return FieldType::STRING;
101+
const arrow::Type::type value_type = dict_type.value_type()->id();
102+
const arrow::Type::type index_type = dict_type.index_type()->id();
103+
if ((value_type == arrow::Type::STRING || value_type == arrow::Type::BINARY) &&
104+
index_type == arrow::Type::INT32) {
105+
return value_type == arrow::Type::STRING ? FieldType::STRING : FieldType::BINARY;
106+
}
107+
if ((value_type == arrow::Type::LARGE_STRING || value_type == arrow::Type::LARGE_BINARY) &&
108+
index_type == arrow::Type::INT64) {
109+
return value_type == arrow::Type::LARGE_STRING ? FieldType::STRING : FieldType::BINARY;
107110
}
108111
return std::nullopt;
109112
}

‎src/paimon/common/predicate/null_false_leaf_binary_function_test.cpp‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,28 @@ TEST_F(NullFalseLeafBinaryFunctionTest, TestLargeStringDictionary) {
339339
std::vector<char>({0, 1, 0, 0}));
340340
}
341341

342+
TEST_F(NullFalseLeafBinaryFunctionTest, TestBinaryDictionary) {
343+
const std::vector<std::pair<std::shared_ptr<arrow::DataType>, std::shared_ptr<arrow::DataType>>>
344+
types = {{arrow::binary(), arrow::int32()}, {arrow::large_binary(), arrow::int64()}};
345+
for (const auto& [value_type, index_type] : types) {
346+
SCOPED_TRACE(value_type->ToString());
347+
auto dictionary =
348+
arrow::ipc::internal::json::ArrayFromJSON(value_type, R"(["a\u0000b", "a", "", null])")
349+
.ValueOrDie();
350+
SCOPED_TRACE(index_type->ToString());
351+
auto indices =
352+
arrow::ipc::internal::json::ArrayFromJSON(index_type, "[1, 0, null, 1, 2, 3, 0]")
353+
.ValueOrDie();
354+
auto array = arrow::DictionaryArray::FromArrays(indices, dictionary).ValueOrDie()->Slice(1);
355+
const auto literal = BinaryLiteral(std::string("a\0b", 3));
356+
ASSERT_EQ(Eval(Equal::Instance(), literal, array), std::vector<char>({1, 0, 0, 0, 0, 1}));
357+
ASSERT_EQ(Eval(LessThan::Instance(), literal, array),
358+
std::vector<char>({0, 0, 1, 1, 0, 0}));
359+
ASSERT_EQ(Eval(NotEqual::Instance(), BinaryLiteral(""), array),
360+
std::vector<char>({1, 0, 1, 0, 0, 1}));
361+
}
362+
}
363+
342364
TEST_F(NullFalseLeafBinaryFunctionTest, TestDictionaryWithNullValue) {
343365
// A kernel decodes the dictionary, so a row pointing at a null dictionary value becomes a null
344366
// row and is false for every comparison. The row by row path reads the slot the index names

‎src/paimon/core/casting/casting_utils.cpp‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,19 @@
2323
#include "paimon/common/utils/checked_cast.h"
2424

2525
namespace paimon {
26+
Result<std::shared_ptr<arrow::Array>> CastingUtils::DecodeDictionary(
27+
const std::shared_ptr<arrow::Array>& array, arrow::MemoryPool* pool) {
28+
if (array->type_id() != arrow::Type::DICTIONARY) {
29+
return array;
30+
}
31+
const auto& dictionary_type = checked_cast<const arrow::DictionaryType&>(*array->type());
32+
std::shared_ptr<arrow::DataType> value_type = dictionary_type.value_type();
33+
if (value_type->id() == arrow::Type::LARGE_STRING) {
34+
value_type = arrow::utf8();
35+
}
36+
return Cast(array, value_type, arrow::compute::CastOptions::Safe(), pool);
37+
}
38+
2639
Result<std::shared_ptr<arrow::Array>> CastingUtils::Cast(
2740
const std::shared_ptr<arrow::Array>& src_array,
2841
const std::shared_ptr<arrow::DataType>& target_type, const arrow::compute::CastOptions& options,

‎src/paimon/core/casting/casting_utils.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,12 @@ class PAIMON_EXPORT CastingUtils {
4747
const std::shared_ptr<arrow::DataType>& target_type,
4848
const arrow::compute::CastOptions& options, arrow::MemoryPool* pool);
4949

50+
/// Decodes a dictionary to its value type, normalizing LARGE_STRING to STRING for readers
51+
/// that widen string offsets. Other value types, including binary offsets, are preserved.
52+
/// Returns a non-dictionary array unchanged.
53+
static Result<std::shared_ptr<arrow::Array>> DecodeDictionary(
54+
const std::shared_ptr<arrow::Array>& array, arrow::MemoryPool* pool);
55+
5056
template <typename SrcScalar, typename SrcDataType, typename TargetScalar,
5157
typename TargetDataType>
5258
static Result<Literal> Cast(const Literal& literal,

‎src/paimon/core/casting/casting_utils_test.cpp‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,59 @@ TEST_F(CastingUtilsTest, TestDictionaryToString) {
5353
ASSERT_TRUE(result_array->Equals(string_array));
5454
}
5555

56+
TEST_F(CastingUtilsTest, TestDecodeDictionaryPreservesValueType) {
57+
for (const auto& value_type : {arrow::utf8(), arrow::large_utf8(), arrow::binary(),
58+
arrow::large_binary(), arrow::int32()}) {
59+
SCOPED_TRACE(value_type->ToString());
60+
const bool is_integer = value_type->id() == arrow::Type::INT32;
61+
auto dictionary =
62+
arrow::ipc::internal::json::ArrayFromJSON(
63+
value_type, is_integer ? "[10, null, 20]" : R"(["a\u0000b", null, ""])")
64+
.ValueOrDie();
65+
const auto target_type =
66+
value_type->id() == arrow::Type::LARGE_STRING ? arrow::utf8() : value_type;
67+
auto expected = arrow::ipc::internal::json::ArrayFromJSON(
68+
target_type, is_integer ? "[10, null, null, 20, 10]"
69+
: R"(["a\u0000b", null, null, "", "a\u0000b"])")
70+
.ValueOrDie();
71+
for (const auto& index_type :
72+
{arrow::int8(), arrow::int16(), arrow::int32(), arrow::int64()}) {
73+
SCOPED_TRACE(index_type->ToString());
74+
auto indices =
75+
arrow::ipc::internal::json::ArrayFromJSON(index_type, "[2, 0, null, 1, 2, 0]")
76+
.ValueOrDie();
77+
auto array = arrow::DictionaryArray::FromArrays(indices, dictionary).ValueOrDie();
78+
ASSERT_OK_AND_ASSIGN(
79+
std::shared_ptr<arrow::Array> decoded,
80+
CastingUtils::DecodeDictionary(array->Slice(1), arrow_pool_.get()));
81+
ASSERT_TRUE(decoded->ValidateFull().ok());
82+
ASSERT_TRUE(decoded->Equals(expected)) << decoded->ToString();
83+
}
84+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> unchanged,
85+
CastingUtils::DecodeDictionary(dictionary, arrow_pool_.get()));
86+
ASSERT_EQ(unchanged, dictionary);
87+
}
88+
}
89+
90+
TEST_F(CastingUtilsTest, TestDecodeBinaryDictionaryPreservesNonUtf8) {
91+
const std::string bytes("\x00\xff\x80", 3);
92+
arrow::BinaryBuilder builder;
93+
ASSERT_TRUE(builder.Append(bytes).ok());
94+
std::shared_ptr<arrow::Array> dictionary;
95+
ASSERT_TRUE(builder.Finish(&dictionary).ok());
96+
auto indices =
97+
arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, null, 0]").ValueOrDie();
98+
auto array = arrow::DictionaryArray::FromArrays(indices, dictionary).ValueOrDie();
99+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::Array> decoded,
100+
CastingUtils::DecodeDictionary(array, arrow_pool_.get()));
101+
ASSERT_TRUE(decoded->ValidateFull().ok());
102+
ASSERT_EQ(decoded->type_id(), arrow::Type::BINARY);
103+
auto binary = checked_pointer_cast<arrow::BinaryArray>(decoded);
104+
ASSERT_EQ(bytes, binary->GetView(0));
105+
ASSERT_TRUE(binary->IsNull(1));
106+
ASSERT_EQ(bytes, binary->GetView(2));
107+
}
108+
56109
TEST_F(CastingUtilsTest, TestTimestampToTimestampWithTimezone) {
57110
// local no tz -> utc tz
58111
auto src_array = arrow::ipc::internal::json::ArrayFromJSON(

‎src/paimon/core/global_index/global_index_write_task.cpp‎

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -164,22 +164,11 @@ Result<std::unique_ptr<BatchReader>> CreateBatchReader(
164164
return table_read->CreateReader(indexed_split);
165165
}
166166

167-
Result<std::shared_ptr<arrow::Array>> CastDictionaryArrayToString(
167+
Result<std::shared_ptr<arrow::Array>> DecodeDictionaryArrays(
168168
const std::shared_ptr<arrow::Array>& array, arrow::MemoryPool* pool) {
169169
arrow::Type::type type_id = array->type_id();
170170
if (type_id == arrow::Type::DICTIONARY) {
171-
const auto* dictionary_type =
172-
checked_cast<const arrow::DictionaryType*>(array->type().get());
173-
arrow::Type::type value_type = dictionary_type->value_type()->id();
174-
if (value_type != arrow::Type::STRING && value_type != arrow::Type::LARGE_STRING) {
175-
return Status::Invalid(fmt::format(
176-
"GlobalIndexWriteTask cannot decode dictionary array with value type {}",
177-
dictionary_type->value_type()->ToString()));
178-
}
179-
PAIMON_ASSIGN_OR_RAISE(
180-
std::shared_ptr<arrow::Array> casted_array,
181-
CastingUtils::Cast(array, arrow::utf8(), arrow::compute::CastOptions::Safe(), pool));
182-
return casted_array;
171+
return CastingUtils::DecodeDictionary(array, pool);
183172
}
184173
if (type_id != arrow::Type::STRUCT && type_id != arrow::Type::MAP &&
185174
type_id != arrow::Type::LIST) {
@@ -193,7 +182,7 @@ Result<std::shared_ptr<arrow::Array>> CastDictionaryArrayToString(
193182
for (int32_t i = 0; i < struct_array->num_fields(); i++) {
194183
std::shared_ptr<arrow::Array> child = struct_array->field(i);
195184
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> casted_child,
196-
CastDictionaryArrayToString(child, pool));
185+
DecodeDictionaryArrays(child, pool));
197186
if (casted_child != child && children.empty()) {
198187
children = struct_array->fields();
199188
}
@@ -221,9 +210,9 @@ Result<std::shared_ptr<arrow::Array>> CastDictionaryArrayToString(
221210
std::shared_ptr<arrow::Array> original_keys = map_array->keys();
222211
std::shared_ptr<arrow::Array> original_items = map_array->items();
223212
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> keys,
224-
CastDictionaryArrayToString(original_keys, pool));
213+
DecodeDictionaryArrays(original_keys, pool));
225214
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> items,
226-
CastDictionaryArrayToString(original_items, pool));
215+
DecodeDictionaryArrays(original_items, pool));
227216
if (keys == original_keys && items == original_items) {
228217
return array;
229218
}
@@ -239,7 +228,7 @@ Result<std::shared_ptr<arrow::Array>> CastDictionaryArrayToString(
239228
std::shared_ptr<arrow::ListArray> list_array = checked_pointer_cast<arrow::ListArray>(array);
240229
std::shared_ptr<arrow::Array> original_values = list_array->values();
241230
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> values,
242-
CastDictionaryArrayToString(original_values, pool));
231+
DecodeDictionaryArrays(original_values, pool));
243232
if (values == original_values) {
244233
return array;
245234
}
@@ -296,7 +285,7 @@ Result<std::vector<GlobalIndexIOMeta>> BuildIndex(
296285
writer_field_name));
297286
}
298287
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> decoded_writer_array,
299-
CastDictionaryArrayToString(writer_array, arrow_pool));
288+
DecodeDictionaryArrays(writer_array, arrow_pool));
300289
writer_arrays.push_back(std::move(decoded_writer_array));
301290
}
302291
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(

0 commit comments

Comments
 (0)