|
41 | 41 | #include "exec/sink/scale_writer_partitioning_exchanger.hpp" |
42 | 42 | #include "exec/sink/tablet_sink_hash_partitioner.h" |
43 | 43 | #include "exprs/vexpr.h" |
| 44 | +#include "format/transformer/merge_partitioner.h" |
44 | 45 | #include "runtime/runtime_profile.h" |
45 | 46 | #include "util/uid_util.h" |
46 | 47 |
|
@@ -175,6 +176,20 @@ Status ExchangeSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& inf |
175 | 176 | RETURN_IF_ERROR(_partitioner->prepare(state, p._row_desc)); |
176 | 177 | custom_profile()->add_info_string( |
177 | 178 | "Partitioner", fmt::format("ScaleWriterPartitioner({})", _partition_count)); |
| 179 | + } else if (_part_type == TPartitionType::MERGE_PARTITIONED) { |
| 180 | + if (!p._has_merge_partition_info) { |
| 181 | + return Status::InternalError("Merge partition info is missing"); |
| 182 | + } |
| 183 | + _partition_count = channels.size(); |
| 184 | + const bool use_new_shuffle_hash_method = |
| 185 | + _state->query_options().__isset.enable_new_shuffle_hash_method && |
| 186 | + _state->query_options().enable_new_shuffle_hash_method; |
| 187 | + _partitioner = std::make_unique<MergePartitioner>(_partition_count, p._merge_partition_info, |
| 188 | + use_new_shuffle_hash_method); |
| 189 | + RETURN_IF_ERROR(_partitioner->init({})); |
| 190 | + RETURN_IF_ERROR(_partitioner->prepare(state, p._row_desc)); |
| 191 | + custom_profile()->add_info_string("Partitioner", |
| 192 | + fmt::format("MergePartitioner({})", _partition_count)); |
178 | 193 | } |
179 | 194 |
|
180 | 195 | return Status::OK(); |
@@ -258,7 +273,8 @@ Status ExchangeSinkLocalState::open(RuntimeState* state) { |
258 | 273 | if (_part_type == TPartitionType::HASH_PARTITIONED || |
259 | 274 | _part_type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || |
260 | 275 | _part_type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || |
261 | | - _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED) { |
| 276 | + _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED || |
| 277 | + _part_type == TPartitionType::MERGE_PARTITIONED) { |
262 | 278 | RETURN_IF_ERROR(_partitioner->open(state)); |
263 | 279 | } |
264 | 280 | return Status::OK(); |
@@ -301,13 +317,18 @@ ExchangeSinkOperatorX::ExchangeSinkOperatorX( |
301 | 317 | sink.output_partition.type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED || |
302 | 318 | sink.output_partition.type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || |
303 | 319 | sink.output_partition.type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || |
304 | | - sink.output_partition.type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED); |
| 320 | + sink.output_partition.type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED || |
| 321 | + sink.output_partition.type == TPartitionType::MERGE_PARTITIONED); |
305 | 322 | #endif |
306 | 323 | _name = "ExchangeSinkOperatorX"; |
307 | 324 | _pool = std::make_shared<ObjectPool>(); |
308 | 325 | if (sink.__isset.output_tuple_id) { |
309 | 326 | _output_tuple_id = sink.output_tuple_id; |
310 | 327 | } |
| 328 | + if (sink.output_partition.__isset.merge_partition_info) { |
| 329 | + _merge_partition_info = sink.output_partition.merge_partition_info; |
| 330 | + _has_merge_partition_info = true; |
| 331 | + } |
311 | 332 |
|
312 | 333 | if (_part_type != TPartitionType::UNPARTITIONED) { |
313 | 334 | // if the destinations only one dest, we need to use broadcast |
@@ -507,8 +528,9 @@ Status ExchangeSinkOperatorX::sink(RuntimeState* state, Block* block, bool eos) |
507 | 528 | (local_state.current_channel_idx + 1) % local_state.channels.size(); |
508 | 529 | } else if (_part_type == TPartitionType::HASH_PARTITIONED || |
509 | 530 | _part_type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || |
| 531 | + _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED || |
510 | 532 | _part_type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || |
511 | | - _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED) { |
| 533 | + _part_type == TPartitionType::MERGE_PARTITIONED) { |
512 | 534 | RETURN_IF_ERROR(local_state._writer->write(state, block, eos)); |
513 | 535 | } else if (_part_type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED) { |
514 | 536 | // Control the number of channels according to the flow, thereby controlling the number of table sink writers. |
|
0 commit comments