be/src/format/arrow/arrow_row_batch.cpp
Line | Count | Source |
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/buffer.h> |
21 | | #include <arrow/io/memory.h> |
22 | | #include <arrow/ipc/writer.h> |
23 | | #include <arrow/record_batch.h> |
24 | | #include <arrow/result.h> |
25 | | #include <arrow/status.h> |
26 | | #include <arrow/type.h> |
27 | | #include <arrow/type_fwd.h> |
28 | | #include <arrow/util/key_value_metadata.h> |
29 | | #include <glog/logging.h> |
30 | | #include <stdint.h> |
31 | | |
32 | | #include <algorithm> |
33 | | #include <cstdlib> |
34 | | #include <memory> |
35 | | #include <utility> |
36 | | #include <vector> |
37 | | |
38 | | #include "core/block/block.h" |
39 | | #include "core/data_type/data_type_agg_state.h" |
40 | | #include "core/data_type/data_type_array.h" |
41 | | #include "core/data_type/data_type_map.h" |
42 | | #include "core/data_type/data_type_struct.h" |
43 | | #include "core/data_type/define_primitive_type.h" |
44 | | #include "exprs/vexpr.h" |
45 | | #include "exprs/vexpr_context.h" |
46 | | #include "format/arrow/arrow_block_convertor.h" |
47 | | #include "runtime/descriptors.h" |
48 | | |
49 | | namespace doris { |
50 | | |
51 | | Status convert_to_arrow_type(const DataTypePtr& origin_type, |
52 | | std::shared_ptr<arrow::DataType>* result, const std::string& timezone, |
53 | 575 | bool datetime_naive) { |
54 | 575 | auto type = get_serialized_type(origin_type); |
55 | 575 | switch (type->get_primitive_type()) { |
56 | 0 | case TYPE_NULL: |
57 | 0 | *result = arrow::null(); |
58 | 0 | break; |
59 | 9 | case TYPE_TINYINT: |
60 | 9 | *result = arrow::int8(); |
61 | 9 | break; |
62 | 9 | case TYPE_SMALLINT: |
63 | 9 | *result = arrow::int16(); |
64 | 9 | break; |
65 | 38 | case TYPE_INT: |
66 | 38 | *result = arrow::int32(); |
67 | 38 | break; |
68 | 9 | case TYPE_BIGINT: |
69 | 9 | *result = arrow::int64(); |
70 | 9 | break; |
71 | 9 | case TYPE_FLOAT: |
72 | 9 | *result = arrow::float32(); |
73 | 9 | break; |
74 | 22 | case TYPE_DOUBLE: |
75 | 22 | *result = arrow::float64(); |
76 | 22 | break; |
77 | 0 | case TYPE_TIMEV2: |
78 | 0 | *result = arrow::float64(); |
79 | 0 | break; |
80 | 25 | case TYPE_IPV4: |
81 | | // ipv4 is uint32, but parquet not uint32, it's will be convert to int64 |
82 | | // so use int32 directly |
83 | 25 | *result = arrow::int32(); |
84 | 25 | break; |
85 | 20 | case TYPE_IPV6: |
86 | 20 | *result = arrow::utf8(); |
87 | 20 | break; |
88 | 23 | case TYPE_LARGEINT: |
89 | 23 | case TYPE_VARCHAR: |
90 | 23 | case TYPE_CHAR: |
91 | 40 | case TYPE_DATE: |
92 | 53 | case TYPE_DATETIME: |
93 | 101 | case TYPE_STRING: |
94 | 101 | case TYPE_JSONB: |
95 | 101 | *result = arrow::utf8(); |
96 | 101 | break; |
97 | 13 | case TYPE_DATEV2: |
98 | 13 | *result = std::make_shared<arrow::Date32Type>(); |
99 | 13 | break; |
100 | 1 | case TYPE_TIMESTAMPTZ: |
101 | 22 | case TYPE_DATETIMEV2: { |
102 | 22 | arrow::TimeUnit::type time_unit; |
103 | 22 | if (type->get_scale() > 3) { |
104 | 4 | time_unit = arrow::TimeUnit::MICRO; |
105 | 18 | } else if (type->get_scale() > 0) { |
106 | 4 | time_unit = arrow::TimeUnit::MILLI; |
107 | 14 | } else { |
108 | 14 | time_unit = arrow::TimeUnit::SECOND; |
109 | 14 | } |
110 | | // Doris DATETIMEV2 represents a wall-clock value without a timezone. Arrow Flight |
111 | | // exposes it as a timezone-naive timestamp so clients do not interpret it as an instant. |
112 | | // This option only changes the DATETIMEV2 output schema. TIMESTAMPTZ remains timezone-aware, |
113 | | // and Arrow-to-Doris conversions are unaffected. |
114 | 22 | if (type->get_primitive_type() == TYPE_DATETIMEV2 && datetime_naive) { |
115 | 2 | *result = std::make_shared<arrow::TimestampType>(time_unit); |
116 | 20 | } else { |
117 | 20 | *result = std::make_shared<arrow::TimestampType>(time_unit, timezone); |
118 | 20 | } |
119 | 22 | break; |
120 | 1 | } |
121 | 5 | case TYPE_DECIMALV2: |
122 | 27 | case TYPE_DECIMAL32: |
123 | 43 | case TYPE_DECIMAL64: |
124 | 55 | case TYPE_DECIMAL128I: |
125 | 55 | *result = std::make_shared<arrow::Decimal128Type>(type->get_precision(), type->get_scale()); |
126 | 55 | break; |
127 | 16 | case TYPE_DECIMAL256: |
128 | 16 | *result = std::make_shared<arrow::Decimal256Type>(type->get_precision(), type->get_scale()); |
129 | 16 | break; |
130 | 9 | case TYPE_BOOLEAN: |
131 | 9 | *result = arrow::boolean(); |
132 | 9 | break; |
133 | 101 | case TYPE_ARRAY: { |
134 | 101 | const auto* type_arr = assert_cast<const DataTypeArray*>(remove_nullable(type).get()); |
135 | 101 | std::shared_ptr<arrow::DataType> item_type; |
136 | 101 | RETURN_IF_ERROR(convert_to_arrow_type(type_arr->get_nested_type(), &item_type, timezone, |
137 | 101 | datetime_naive)); |
138 | 101 | *result = std::make_shared<arrow::ListType>(item_type); |
139 | 101 | break; |
140 | 101 | } |
141 | 62 | case TYPE_MAP: { |
142 | 62 | const auto* type_map = assert_cast<const DataTypeMap*>(remove_nullable(type).get()); |
143 | 62 | std::shared_ptr<arrow::DataType> key_type; |
144 | 62 | std::shared_ptr<arrow::DataType> val_type; |
145 | 62 | RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_key_type(), &key_type, timezone, |
146 | 62 | datetime_naive)); |
147 | 62 | RETURN_IF_ERROR(convert_to_arrow_type(type_map->get_value_type(), &val_type, timezone, |
148 | 62 | datetime_naive)); |
149 | 62 | *result = std::make_shared<arrow::MapType>(key_type, val_type); |
150 | 62 | break; |
151 | 62 | } |
152 | 53 | case TYPE_STRUCT: { |
153 | 53 | const auto* type_struct = assert_cast<const DataTypeStruct*>(remove_nullable(type).get()); |
154 | 53 | std::vector<std::shared_ptr<arrow::Field>> fields; |
155 | 175 | for (size_t i = 0; i < type_struct->get_elements().size(); i++) { |
156 | 122 | std::shared_ptr<arrow::DataType> field_type; |
157 | 122 | RETURN_IF_ERROR(convert_to_arrow_type(type_struct->get_element(i), &field_type, |
158 | 122 | timezone, datetime_naive)); |
159 | 122 | fields.push_back( |
160 | 122 | std::make_shared<arrow::Field>(type_struct->get_element_name(i), field_type, |
161 | 122 | type_struct->get_element(i)->is_nullable())); |
162 | 122 | } |
163 | 53 | *result = std::make_shared<arrow::StructType>(fields); |
164 | 53 | break; |
165 | 53 | } |
166 | 0 | case TYPE_VARIANT: { |
167 | 0 | *result = arrow::utf8(); |
168 | 0 | break; |
169 | 53 | } |
170 | 0 | case TYPE_QUANTILE_STATE: |
171 | 1 | case TYPE_BITMAP: |
172 | 2 | case TYPE_HLL: { |
173 | 2 | *result = arrow::binary(); |
174 | 2 | break; |
175 | 1 | } |
176 | 0 | case TYPE_VARBINARY: { |
177 | 0 | *result = arrow::binary(); |
178 | 0 | break; |
179 | 1 | } |
180 | 0 | default: |
181 | 0 | return Status::InvalidArgument("Unknown primitive type({}) convert to Arrow type", |
182 | 0 | type->get_name()); |
183 | 575 | } |
184 | 575 | return Status::OK(); |
185 | 575 | } |
186 | | |
187 | | // Helper function to create an Arrow Field with type metadata if applicable, such as IP types |
188 | | std::shared_ptr<arrow::Field> create_arrow_field_with_metadata( |
189 | | const std::string& field_name, const std::shared_ptr<arrow::DataType>& arrow_type, |
190 | 208 | bool is_nullable, PrimitiveType primitive_type) { |
191 | 208 | if (primitive_type == PrimitiveType::TYPE_IPV4) { |
192 | 4 | auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV4"}); |
193 | 4 | return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); |
194 | 204 | } else if (primitive_type == PrimitiveType::TYPE_IPV6) { |
195 | 4 | auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"IPV6"}); |
196 | 4 | return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); |
197 | 200 | } else if (primitive_type == PrimitiveType::TYPE_LARGEINT) { |
198 | 4 | auto metadata = arrow::KeyValueMetadata::Make({"doris_type"}, {"LARGEINT"}); |
199 | 4 | return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable, metadata); |
200 | 196 | } else { |
201 | 196 | return std::make_shared<arrow::Field>(field_name, arrow_type, is_nullable); |
202 | 196 | } |
203 | 208 | } |
204 | | |
205 | | Status get_arrow_schema_from_block(const Block& block, std::shared_ptr<arrow::Schema>* result, |
206 | 30 | const std::string& timezone) { |
207 | 30 | std::vector<std::shared_ptr<arrow::Field>> fields; |
208 | 208 | for (const auto& type_and_name : block) { |
209 | 208 | std::shared_ptr<arrow::DataType> arrow_type; |
210 | 208 | RETURN_IF_ERROR(convert_to_arrow_type(type_and_name.type, &arrow_type, timezone)); |
211 | 208 | auto field = create_arrow_field_with_metadata(type_and_name.name, arrow_type, |
212 | 208 | type_and_name.type->is_nullable(), |
213 | 208 | type_and_name.type->get_primitive_type()); |
214 | 208 | fields.push_back(field); |
215 | 208 | } |
216 | 30 | *result = arrow::schema(std::move(fields)); |
217 | 30 | return Status::OK(); |
218 | 30 | } |
219 | | |
220 | | Status get_arrow_schema_from_expr_ctxs(const VExprContextSPtrs& output_vexpr_ctxs, |
221 | | std::shared_ptr<arrow::Schema>* result, |
222 | 0 | const std::string& timezone, bool datetime_naive) { |
223 | 0 | std::vector<std::shared_ptr<arrow::Field>> fields; |
224 | 0 | for (int i = 0; i < output_vexpr_ctxs.size(); i++) { |
225 | 0 | std::shared_ptr<arrow::DataType> arrow_type; |
226 | 0 | auto root_expr = output_vexpr_ctxs.at(i)->root(); |
227 | 0 | RETURN_IF_ERROR(convert_to_arrow_type(root_expr->data_type(), &arrow_type, timezone, |
228 | 0 | datetime_naive)); |
229 | 0 | auto field_name = root_expr->is_slot_ref() && !root_expr->expr_label().empty() |
230 | 0 | ? root_expr->expr_label() |
231 | 0 | : fmt::format("{}_{}", root_expr->data_type()->get_name(), i); |
232 | 0 | auto field = |
233 | 0 | create_arrow_field_with_metadata(field_name, arrow_type, root_expr->is_nullable(), |
234 | 0 | root_expr->data_type()->get_primitive_type()); |
235 | 0 | fields.push_back(field); |
236 | 0 | } |
237 | 0 | *result = arrow::schema(std::move(fields)); |
238 | 0 | return Status::OK(); |
239 | 0 | } |
240 | | |
241 | 0 | Status serialize_record_batch(const arrow::RecordBatch& record_batch, std::string* result) { |
242 | | // create sink memory buffer outputstream with the computed capacity |
243 | 0 | int64_t capacity; |
244 | 0 | arrow::Status a_st = arrow::ipc::GetRecordBatchSize(record_batch, &capacity); |
245 | 0 | if (!a_st.ok()) { |
246 | 0 | return Status::InternalError("GetRecordBatchSize failure, reason: {}", a_st.ToString()); |
247 | 0 | } |
248 | 0 | auto sink_res = arrow::io::BufferOutputStream::Create(capacity, arrow::default_memory_pool()); |
249 | 0 | if (!sink_res.ok()) { |
250 | 0 | return Status::InternalError("create BufferOutputStream failure, reason: {}", |
251 | 0 | sink_res.status().ToString()); |
252 | 0 | } |
253 | 0 | std::shared_ptr<arrow::io::BufferOutputStream> sink = sink_res.ValueOrDie(); |
254 | | // create RecordBatch Writer |
255 | 0 | auto res = arrow::ipc::MakeStreamWriter(sink.get(), record_batch.schema()); |
256 | 0 | if (!res.ok()) { |
257 | 0 | return Status::InternalError("open RecordBatchStreamWriter failure, reason: {}", |
258 | 0 | res.status().ToString()); |
259 | 0 | } |
260 | | // write RecordBatch to memory buffer outputstream |
261 | 0 | std::shared_ptr<arrow::ipc::RecordBatchWriter> record_batch_writer = res.ValueOrDie(); |
262 | 0 | a_st = record_batch_writer->WriteRecordBatch(record_batch); |
263 | 0 | if (!a_st.ok()) { |
264 | 0 | return Status::InternalError("write record batch failure, reason: {}", a_st.ToString()); |
265 | 0 | } |
266 | 0 | a_st = record_batch_writer->Close(); |
267 | 0 | if (!a_st.ok()) { |
268 | 0 | return Status::InternalError("Close failed, reason: {}", a_st.ToString()); |
269 | 0 | } |
270 | 0 | auto finish_res = sink->Finish(); |
271 | 0 | if (!finish_res.ok()) { |
272 | 0 | return Status::InternalError("allocate result buffer failure, reason: {}", |
273 | 0 | finish_res.status().ToString()); |
274 | 0 | } |
275 | 0 | *result = finish_res.ValueOrDie()->ToString(); |
276 | | // close the sink |
277 | 0 | a_st = sink->Close(); |
278 | 0 | if (!a_st.ok()) { |
279 | 0 | return Status::InternalError("Close failed, reason: {}", a_st.ToString()); |
280 | 0 | } |
281 | 0 | return Status::OK(); |
282 | 0 | } |
283 | | |
284 | 0 | Status serialize_arrow_schema(std::shared_ptr<arrow::Schema>* schema, std::string* result) { |
285 | 0 | auto make_empty_result = arrow::RecordBatch::MakeEmpty(*schema); |
286 | 0 | if (!make_empty_result.ok()) { |
287 | 0 | return Status::InternalError("serialize_arrow_schema failed, reason: {}", |
288 | 0 | make_empty_result.status().ToString()); |
289 | 0 | } |
290 | 0 | auto batch = make_empty_result.ValueOrDie(); |
291 | 0 | return serialize_record_batch(*batch, result); |
292 | 0 | } |
293 | | |
294 | | } // namespace doris |