Skip to content

Commit a263262

Browse files
Gabriel39Gabriel39
committed
[fix](arrow-flight) Preserve nested logical type metadata
1 parent 4750b1c commit a263262

4 files changed

Lines changed: 339 additions & 16 deletions

File tree

‎be/src/format/arrow/arrow_row_batch.cpp‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,9 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type,
138138
std::shared_ptr<arrow::DataType> item_type;
139139
RETURN_IF_ERROR(convert_to_arrow_type(type_arr->get_nested_type(), &item_type, timezone,
140140
datetime_naive));
141-
*result = std::make_shared<arrow::ListType>(item_type);
141+
// Arrow stores metadata on fields, so implicit child fields lose the Doris logical type.
142+
*result = std::make_shared<arrow::ListType>(create_arrow_field_with_metadata(
143+
"item", item_type, true, type_arr->get_nested_type()->get_primitive_type()));
142144
break;
143145
}
144146
case TYPE_MAP: {
@@ -149,7 +151,11 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type,
149151
datetime_naive));
150152
RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_value_type(), &val_type, timezone,
151153
datetime_naive));
152-
*result = std::make_shared<arrow::MapType>(key_type, val_type);
154+
auto key_field = create_arrow_field_with_metadata(
155+
"key", key_type, false, type_map->get_key_type()->get_primitive_type());
156+
auto value_field = create_arrow_field_with_metadata(
157+
"value", val_type, true, type_map->get_value_type()->get_primitive_type());
158+
*result = std::make_shared<arrow::MapType>(key_field, value_field);
153159
break;
154160
}
155161
case TYPE_STRUCT: {
@@ -159,9 +165,10 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type,
159165
std::shared_ptr<arrow::DataType> field_type;
160166
RETURN_IF_ERROR(convert_to_arrow_type(type_struct->get_element(i), &field_type,
161167
timezone, datetime_naive));
162-
fields.push_back(
163-
std::make_shared<arrow::Field>(type_struct->get_element_name(i), field_type,
164-
type_struct->get_element(i)->is_nullable()));
168+
fields.push_back(create_arrow_field_with_metadata(
169+
type_struct->get_element_name(i), field_type,
170+
type_struct->get_element(i)->is_nullable(),
171+
type_struct->get_element(i)->get_primitive_type()));
165172
}
166173
*result = std::make_shared<arrow::StructType>(fields);
167174
break;
@@ -187,22 +194,32 @@ Status convert_to_arrow_type(const DataTypePtr& origin_type,
187194
return Status::OK();
188195
}
189196

190-
// Helper function to create an Arrow Field with type metadata if applicable, such as IP types
197+
// Logical types sharing Arrow storage need the same marker at the root and every child field.
191198
std::shared_ptr<arrow::Field> create_arrow_field_with_metadata(
192199
const std::string& field_name, const std::shared_ptr<arrow::DataType>& arrow_type,
193200
bool is_nullable, PrimitiveType primitive_type) {
194-
if (primitive_type == PrimitiveType::TYPE_IPV4) {
195-
auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV4"});
196-
return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata);
197-
} else if (primitive_type == PrimitiveType::TYPE_IPV6) {
198-
auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV6"});
199-
return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata);
200-
} else if (primitive_type == PrimitiveType::TYPE_LARGEINT) {
201-
auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"LARGEINT"});
202-
return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata);
203-
} else {
201+
const char* type_name;
202+
switch (primitive_type) {
203+
case TYPE_IPV4:
204+
type_name = "IPV4";
205+
break;
206+
case TYPE_IPV6:
207+
type_name = "IPV6";
208+
break;
209+
case TYPE_LARGEINT:
210+
type_name = "LARGEINT";
211+
break;
212+
case TYPE_JSONB:
213+
type_name = "JSON";
214+
break;
215+
case TYPE_VARIANT:
216+
type_name = "VARIANT";
217+
break;
218+
default:
204219
return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable);
205220
}
221+
auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {type_name});
222+
return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata);
206223
}
207224

208225
Status get_arrow_schema_from_block(const Block& block, std::shared_ptr<arrow::Schema>* result,

‎be/test/format/arrow/arrow_row_batch_test.cpp‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include "format/arrow/arrow_row_batch.h"
19+
20+
#include <arrow/api.h>
21+
#include <arrow/io/api.h>
22+
#include <arrow/ipc/api.h>
23+
#include <gtest/gtest.h>
24+
25+
#include "core/data_type/data_type_array.h"
26+
#include "core/data_type/data_type_factory.hpp"
27+
#include "core/data_type/data_type_map.h"
28+
#include "core/data_type/data_type_nullable.h"
29+
#include "core/data_type/data_type_number.h"
30+
#include "core/data_type/data_type_string.h"
31+
#include "core/data_type/data_type_struct.h"
32+
#include "format/arrow/arrow_block_convertor.h"
33+
34+
namespace doris {
35+
namespace {
36+
37+
void expect_logical_type(const std::shared_ptr<arrow::Field>& field, const std::string& name) {
38+
ASSERT_NE(nullptr, field->metadata()) << field->ToString();
39+
auto value = field->metadata()->Get("doris_type");
40+
ASSERT_TRUE(value.ok()) << value.status();
41+
EXPECT_EQ(name, *value);
42+
}
43+
44+
class ArrowLogicalTypeMetadataTest
45+
: public testing::TestWithParam<std::pair<PrimitiveType, const char*>> {};
46+
47+
TEST_P(ArrowLogicalTypeMetadataTest, PreservesTopLevelAndNestedFields) {
48+
const auto [primitive, name] = GetParam();
49+
auto type = make_nullable(DataTypeFactory::instance().create_data_type(primitive, false));
50+
auto string_type = std::make_shared<DataTypeString>();
51+
DataTypes types {type, std::make_shared<DataTypeArray>(type),
52+
std::make_shared<DataTypeStruct>(DataTypes {type, string_type},
53+
Strings {"typed", "text"}),
54+
std::make_shared<DataTypeMap>(string_type, type)};
55+
Block block;
56+
for (size_t i = 0; i < types.size(); ++i) {
57+
block.insert({types[i]->create_column(), types[i], std::to_string(i)});
58+
}
59+
std::shared_ptr<arrow::Schema> schema;
60+
ASSERT_TRUE(get_arrow_schema_from_block(block, &schema, "UTC", true).ok());
61+
expect_logical_type(schema->field(0), name);
62+
expect_logical_type(schema->field(1)->type()->field(0), name);
63+
expect_logical_type(schema->field(2)->type()->field(0), name);
64+
EXPECT_EQ(nullptr, schema->field(2)->type()->field(1)->metadata());
65+
const auto& map = static_cast<const arrow::MapType&>(*schema->field(3)->type());
66+
expect_logical_type(map.item_field(), name);
67+
EXPECT_FALSE(map.key_field()->nullable());
68+
EXPECT_TRUE(map.item_field()->nullable());
69+
EXPECT_EQ(nullptr, map.key_field()->metadata());
70+
}
71+
72+
INSTANTIATE_TEST_SUITE_P(LogicalTypes, ArrowLogicalTypeMetadataTest,
73+
testing::Values(std::make_pair(TYPE_LARGEINT, "LARGEINT"),
74+
std::make_pair(TYPE_IPV4, "IPV4"),
75+
std::make_pair(TYPE_IPV6, "IPV6"),
76+
std::make_pair(TYPE_JSONB, "JSON"),
77+
std::make_pair(TYPE_VARIANT, "VARIANT")));
78+
79+
TEST(ArrowRowBatchMetadataTest, PreservesMapKeysAndDeepNestingThroughIpc) {
80+
auto integer = std::make_shared<DataTypeInt128>();
81+
auto structure = std::make_shared<DataTypeStruct>(DataTypes {integer}, Strings {"number"});
82+
auto array = std::make_shared<DataTypeArray>(structure);
83+
auto type = std::make_shared<DataTypeMap>(integer, array);
84+
Block block;
85+
block.insert({type->create_column(), type, "m"});
86+
std::shared_ptr<arrow::Schema> schema;
87+
ASSERT_TRUE(get_arrow_schema_from_block(block, &schema, "UTC").ok());
88+
std::string serialized;
89+
ASSERT_TRUE(serialize_arrow_schema(&schema, &serialized).ok());
90+
auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized));
91+
auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie();
92+
EXPECT_TRUE(schema->Equals(*reader->schema(), true));
93+
const auto& map = static_cast<const arrow::MapType&>(*reader->schema()->field(0)->type());
94+
expect_logical_type(map.key_field(), "LARGEINT");
95+
EXPECT_EQ("key", map.key_field()->name());
96+
EXPECT_FALSE(map.key_field()->nullable());
97+
EXPECT_EQ("value", map.item_field()->name());
98+
auto item = map.item_type()->field(0);
99+
EXPECT_EQ("item", item->name());
100+
EXPECT_TRUE(item->nullable());
101+
expect_logical_type(item->type()->field(0), "LARGEINT");
102+
EXPECT_FALSE(item->type()->field(0)->nullable());
103+
auto physical =
104+
arrow::map(arrow::utf8(),
105+
arrow::list(arrow::struct_({arrow::field("number", arrow::utf8(), false)})));
106+
EXPECT_TRUE(map.Equals(physical));
107+
EXPECT_FALSE(map.Equals(physical, true));
108+
}
109+
110+
TEST(ArrowRowBatchMetadataTest, KeepsNativeTypesUnannotated) {
111+
auto type = std::make_shared<DataTypeStruct>(
112+
DataTypes {std::make_shared<DataTypeInt32>(), std::make_shared<DataTypeString>()},
113+
Strings {"number", "text"});
114+
std::shared_ptr<arrow::DataType> arrow_type;
115+
ASSERT_TRUE(convert_to_arrow_type(type, &arrow_type, "UTC").ok());
116+
for (const auto& field : arrow_type->fields()) {
117+
EXPECT_EQ(nullptr, field->metadata());
118+
EXPECT_FALSE(field->nullable());
119+
}
120+
}
121+
122+
TEST(ArrowRowBatchMetadataTest, PreservesLargeintExtremesAndNullsInRecordBatches) {
123+
auto integer = make_nullable(std::make_shared<DataTypeInt128>());
124+
auto array_type = std::make_shared<DataTypeArray>(integer);
125+
auto column = array_type->create_column();
126+
column->insert(Field::create_field<TYPE_ARRAY>(
127+
Array {Field::create_field<TYPE_LARGEINT>(MAX_INT128),
128+
Field::create_field<TYPE_LARGEINT>(MIN_INT128), Field()}));
129+
column->insert(Field::create_field<TYPE_ARRAY>(Array {}));
130+
Block block;
131+
block.insert({std::move(column), array_type, "numbers"});
132+
std::shared_ptr<arrow::Schema> schema;
133+
ASSERT_TRUE(get_arrow_schema_from_block(block, &schema, "UTC", true).ok());
134+
ArrowFlightArrowBlockConvertor converter(schema, cctz::utc_time_zone());
135+
std::shared_ptr<arrow::RecordBatch> batch;
136+
auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch);
137+
ASSERT_TRUE(status.ok()) << status;
138+
ASSERT_TRUE(batch->ValidateFull().ok());
139+
std::string serialized;
140+
ASSERT_TRUE(serialize_record_batch(*batch, &serialized).ok());
141+
auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized));
142+
auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie();
143+
ASSERT_TRUE(reader->ReadNext(&batch).ok());
144+
EXPECT_TRUE(batch->schema()->Equals(*schema, true));
145+
expect_logical_type(batch->schema()->field(0)->type()->field(0), "LARGEINT");
146+
const auto& lists = static_cast<const arrow::ListArray&>(*batch->column(0));
147+
EXPECT_EQ(0, lists.value_length(1));
148+
const auto& values = static_cast<const arrow::StringArray&>(*lists.values());
149+
EXPECT_EQ("170141183460469231731687303715884105727", values.GetString(0));
150+
EXPECT_EQ("-170141183460469231731687303715884105728", values.GetString(1));
151+
EXPECT_TRUE(values.IsNull(2));
152+
}
153+
154+
TEST(ArrowRowBatchMetadataTest, PreservesNestedMapMetadataAndValuesInRecordBatches) {
155+
auto integer = make_nullable(std::make_shared<DataTypeInt128>());
156+
auto map_type = make_nullable(std::make_shared<DataTypeMap>(integer, integer));
157+
auto type = make_nullable(std::make_shared<DataTypeStruct>(
158+
DataTypes {map_type, std::make_shared<DataTypeString>()}, Strings {"mapping", "text"}));
159+
Map map;
160+
map.push_back(Field::create_field<TYPE_ARRAY>(
161+
Array {Field::create_field<TYPE_LARGEINT>(MAX_INT128)}));
162+
map.push_back(Field::create_field<TYPE_ARRAY>(
163+
Array {Field::create_field<TYPE_LARGEINT>(MIN_INT128)}));
164+
auto column = type->create_column();
165+
column->insert(Field::create_field<TYPE_STRUCT>(
166+
Struct {Field::create_field<TYPE_MAP>(map), Field::create_field<TYPE_STRING>("17")}));
167+
column->insert_default();
168+
Block block;
169+
block.insert({std::move(column), type, "nested"});
170+
std::shared_ptr<arrow::Schema> schema;
171+
ASSERT_TRUE(get_arrow_schema_from_block(block, &schema, "UTC", true).ok());
172+
ArrowFlightArrowBlockConvertor converter(schema, cctz::utc_time_zone());
173+
std::shared_ptr<arrow::RecordBatch> batch;
174+
auto status = converter.convert_to_arrow(block, arrow::default_memory_pool(), &batch);
175+
ASSERT_TRUE(status.ok()) << status;
176+
ASSERT_TRUE(batch->ValidateFull().ok());
177+
std::string serialized;
178+
ASSERT_TRUE(serialize_record_batch(*batch, &serialized).ok());
179+
auto source = std::make_shared<arrow::io::BufferReader>(arrow::Buffer::FromString(serialized));
180+
auto reader = arrow::ipc::RecordBatchStreamReader::Open(source).ValueOrDie();
181+
ASSERT_TRUE(reader->ReadNext(&batch).ok());
182+
const auto& structure = static_cast<const arrow::StructArray&>(*batch->column(0));
183+
EXPECT_TRUE(structure.IsNull(1));
184+
const auto& maps = static_cast<const arrow::MapArray&>(*structure.field(0));
185+
const auto& arrow_map_type = static_cast<const arrow::MapType&>(*maps.type());
186+
expect_logical_type(arrow_map_type.key_field(), "LARGEINT");
187+
expect_logical_type(arrow_map_type.item_field(), "LARGEINT");
188+
EXPECT_FALSE(arrow_map_type.key_field()->nullable());
189+
EXPECT_EQ("170141183460469231731687303715884105727",
190+
static_cast<const arrow::StringArray&>(*maps.keys()).GetString(0));
191+
EXPECT_EQ("-170141183460469231731687303715884105728",
192+
static_cast<const arrow::StringArray&>(*maps.items()).GetString(0));
193+
EXPECT_EQ("17", static_cast<const arrow::StringArray&>(*structure.field(1)).GetString(0));
194+
EXPECT_EQ(nullptr, structure.type()->field(1)->metadata());
195+
}
196+
197+
} // namespace
198+
} // namespace doris

‎samples/arrow-flight-sql/python/README.md‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,22 @@ under the License.
3535
Section 6.1 of https://github.com/apache/doris/issues/25514 is the performance test
3636
results of the doris arrow flight sql using python.
3737

38+
## Logical type metadata
39+
40+
Some Doris types share an Arrow storage type with ordinary strings or integers.
41+
The `doris_type` field metadata identifies `LARGEINT`, `IPV4`, `IPV6`, `JSON`, and
42+
`VARIANT`, including fields nested in arrays, maps, and structs. Consumers should
43+
inspect each nested Arrow field instead of inferring a type from its value.
44+
45+
LARGEINT retains its Arrow string encoding, including the full signed 128-bit
46+
range. PyArrow does not automatically convert custom metadata into Python types;
47+
a client can use `doris_type=LARGEINT` to safely convert that field's non-NULL
48+
values with `int(value)`, while leaving ordinary STRING fields unchanged.
49+
50+
Run `python test_nested_type_metadata.py` with `DORIS_FLIGHT_SQL_URI` and optional
51+
`DORIS_USER` / `DORIS_PASSWORD` to verify the metadata and values against a cluster.
52+
The tests execute only read-only queries.
53+
3854
# Notes
3955

4056
For more details, refer to [Python Usage] in the document https://doris.apache.org/zh-CN/docs/dev/db-connect/arrow-flight-sql-connect

0 commit comments

Comments
 (0)