be/src/exec/sink/viceberg_merge_sink.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 "exec/sink/viceberg_merge_sink.h" |
19 | | |
20 | | #include <fmt/format.h> |
21 | | |
22 | | #include "common/consts.h" |
23 | | #include "common/exception.h" |
24 | | #include "common/logging.h" |
25 | | #include "core/block/block.h" |
26 | | #include "core/column/column_nullable.h" |
27 | | #include "core/column/column_vector.h" |
28 | | #include "exec/sink/sink_common.h" |
29 | | #include "exec/sink/viceberg_delete_sink.h" |
30 | | #include "exec/sink/writer/iceberg/viceberg_table_writer.h" |
31 | | #include "format/table/iceberg/schema.h" |
32 | | #include "format/table/iceberg/schema_parser.h" |
33 | | #include "runtime/runtime_state.h" |
34 | | #include "util/string_util.h" |
35 | | |
36 | | namespace doris { |
37 | | |
38 | | namespace {} // namespace |
39 | | |
40 | | VIcebergMergeSink::VIcebergMergeSink(const TDataSink& t_sink, const VExprContextSPtrs& output_exprs, |
41 | | std::shared_ptr<Dependency> dep, |
42 | | std::shared_ptr<Dependency> fin_dep) |
43 | 6 | : AsyncResultWriter(output_exprs, dep, fin_dep), _t_sink(t_sink) { |
44 | 6 | DCHECK(_t_sink.__isset.iceberg_merge_sink); |
45 | 6 | } |
46 | | |
47 | 6 | VIcebergMergeSink::~VIcebergMergeSink() = default; |
48 | | |
49 | 6 | Status VIcebergMergeSink::init_properties(ObjectPool* pool, const RowDescriptor& row_desc) { |
50 | 6 | RETURN_IF_ERROR(_build_inner_sinks()); |
51 | | |
52 | 6 | _table_writer = std::make_unique<VIcebergTableWriter>(_table_sink, _table_output_expr_ctxs, |
53 | 6 | nullptr, nullptr); |
54 | 6 | _delete_writer = std::make_unique<VIcebergDeleteSink>(_delete_sink, _delete_output_expr_ctxs, |
55 | 6 | nullptr, nullptr); |
56 | 6 | RETURN_IF_ERROR(_table_writer->init_properties(pool, row_desc)); |
57 | 6 | RETURN_IF_ERROR(_delete_writer->init_properties(pool)); |
58 | 6 | return Status::OK(); |
59 | 6 | } |
60 | | |
61 | 6 | Status VIcebergMergeSink::open(RuntimeState* state, RuntimeProfile* profile) { |
62 | 6 | _state = state; |
63 | | |
64 | 6 | _written_rows_counter = ADD_COUNTER(profile, "RowsWritten", TUnit::UNIT); |
65 | 6 | _insert_rows_counter = ADD_COUNTER(profile, "InsertRows", TUnit::UNIT); |
66 | 6 | _delete_rows_counter = ADD_COUNTER(profile, "DeleteRows", TUnit::UNIT); |
67 | 6 | _send_data_timer = ADD_TIMER(profile, "SendDataTime"); |
68 | 6 | _open_timer = ADD_TIMER(profile, "OpenTime"); |
69 | 6 | _close_timer = ADD_TIMER(profile, "CloseTime"); |
70 | | |
71 | 6 | SCOPED_TIMER(_open_timer); |
72 | | |
73 | 6 | RETURN_IF_ERROR(_prepare_output_layout()); |
74 | | |
75 | 4 | RuntimeProfile* table_profile = profile->create_child("IcebergMergeTableWriter", true, true); |
76 | 4 | RuntimeProfile* delete_profile = profile->create_child("IcebergMergeDeleteWriter", true, true); |
77 | | |
78 | 4 | RETURN_IF_ERROR(_table_writer->open(state, table_profile)); |
79 | 3 | RETURN_IF_ERROR(_delete_writer->open(state, delete_profile)); |
80 | | |
81 | 3 | return Status::OK(); |
82 | 3 | } |
83 | | |
84 | 3 | Status VIcebergMergeSink::write(RuntimeState* state, Block& block) { |
85 | 3 | SCOPED_TIMER(_send_data_timer); |
86 | 3 | if (block.rows() == 0) { |
87 | 0 | return Status::OK(); |
88 | 0 | } |
89 | | |
90 | 3 | Block output_block; |
91 | 3 | RETURN_IF_ERROR(_projection_block(block, &output_block)); |
92 | 3 | if (output_block.rows() == 0) { |
93 | 0 | return Status::OK(); |
94 | 0 | } |
95 | | |
96 | 3 | _row_count += output_block.rows(); |
97 | | |
98 | 3 | if (_operation_idx < 0 || _row_id_idx < 0) { |
99 | 0 | return Status::InternalError("Iceberg merge sink missing operation/row_id columns"); |
100 | 0 | } |
101 | | |
102 | 3 | const auto& op_column = output_block.get_by_position(_operation_idx).column; |
103 | 3 | const auto* op_data = remove_nullable(op_column).get(); |
104 | | |
105 | 3 | IColumn::Filter delete_filter(output_block.rows(), 0); |
106 | 3 | IColumn::Filter insert_filter(output_block.rows(), 0); |
107 | 3 | bool has_delete = false; |
108 | 3 | bool has_insert = false; |
109 | 3 | size_t delete_rows = 0; |
110 | 3 | size_t insert_rows = 0; |
111 | | |
112 | 8 | for (size_t i = 0; i < output_block.rows(); ++i) { |
113 | 6 | int8_t op = static_cast<int8_t>(op_data->get_int(i)); |
114 | 6 | bool delete_op = is_delete_op(op); |
115 | 6 | bool insert_op = is_insert_op(op); |
116 | 6 | if (!delete_op && !insert_op) { |
117 | 1 | return Status::InternalError("Unknown Iceberg merge operation {}", op); |
118 | 1 | } |
119 | 5 | if (delete_op) { |
120 | 3 | delete_filter[i] = 1; |
121 | 3 | has_delete = true; |
122 | 3 | ++_delete_row_count; |
123 | 3 | ++delete_rows; |
124 | 3 | } |
125 | 5 | if (insert_op) { |
126 | 3 | insert_filter[i] = 1; |
127 | 3 | has_insert = true; |
128 | 3 | ++_insert_row_count; |
129 | 3 | ++insert_rows; |
130 | 3 | } |
131 | 5 | } |
132 | | |
133 | 2 | bool skip_io = false; |
134 | 2 | #ifdef BE_TEST |
135 | 2 | skip_io = _skip_io; |
136 | 2 | #endif |
137 | | |
138 | 2 | if (has_delete && !skip_io) { |
139 | 0 | Block delete_block = output_block; |
140 | 0 | std::vector<int> delete_indices {_row_id_idx}; |
141 | 0 | delete_block.erase_not_in(delete_indices); |
142 | 0 | Block::filter_block_internal(&delete_block, delete_filter); |
143 | 0 | RETURN_IF_ERROR(_delete_writer->write(state, delete_block)); |
144 | 0 | } |
145 | | |
146 | 2 | if (has_insert && !skip_io) { |
147 | 0 | if (_data_column_indices.empty()) { |
148 | 0 | return Status::InternalError("Iceberg merge sink has no data columns for insert"); |
149 | 0 | } |
150 | 0 | Block insert_block = output_block; |
151 | 0 | insert_block.erase_not_in(_data_column_indices); |
152 | 0 | Block::filter_block_internal(&insert_block, insert_filter); |
153 | 0 | RETURN_IF_ERROR(_table_writer->write_prepared_block(insert_block)); |
154 | 0 | } |
155 | | |
156 | 2 | if (_written_rows_counter != nullptr) { |
157 | 2 | COUNTER_UPDATE(_written_rows_counter, output_block.rows()); |
158 | 2 | } |
159 | 2 | if (_insert_rows_counter != nullptr) { |
160 | 2 | COUNTER_UPDATE(_insert_rows_counter, insert_rows); |
161 | 2 | } |
162 | 2 | if (_delete_rows_counter != nullptr) { |
163 | 2 | COUNTER_UPDATE(_delete_rows_counter, delete_rows); |
164 | 2 | } |
165 | | |
166 | 2 | return Status::OK(); |
167 | 2 | } |
168 | | |
169 | 2 | Status VIcebergMergeSink::close(Status close_status) { |
170 | 2 | SCOPED_TIMER(_close_timer); |
171 | | |
172 | 2 | if (!close_status.ok()) { |
173 | 0 | LOG(WARNING) << fmt::format("VIcebergMergeSink close with error: {}", |
174 | 0 | close_status.to_string()); |
175 | 0 | if (_table_writer) { |
176 | 0 | static_cast<void>(_table_writer->close(close_status)); |
177 | 0 | } |
178 | 0 | if (_delete_writer) { |
179 | 0 | static_cast<void>(_delete_writer->close(close_status)); |
180 | 0 | } |
181 | 0 | return close_status; |
182 | 0 | } |
183 | | |
184 | 2 | Status table_status = Status::OK(); |
185 | 2 | Status delete_status = Status::OK(); |
186 | 2 | if (_table_writer) { |
187 | 2 | table_status = _table_writer->close(close_status); |
188 | 2 | } |
189 | 2 | if (_delete_writer) { |
190 | 2 | delete_status = _delete_writer->close(close_status); |
191 | 2 | } |
192 | | |
193 | 2 | if (_written_rows_counter != nullptr) { |
194 | 2 | COUNTER_SET(_written_rows_counter, static_cast<int64_t>(_row_count)); |
195 | 2 | } |
196 | 2 | if (_insert_rows_counter != nullptr) { |
197 | 2 | COUNTER_SET(_insert_rows_counter, static_cast<int64_t>(_insert_row_count)); |
198 | 2 | } |
199 | 2 | if (_delete_rows_counter != nullptr) { |
200 | 2 | COUNTER_SET(_delete_rows_counter, static_cast<int64_t>(_delete_row_count)); |
201 | 2 | } |
202 | | |
203 | 2 | if (!table_status.ok()) { |
204 | 0 | return table_status; |
205 | 0 | } |
206 | 2 | return delete_status; |
207 | 2 | } |
208 | | |
209 | 6 | Status VIcebergMergeSink::_build_inner_sinks() { |
210 | 6 | if (!_t_sink.__isset.iceberg_merge_sink) { |
211 | 0 | return Status::InternalError("Missing iceberg merge sink config"); |
212 | 0 | } |
213 | | |
214 | 6 | const auto& merge_sink = _t_sink.iceberg_merge_sink; |
215 | | |
216 | 6 | TIcebergTableSink table_sink; |
217 | 6 | if (merge_sink.__isset.db_name) { |
218 | 6 | table_sink.__set_db_name(merge_sink.db_name); |
219 | 6 | } |
220 | 6 | if (merge_sink.__isset.tb_name) { |
221 | 6 | table_sink.__set_tb_name(merge_sink.tb_name); |
222 | 6 | } |
223 | 6 | if (merge_sink.__isset.schema_json) { |
224 | 6 | table_sink.__set_schema_json(merge_sink.schema_json); |
225 | 6 | } |
226 | 6 | if (merge_sink.__isset.partition_specs_json) { |
227 | 0 | table_sink.__set_partition_specs_json(merge_sink.partition_specs_json); |
228 | 0 | } |
229 | 6 | if (merge_sink.__isset.partition_spec_id) { |
230 | 6 | table_sink.__set_partition_spec_id(merge_sink.partition_spec_id); |
231 | 6 | } |
232 | 6 | if (merge_sink.__isset.sort_fields) { |
233 | 0 | table_sink.__set_sort_fields(merge_sink.sort_fields); |
234 | 0 | } |
235 | 6 | if (merge_sink.__isset.file_format) { |
236 | 6 | table_sink.__set_file_format(merge_sink.file_format); |
237 | 6 | } |
238 | 6 | if (merge_sink.__isset.compression_type) { |
239 | 6 | table_sink.__set_compression_type(merge_sink.compression_type); |
240 | 6 | } |
241 | 6 | if (merge_sink.__isset.output_path) { |
242 | 6 | table_sink.__set_output_path(merge_sink.output_path); |
243 | 6 | } |
244 | 6 | if (merge_sink.__isset.original_output_path) { |
245 | 6 | table_sink.__set_original_output_path(merge_sink.original_output_path); |
246 | 6 | } |
247 | 6 | if (merge_sink.__isset.hadoop_config) { |
248 | 0 | table_sink.__set_hadoop_config(merge_sink.hadoop_config); |
249 | 0 | } |
250 | 6 | if (merge_sink.__isset.file_type) { |
251 | 6 | table_sink.__set_file_type(merge_sink.file_type); |
252 | 6 | } |
253 | 6 | if (merge_sink.__isset.broker_addresses) { |
254 | 0 | table_sink.__set_broker_addresses(merge_sink.broker_addresses); |
255 | 0 | } |
256 | 6 | _table_sink.__set_type(TDataSinkType::ICEBERG_TABLE_SINK); |
257 | 6 | _table_sink.__set_iceberg_table_sink(table_sink); |
258 | | |
259 | 6 | TIcebergDeleteSink delete_sink; |
260 | 6 | if (merge_sink.__isset.db_name) { |
261 | 6 | delete_sink.__set_db_name(merge_sink.db_name); |
262 | 6 | } |
263 | 6 | if (merge_sink.__isset.tb_name) { |
264 | 6 | delete_sink.__set_tb_name(merge_sink.tb_name); |
265 | 6 | } |
266 | 6 | if (merge_sink.__isset.delete_type) { |
267 | 6 | delete_sink.__set_delete_type(merge_sink.delete_type); |
268 | 6 | } |
269 | 6 | if (merge_sink.__isset.file_format) { |
270 | 6 | delete_sink.__set_file_format(merge_sink.file_format); |
271 | 6 | } |
272 | 6 | if (merge_sink.__isset.compression_type) { |
273 | 6 | delete_sink.__set_compress_type(merge_sink.compression_type); |
274 | 6 | } |
275 | 6 | if (merge_sink.__isset.output_path) { |
276 | 6 | delete_sink.__set_output_path(merge_sink.output_path); |
277 | 6 | } |
278 | 6 | if (merge_sink.__isset.table_location) { |
279 | 6 | delete_sink.__set_table_location(merge_sink.table_location); |
280 | 6 | } |
281 | 6 | if (merge_sink.__isset.hadoop_config) { |
282 | 0 | delete_sink.__set_hadoop_config(merge_sink.hadoop_config); |
283 | 0 | } |
284 | 6 | if (merge_sink.__isset.file_type) { |
285 | 6 | delete_sink.__set_file_type(merge_sink.file_type); |
286 | 6 | } |
287 | 6 | if (merge_sink.__isset.partition_spec_id_for_delete) { |
288 | 6 | delete_sink.__set_partition_spec_id(merge_sink.partition_spec_id_for_delete); |
289 | 6 | } |
290 | 6 | if (merge_sink.__isset.partition_data_json_for_delete) { |
291 | 0 | delete_sink.__set_partition_data_json(merge_sink.partition_data_json_for_delete); |
292 | 0 | } |
293 | 6 | if (merge_sink.__isset.broker_addresses) { |
294 | 0 | delete_sink.__set_broker_addresses(merge_sink.broker_addresses); |
295 | 0 | } |
296 | 6 | _delete_sink.__set_type(TDataSinkType::ICEBERG_DELETE_SINK); |
297 | 6 | _delete_sink.__set_iceberg_delete_sink(delete_sink); |
298 | | |
299 | 6 | return Status::OK(); |
300 | 6 | } |
301 | | |
302 | 6 | Status VIcebergMergeSink::_prepare_output_layout() { |
303 | 6 | if (_vec_output_expr_ctxs.empty()) { |
304 | 0 | return Status::InternalError("Iceberg merge sink has empty output expressions"); |
305 | 0 | } |
306 | | |
307 | 6 | std::string row_id_name = doris::to_lower(BeConsts::ICEBERG_ROWID_COL); |
308 | 6 | std::string op_name = doris::to_lower(kOperationColumnName); |
309 | | |
310 | 6 | _operation_idx = -1; |
311 | 6 | _row_id_idx = -1; |
312 | 28 | for (size_t i = 0; i < _vec_output_expr_ctxs.size(); ++i) { |
313 | 22 | std::string expr_name = doris::to_lower(_vec_output_expr_ctxs[i]->expr_name()); |
314 | 22 | if (_operation_idx < 0 && expr_name == op_name) { |
315 | 5 | _operation_idx = static_cast<int>(i); |
316 | 17 | } else if (_row_id_idx < 0 && expr_name == row_id_name) { |
317 | 5 | _row_id_idx = static_cast<int>(i); |
318 | 5 | } |
319 | 22 | } |
320 | | |
321 | 6 | if (_operation_idx < 0) { |
322 | 1 | return Status::InternalError("Iceberg merge sink missing operation column"); |
323 | 1 | } |
324 | 5 | if (_row_id_idx < 0) { |
325 | 1 | return Status::InternalError("Iceberg merge sink missing row_id column"); |
326 | 1 | } |
327 | | |
328 | 4 | _data_column_indices.clear(); |
329 | 4 | _table_output_expr_ctxs.clear(); |
330 | 20 | for (size_t i = 0; i < _vec_output_expr_ctxs.size(); ++i) { |
331 | 16 | if (static_cast<int>(i) == _operation_idx || static_cast<int>(i) == _row_id_idx) { |
332 | 8 | continue; |
333 | 8 | } |
334 | 8 | _data_column_indices.push_back(static_cast<int>(i)); |
335 | 8 | _table_output_expr_ctxs.emplace_back(_vec_output_expr_ctxs[i]); |
336 | 8 | } |
337 | | |
338 | 4 | return Status::OK(); |
339 | 5 | } |
340 | | |
341 | | } // namespace doris |