Skip to content

Commit e52da51

Browse files
morningmanGabriel39hubgetersuxiaogang223claude
committed
[test](external) Port the branch-4.x Arrow convertor and Parquet timestamp tests
### What problem does this PR solve? Issue Number: None Related PR: #68381, #68441, #68396, #68883, #68516 Problem Summary: branch-4.x #68381 made the Arrow block convertors explicit and schema-owned, and #68441 aligned the Parquet and ORC timestamp semantics. master has their production code through #68396 and #68883, and #65446 and #68516, but some of their tests did not come along: - BE, Iceberg Variant writes through IcebergArrowBlockConvertor: a nullable Variant column and Variants nested in LIST, MAP and STRUCT become arrow.parquet.variant extension arrays with the encoded bytes unchanged, and the Parquet file carries the Variant logical type and the Iceberg field ids. These two cases came with #66620, but they build their Arrow schema themselves and test code master has since #68883. - BE, PaimonArrowBlockConvertor: encoded and typed Variants with outer NULLs keep their bytes, ARRAY<VARIANT> goes through the Paimon SerDe recursively, and an unsupported primitive is rejected at the Paimon hook. master tested these only through the generic Arrow path or without byte comparison. - FE: an ALTER CATALOG to an invalid hive.parquet.time-zone is rejected and a valid +05:45 is accepted; a Hive write without the property still sends an empty time zone; only a Parquet Hudi MOR table advertises native Parquet slices; Paimon sets timestamp_is_adjusted_to_utc for wall clock and instant timestamps, nested ones included; and a connector that does not use the Hive INT96 time zone does not get it from an HMS catalog. ### Release note None ### Check List (For Author) - Test: Unit Test - BE UT: DataTypeSerDeArrowTest, DataTypeVariantV2SerdeOutputTest and PaimonArrowBlockConvertorTest pass (42 tests, 5 of them new). - FE UT: HMSExternalCatalogParquetTimeZoneTest, PluginDrivenScanNodeScanProviderSelectionTest, HiveWritePlanProviderTest, HudiForceJniTest and PaimonScanPlanProviderTest pass. Each new FE assertion fails when the check it guards is removed. - Behavior changed: No - Does this need documentation: No Co-authored-by: Gabriel <[email protected]> Co-authored-by: daidai <[email protected]> Co-authored-by: Socrates <[email protected]> Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
1 parent cad4abb commit e52da51

7 files changed

Lines changed: 468 additions & 16 deletions

File tree

‎be/test/core/data_type_serde/data_type_serde_arrow_test.cpp‎

Lines changed: 255 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@
2222
#include <arrow/array/builder_decimal.h>
2323
#include <arrow/array/builder_nested.h>
2424
#include <arrow/array/builder_primitive.h>
25+
#include <arrow/extension/parquet_variant.h>
26+
#include <arrow/extension_type.h>
27+
#include <arrow/io/memory.h>
2528
#include <arrow/record_batch.h>
2629
#include <arrow/status.h>
2730
#include <arrow/type.h>
@@ -35,6 +38,8 @@
3538
#include <gtest/gtest-message.h>
3639
#include <gtest/gtest-test-part.h>
3740
#include <gtest/gtest.h>
41+
#include <parquet/api/reader.h>
42+
#include <parquet/arrow/writer.h>
3843

3944
#include <array>
4045
#include <chrono>
@@ -59,6 +64,7 @@
5964
#include "core/column/column_struct.h"
6065
#include "core/column/column_varbinary.h"
6166
#include "core/column/column_vector.h"
67+
#include "core/column/variant_v2/column_variant_v2.h"
6268
#include "core/data_type/common_data_type_serder_test.h"
6369
#include "core/data_type/data_type.h"
6470
#include "core/data_type/data_type_array.h"
@@ -88,6 +94,7 @@
8894
#include "core/value/vdatetime_value.h"
8995
#include "exec/common/arrow_column_to_doris_column.h"
9096
#include "exprs/function/cast/cast_to_datetimev2_impl.hpp"
97+
#include "exprs/function/parse/variant_string_parse.h"
9198
#include "format/arrow/arrow_block_convertor.h"
9299
#include "format/arrow/arrow_row_batch.h"
93100
#include "format/table/iceberg/iceberg_arrow_block_convertor.h"
@@ -948,6 +955,254 @@ TEST(DataTypeSerDeArrowTest, PlainArrowConverterDoesNotInferIcebergUuid) {
948955
EXPECT_NE(std::string::npos, status.to_string().find("Plain Arrow writer is not bound"));
949956
}
950957

958+
TEST(DataTypeSerDeArrowTest, IcebergVariantExtensionAndParquetSchema) {
959+
auto make_variant_column = []() {
960+
JsonStringToVariantEncoder encoder({.max_json_key_length = 1024,
961+
.throw_on_invalid_json = true,
962+
.check_duplicate_json_path = false});
963+
for (std::string_view json :
964+
{std::string_view {R"({"name":"doris","n":1})"}, std::string_view {"null"},
965+
std::string_view {R"([1,true,"x"])"}}) {
966+
encoder.add_json({json.data(), json.size()});
967+
}
968+
auto column = ColumnVariantV2::create();
969+
column->insert_encoded_batch(encoder.finish_batch());
970+
return column;
971+
};
972+
973+
auto variant_column = make_variant_column();
974+
std::vector<std::string> expected_metadata;
975+
std::vector<std::string> expected_values;
976+
const auto source_view = variant_column->read_view();
977+
for (size_t row = 0; row < variant_column->size(); ++row) {
978+
const VariantRef value = source_view.value_at(row);
979+
expected_metadata.emplace_back(value.metadata.data, value.metadata.size);
980+
expected_values.emplace_back(value.value.data, value.value.size);
981+
}
982+
983+
auto null_map = ColumnUInt8::create();
984+
null_map->get_data().assign({0, 0, 1});
985+
auto nullable_variant = ColumnNullable::create(std::move(variant_column), std::move(null_map));
986+
DataTypePtr variant_type = std::make_shared<DataTypeVariantV2>();
987+
DataTypePtr nullable_variant_type = make_nullable(variant_type);
988+
989+
Block block;
990+
block.insert(
991+
ColumnWithTypeAndName(nullable_variant->get_ptr(), nullable_variant_type, "payload"));
992+
993+
auto variant_storage = arrow::struct_({
994+
arrow::field("metadata", arrow::binary(), false),
995+
arrow::field("value", arrow::binary(), false),
996+
});
997+
auto arrow_variant = arrow::extension::variant(variant_storage);
998+
auto field_id_21 = arrow::KeyValueMetadata::Make({"PARQUET:field_id"}, {"21"});
999+
auto schema = arrow::schema({
1000+
arrow::field("payload", arrow_variant, true, field_id_21),
1001+
});
1002+
1003+
std::shared_ptr<arrow::RecordBatch> record_batch;
1004+
cctz::time_zone default_timezone;
1005+
Status status = convert_to_arrow_batch_for_test(
1006+
block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0,
1007+
block.rows(), make_arrow_convertor<iceberg::IcebergArrowBlockConvertor>);
1008+
ASSERT_TRUE(status.ok()) << status;
1009+
ASSERT_NE(nullptr, record_batch);
1010+
ASSERT_TRUE(record_batch->ValidateFull().ok());
1011+
1012+
auto extension_array =
1013+
std::dynamic_pointer_cast<arrow::ExtensionArray>(record_batch->column(0));
1014+
ASSERT_NE(nullptr, extension_array);
1015+
EXPECT_EQ("arrow.parquet.variant", extension_array->extension_type()->extension_name());
1016+
EXPECT_FALSE(extension_array->IsNull(0));
1017+
EXPECT_FALSE(extension_array->IsNull(1));
1018+
EXPECT_TRUE(extension_array->IsNull(2));
1019+
auto storage_array = std::static_pointer_cast<arrow::StructArray>(extension_array->storage());
1020+
auto metadata_array = std::static_pointer_cast<arrow::BinaryArray>(storage_array->field(0));
1021+
auto value_array = std::static_pointer_cast<arrow::BinaryArray>(storage_array->field(1));
1022+
for (int64_t row = 0; row < 2; ++row) {
1023+
EXPECT_EQ(expected_metadata[row], metadata_array->GetView(row));
1024+
EXPECT_EQ(expected_values[row], value_array->GetView(row));
1025+
}
1026+
1027+
auto sink_result = arrow::io::BufferOutputStream::Create();
1028+
ASSERT_TRUE(sink_result.ok()) << sink_result.status();
1029+
auto sink = std::move(sink_result).ValueUnsafe();
1030+
auto writer_result =
1031+
::parquet::arrow::FileWriter::Open(*schema, arrow::default_memory_pool(), sink);
1032+
ASSERT_TRUE(writer_result.ok()) << writer_result.status();
1033+
auto writer = std::move(writer_result).ValueUnsafe();
1034+
ASSERT_TRUE(writer->WriteRecordBatch(*record_batch).ok());
1035+
ASSERT_TRUE(writer->Close().ok());
1036+
auto buffer_result = sink->Finish();
1037+
ASSERT_TRUE(buffer_result.ok()) << buffer_result.status();
1038+
1039+
auto reader = ::parquet::ParquetFileReader::Open(
1040+
std::make_shared<arrow::io::BufferReader>(std::move(buffer_result).ValueUnsafe()));
1041+
const auto* root = reader->metadata()->schema()->group_node();
1042+
ASSERT_EQ(1, root->field_count());
1043+
const auto& payload_node = root->field(0);
1044+
ASSERT_NE(nullptr, payload_node->logical_type());
1045+
EXPECT_TRUE(payload_node->logical_type()->is_variant());
1046+
EXPECT_EQ(21, payload_node->field_id());
1047+
const auto& payload_group = static_cast<const ::parquet::schema::GroupNode&>(*payload_node);
1048+
ASSERT_EQ(2, payload_group.field_count());
1049+
EXPECT_EQ("metadata", payload_group.field(0)->name());
1050+
EXPECT_EQ(-1, payload_group.field(0)->field_id());
1051+
EXPECT_EQ("value", payload_group.field(1)->name());
1052+
EXPECT_EQ(-1, payload_group.field(1)->field_id());
1053+
}
1054+
1055+
TEST(DataTypeSerDeArrowTest, NestedIcebergVariantExtensionsAndParquetSchema) {
1056+
auto make_variant_column = [](std::initializer_list<std::string_view> json_values) {
1057+
JsonStringToVariantEncoder encoder({.max_json_key_length = 1024,
1058+
.throw_on_invalid_json = true,
1059+
.check_duplicate_json_path = false});
1060+
for (std::string_view json : json_values) {
1061+
encoder.add_json({json.data(), json.size()});
1062+
}
1063+
auto column = ColumnVariantV2::create();
1064+
column->insert_encoded_batch(encoder.finish_batch());
1065+
return column;
1066+
};
1067+
auto make_null_map = [](std::initializer_list<uint8_t> values) {
1068+
auto null_map = ColumnUInt8::create();
1069+
null_map->get_data().assign(values);
1070+
return null_map;
1071+
};
1072+
1073+
DataTypePtr variant_type = std::make_shared<DataTypeVariantV2>();
1074+
DataTypePtr nullable_variant_type = make_nullable(variant_type);
1075+
1076+
auto array_values =
1077+
ColumnNullable::create(make_variant_column({R"({"kind":"array"})", "null", "{}", "7"}),
1078+
make_null_map({0, 0, 1, 0}));
1079+
auto array_offsets = ColumnArray::ColumnOffsets::create();
1080+
array_offsets->get_data().assign({3, 4});
1081+
auto array_column = ColumnArray::create(std::move(array_values), std::move(array_offsets));
1082+
DataTypePtr array_type = std::make_shared<DataTypeArray>(variant_type);
1083+
1084+
auto map_keys_data = ColumnString::create();
1085+
for (std::string_view key : {std::string_view {"object"}, std::string_view {"json_null"},
1086+
std::string_view {"sql_null"}}) {
1087+
map_keys_data->insert_data(key.data(), key.size());
1088+
}
1089+
auto map_keys = ColumnNullable::create(std::move(map_keys_data), make_null_map({0, 0, 0}));
1090+
auto map_values = ColumnNullable::create(
1091+
make_variant_column({R"({"kind":"map"})", "null", "{}"}), make_null_map({0, 0, 1}));
1092+
auto map_offsets = ColumnArray::ColumnOffsets::create();
1093+
map_offsets->get_data().assign({3, 3});
1094+
auto map_column =
1095+
ColumnMap::create(std::move(map_keys), std::move(map_values), std::move(map_offsets));
1096+
DataTypePtr map_type = std::make_shared<DataTypeMap>(
1097+
make_nullable(std::make_shared<DataTypeString>()), nullable_variant_type);
1098+
1099+
auto labels = ColumnString::create();
1100+
labels->insert_data("first", 5);
1101+
labels->insert_data("second", 6);
1102+
auto struct_payloads = ColumnNullable::create(
1103+
make_variant_column({R"({"kind":"struct"})", "{}"}), make_null_map({0, 1}));
1104+
MutableColumns struct_children;
1105+
struct_children.emplace_back(std::move(labels));
1106+
struct_children.emplace_back(std::move(struct_payloads));
1107+
auto struct_column = ColumnStruct::create(std::move(struct_children));
1108+
DataTypePtr struct_type = std::make_shared<DataTypeStruct>(
1109+
DataTypes {std::make_shared<DataTypeString>(), nullable_variant_type},
1110+
Strings {"label", "payload"});
1111+
1112+
Block block;
1113+
block.insert(ColumnWithTypeAndName(std::move(array_column), array_type, "events"));
1114+
block.insert(ColumnWithTypeAndName(std::move(map_column), map_type, "attrs"));
1115+
block.insert(ColumnWithTypeAndName(std::move(struct_column), struct_type, "info"));
1116+
1117+
auto field_id = [](int id) {
1118+
return arrow::KeyValueMetadata::Make({"PARQUET:field_id"}, {std::to_string(id)});
1119+
};
1120+
auto variant_storage = arrow::struct_({
1121+
arrow::field("metadata", arrow::binary(), false),
1122+
arrow::field("value", arrow::binary(), false),
1123+
});
1124+
auto arrow_variant = arrow::extension::variant(variant_storage);
1125+
auto schema = arrow::schema({
1126+
arrow::field("events",
1127+
arrow::list(arrow::field("element", arrow_variant, true, field_id(31))),
1128+
true, field_id(30)),
1129+
arrow::field("attrs",
1130+
std::make_shared<arrow::MapType>(
1131+
arrow::field("key", arrow::utf8(), false, field_id(41)),
1132+
arrow::field("value", arrow_variant, true, field_id(42))),
1133+
true, field_id(40)),
1134+
arrow::field("info",
1135+
arrow::struct_({
1136+
arrow::field("label", arrow::utf8(), true, field_id(51)),
1137+
arrow::field("payload", arrow_variant, true, field_id(52)),
1138+
}),
1139+
true, field_id(50)),
1140+
});
1141+
1142+
std::shared_ptr<arrow::RecordBatch> record_batch;
1143+
cctz::time_zone default_timezone;
1144+
Status status = convert_to_arrow_batch_for_test(
1145+
block, schema, arrow::default_memory_pool(), &record_batch, default_timezone, 0,
1146+
block.rows(), make_arrow_convertor<iceberg::IcebergArrowBlockConvertor>);
1147+
ASSERT_TRUE(status.ok()) << status;
1148+
ASSERT_NE(nullptr, record_batch);
1149+
ASSERT_TRUE(record_batch->ValidateFull().ok()) << record_batch->ValidateFull();
1150+
1151+
auto events = std::static_pointer_cast<arrow::ListArray>(record_batch->column(0));
1152+
auto event_values = std::dynamic_pointer_cast<arrow::ExtensionArray>(events->values());
1153+
ASSERT_NE(nullptr, event_values);
1154+
EXPECT_EQ("arrow.parquet.variant", event_values->extension_type()->extension_name());
1155+
EXPECT_TRUE(event_values->IsNull(2));
1156+
1157+
auto attrs = std::static_pointer_cast<arrow::MapArray>(record_batch->column(1));
1158+
auto attr_values = std::dynamic_pointer_cast<arrow::ExtensionArray>(attrs->items());
1159+
ASSERT_NE(nullptr, attr_values);
1160+
EXPECT_EQ("arrow.parquet.variant", attr_values->extension_type()->extension_name());
1161+
EXPECT_TRUE(attr_values->IsNull(2));
1162+
1163+
auto info = std::static_pointer_cast<arrow::StructArray>(record_batch->column(2));
1164+
auto info_payload = std::dynamic_pointer_cast<arrow::ExtensionArray>(info->field(1));
1165+
ASSERT_NE(nullptr, info_payload);
1166+
EXPECT_EQ("arrow.parquet.variant", info_payload->extension_type()->extension_name());
1167+
EXPECT_TRUE(info_payload->IsNull(1));
1168+
1169+
auto sink_result = arrow::io::BufferOutputStream::Create();
1170+
ASSERT_TRUE(sink_result.ok()) << sink_result.status();
1171+
auto sink = std::move(sink_result).ValueUnsafe();
1172+
auto writer_result =
1173+
::parquet::arrow::FileWriter::Open(*schema, arrow::default_memory_pool(), sink);
1174+
ASSERT_TRUE(writer_result.ok()) << writer_result.status();
1175+
auto writer = std::move(writer_result).ValueUnsafe();
1176+
ASSERT_TRUE(writer->WriteRecordBatch(*record_batch).ok());
1177+
ASSERT_TRUE(writer->Close().ok());
1178+
auto buffer_result = sink->Finish();
1179+
ASSERT_TRUE(buffer_result.ok()) << buffer_result.status();
1180+
1181+
auto reader = ::parquet::ParquetFileReader::Open(
1182+
std::make_shared<arrow::io::BufferReader>(std::move(buffer_result).ValueUnsafe()));
1183+
const auto* root = reader->metadata()->schema()->group_node();
1184+
const auto& events_group = static_cast<const ::parquet::schema::GroupNode&>(*root->field(0));
1185+
const auto& event_list =
1186+
static_cast<const ::parquet::schema::GroupNode&>(*events_group.field(0));
1187+
const auto& event_variant = event_list.field(0);
1188+
ASSERT_NE(nullptr, event_variant->logical_type());
1189+
EXPECT_TRUE(event_variant->logical_type()->is_variant());
1190+
EXPECT_EQ(31, event_variant->field_id());
1191+
1192+
const auto& attrs_group = static_cast<const ::parquet::schema::GroupNode&>(*root->field(1));
1193+
const auto& key_value = static_cast<const ::parquet::schema::GroupNode&>(*attrs_group.field(0));
1194+
const auto& map_variant = key_value.field(1);
1195+
ASSERT_NE(nullptr, map_variant->logical_type());
1196+
EXPECT_TRUE(map_variant->logical_type()->is_variant());
1197+
EXPECT_EQ(42, map_variant->field_id());
1198+
1199+
const auto& info_group = static_cast<const ::parquet::schema::GroupNode&>(*root->field(2));
1200+
const auto& struct_variant = info_group.field(1);
1201+
ASSERT_NE(nullptr, struct_variant->logical_type());
1202+
EXPECT_TRUE(struct_variant->logical_type()->is_variant());
1203+
EXPECT_EQ(52, struct_variant->field_id());
1204+
}
1205+
9511206
TEST(DataTypeSerDeArrowTest, NestedIcebergUuidStringToFixedSizeBinary) {
9521207
auto block = std::make_shared<Block>();
9531208
DataTypePtr data_type = std::make_shared<DataTypeStruct>(

0 commit comments

Comments
 (0)