| // Licensed to the Apache Software Foundation (ASF) under one |
| // or more contributor license agreements. See the NOTICE file |
| // distributed with this work for additional information |
| // regarding copyright ownership. The ASF licenses this file |
| // to you under the Apache License, Version 2.0 (the |
| // "License"); you may not use this file except in compliance |
| // with the License. You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, |
| // software distributed under the License is distributed on an |
| // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| // KIND, either express or implied. See the License for the |
| // specific language governing permissions and limitations |
| // under the License. |
| |
| #include "exec/sink/viceberg_merge_sink.h" |
| |
| #include <fmt/format.h> |
| |
| #include "agent/be_exec_version_manager.h" |
| #include "common/consts.h" |
| #include "common/exception.h" |
| #include "common/logging.h" |
| #include "core/block/block.h" |
| #include "core/column/column_nullable.h" |
| #include "core/column/column_string.h" |
| #include "core/column/column_struct.h" |
| #include "core/column/column_vector.h" |
| #include "core/data_type/data_type_nullable.h" |
| #include "core/data_type/data_type_struct.h" |
| #include "exec/sink/sink_common.h" |
| #include "exec/sink/viceberg_delete_sink.h" |
| #include "exec/sink/writer/iceberg/viceberg_table_writer.h" |
| #include "exprs/vexpr_context.h" |
| #include "format/table/iceberg/schema.h" |
| #include "format/table/iceberg/schema_parser.h" |
| #include "runtime/runtime_state.h" |
| #include "util/string_util.h" |
| |
| namespace doris { |
| |
| namespace {} // namespace |
| |
| VIcebergMergeSink::VIcebergMergeSink(const TDataSink& t_sink, const VExprContextSPtrs& output_exprs, |
| std::shared_ptr<Dependency> dep, |
| std::shared_ptr<Dependency> fin_dep) |
| : AsyncResultWriter(output_exprs, dep, fin_dep), _t_sink(t_sink) { |
| DCHECK(_t_sink.__isset.iceberg_merge_sink); |
| } |
| |
| VIcebergMergeSink::~VIcebergMergeSink() = default; |
| |
| Status VIcebergMergeSink::init_properties(ObjectPool* pool, const RowDescriptor& row_desc) { |
| RETURN_IF_ERROR(_build_inner_sinks()); |
| |
| _table_writer = std::make_unique<VIcebergTableWriter>(_table_sink, _table_output_expr_ctxs, |
| nullptr, nullptr); |
| _table_writer->defer_file_cleanup_until_outer_close(); |
| _delete_writer = std::make_unique<VIcebergDeleteSink>(_delete_sink, _delete_output_expr_ctxs, |
| nullptr, nullptr); |
| _delete_writer->defer_file_cleanup_until_outer_close(); |
| RETURN_IF_ERROR(_table_writer->init_properties(pool, row_desc)); |
| RETURN_IF_ERROR(_delete_writer->init_properties(pool)); |
| return Status::OK(); |
| } |
| |
| Status VIcebergMergeSink::open(RuntimeState* state, RuntimeProfile* profile) { |
| _state = state; |
| |
| _written_rows_counter = ADD_COUNTER(profile, "RowsWritten", TUnit::UNIT); |
| _insert_rows_counter = ADD_COUNTER(profile, "InsertRows", TUnit::UNIT); |
| _delete_rows_counter = ADD_COUNTER(profile, "DeleteRows", TUnit::UNIT); |
| // The query-wide version keeps validation all-or-nothing during a rolling BE upgrade. |
| _require_merge_cardinality_check = |
| _require_merge_cardinality_check && |
| state->be_exec_version() >= SUPPORT_ICEBERG_MERGE_CARDINALITY_VERSION; |
| if (_require_merge_cardinality_check) { |
| _matched_row_id_state_bytes_counter = |
| ADD_COUNTER(profile, "MatchedRowIdStateBytes", TUnit::BYTES); |
| } |
| _send_data_timer = ADD_TIMER(profile, "SendDataTime"); |
| _open_timer = ADD_TIMER(profile, "OpenTime"); |
| _close_timer = ADD_TIMER(profile, "CloseTime"); |
| |
| SCOPED_TIMER(_open_timer); |
| |
| RETURN_IF_ERROR(_prepare_output_layout()); |
| |
| RuntimeProfile* table_profile = profile->create_child("IcebergMergeTableWriter", true, true); |
| RuntimeProfile* delete_profile = profile->create_child("IcebergMergeDeleteWriter", true, true); |
| |
| RETURN_IF_ERROR(_table_writer->open(state, table_profile)); |
| RETURN_IF_ERROR(_delete_writer->open(state, delete_profile)); |
| |
| return Status::OK(); |
| } |
| |
| Status VIcebergMergeSink::write(RuntimeState* state, Block& block) { |
| SCOPED_TIMER(_send_data_timer); |
| if (block.rows() == 0) { |
| return Status::OK(); |
| } |
| |
| Block output_block; |
| RETURN_IF_ERROR(_projection_block(block, &output_block)); |
| if (output_block.rows() == 0) { |
| return Status::OK(); |
| } |
| |
| if (_operation_idx < 0 || _row_id_idx < 0) { |
| return Status::InternalError("Iceberg merge sink missing operation/row_id columns"); |
| } |
| |
| const auto& op_column = output_block.get_by_position(_operation_idx).column; |
| const auto* op_data = remove_nullable(op_column).get(); |
| |
| IColumn::Filter delete_filter(output_block.rows(), 0); |
| IColumn::Filter insert_filter(output_block.rows(), 0); |
| bool has_delete = false; |
| bool has_insert = false; |
| size_t delete_rows = 0; |
| size_t insert_rows = 0; |
| |
| for (size_t i = 0; i < output_block.rows(); ++i) { |
| int8_t op = static_cast<int8_t>(op_data->get_int(i)); |
| bool delete_op = is_delete_op(op); |
| bool insert_op = is_insert_op(op); |
| if (!delete_op && !insert_op) { |
| return Status::InternalError("Unknown Iceberg merge operation {}", op); |
| } |
| if (delete_op) { |
| delete_filter[i] = 1; |
| has_delete = true; |
| ++delete_rows; |
| } |
| if (insert_op) { |
| insert_filter[i] = 1; |
| has_insert = true; |
| ++insert_rows; |
| } |
| } |
| |
| if (_require_merge_cardinality_check) { |
| // The physical sink hashes matched rows by row_id, so exact state retained across blocks |
| // enforces SQL MERGE cardinality for the whole query without changing UPDATE semantics. |
| RETURN_IF_ERROR(_validate_matched_row_ids(output_block, delete_filter.data())); |
| COUNTER_SET(_matched_row_id_state_bytes_counter, |
| static_cast<int64_t>(_matched_row_id_state_size)); |
| } |
| _row_count += output_block.rows(); |
| _delete_row_count += delete_rows; |
| _insert_row_count += insert_rows; |
| |
| bool skip_io = false; |
| #ifdef BE_TEST |
| skip_io = _skip_io; |
| #endif |
| |
| if (has_delete && !skip_io) { |
| Block delete_block = output_block; |
| std::vector<int> delete_indices {_row_id_idx}; |
| delete_block.erase_not_in(delete_indices); |
| Block::filter_block_internal(&delete_block, delete_filter); |
| RETURN_IF_ERROR(_delete_writer->write(state, delete_block)); |
| } |
| |
| if (has_insert && !skip_io) { |
| if (_data_column_indices.empty()) { |
| return Status::InternalError("Iceberg merge sink has no data columns for insert"); |
| } |
| Block insert_block = output_block; |
| insert_block.erase_not_in(_data_column_indices); |
| Block::filter_block_internal(&insert_block, insert_filter); |
| RETURN_IF_ERROR(_table_writer->write_prepared_block(insert_block)); |
| } |
| |
| if (_written_rows_counter != nullptr) { |
| COUNTER_UPDATE(_written_rows_counter, output_block.rows()); |
| } |
| if (_insert_rows_counter != nullptr) { |
| COUNTER_UPDATE(_insert_rows_counter, insert_rows); |
| } |
| if (_delete_rows_counter != nullptr) { |
| COUNTER_UPDATE(_delete_rows_counter, delete_rows); |
| } |
| |
| return Status::OK(); |
| } |
| |
| Status VIcebergMergeSink::_validate_matched_row_ids(const Block& block, |
| const uint8_t* delete_filter) { |
| const auto& row_id = block.get_by_position(_row_id_idx); |
| const IColumn* row_id_data = row_id.column.get(); |
| const IDataType* row_id_type = row_id.type.get(); |
| const auto* nullable_row_id = check_and_get_column<ColumnNullable>(row_id_data); |
| if (nullable_row_id != nullptr) { |
| row_id_data = nullable_row_id->get_nested_column_ptr().get(); |
| } |
| if (const auto* nullable_type = check_and_get_data_type<DataTypeNullable>(row_id_type)) { |
| row_id_type = nullable_type->get_nested_type().get(); |
| } |
| |
| const auto* struct_column = check_and_get_column<ColumnStruct>(row_id_data); |
| const auto* struct_type = check_and_get_data_type<DataTypeStruct>(row_id_type); |
| if (struct_column == nullptr || struct_type == nullptr) { |
| return Status::InternalError("Iceberg merge row_id column is not a struct"); |
| } |
| |
| int file_path_idx = -1; |
| int row_position_idx = -1; |
| const auto& field_names = struct_type->get_element_names(); |
| for (size_t i = 0; i < field_names.size(); ++i) { |
| std::string field_name = doris::to_lower(field_names[i]); |
| if (field_name == "file_path") { |
| file_path_idx = static_cast<int>(i); |
| } else if (field_name == "row_position") { |
| row_position_idx = static_cast<int>(i); |
| } |
| } |
| if (file_path_idx < 0 || row_position_idx < 0) { |
| return Status::InternalError( |
| "Iceberg merge row_id must contain file_path and row_position fields"); |
| } |
| |
| const auto& file_path_column = struct_column->get_column_ptr(file_path_idx); |
| const auto& row_position_column = struct_column->get_column_ptr(row_position_idx); |
| const auto* nullable_file_path = check_and_get_column<ColumnNullable>(file_path_column.get()); |
| const auto* nullable_row_position = |
| check_and_get_column<ColumnNullable>(row_position_column.get()); |
| const auto* file_paths = |
| check_and_get_column<ColumnString>(remove_nullable(file_path_column).get()); |
| const auto* row_positions = check_and_get_column<ColumnVector<TYPE_BIGINT>>( |
| remove_nullable(row_position_column).get()); |
| if (file_paths == nullptr || row_positions == nullptr) { |
| return Status::InternalError("Iceberg merge row_id fields have incorrect types"); |
| } |
| |
| std::map<roaring::Roaring64Map*, size_t> touched_bitmap_sizes; |
| for (size_t i = 0; i < block.rows(); ++i) { |
| if (delete_filter[i] == 0) { |
| continue; |
| } |
| if ((nullable_row_id != nullptr && nullable_row_id->is_null_at(i)) || |
| (nullable_file_path != nullptr && nullable_file_path->is_null_at(i)) || |
| (nullable_row_position != nullptr && nullable_row_position->is_null_at(i))) { |
| return Status::InternalError("Iceberg merge matched row_id cannot be null"); |
| } |
| |
| int64_t row_position = row_positions->get_element(i); |
| if (row_position < 0) { |
| return Status::InternalError("Invalid row_position {} in Iceberg merge row_id", |
| row_position); |
| } |
| // Intern each file path once and keep exact positions in a compressed bitmap; retaining a |
| // full path string per matched row makes MERGE memory grow with path_length * row_count. |
| auto [file_it, inserted] = |
| _matched_row_positions.try_emplace(file_paths->get_data_at(i).to_string()); |
| auto* positions = &file_it->second; |
| auto touched_it = touched_bitmap_sizes.find(positions); |
| if (touched_it == touched_bitmap_sizes.end()) { |
| touched_it = touched_bitmap_sizes.emplace(positions, positions->getSizeInBytes()).first; |
| } |
| if (inserted) { |
| _matched_row_id_state_size += |
| sizeof(std::pair<const std::string, roaring::Roaring64Map>); |
| _matched_row_id_state_size += file_it->first.capacity(); |
| _matched_row_id_state_size += touched_it->second; |
| } |
| if (!positions->addChecked(static_cast<uint64_t>(row_position))) { |
| return Status::InvalidArgument( |
| "Iceberg MERGE failed because multiple source rows matched the same target " |
| "row"); |
| } |
| } |
| |
| // Measure only bitmaps touched by this block; rescanning all retained files on every write |
| // makes a many-file MERGE quadratic in the number of input blocks. |
| for (const auto& [positions, previous_size] : touched_bitmap_sizes) { |
| size_t current_size = positions->getSizeInBytes(); |
| if (current_size >= previous_size) { |
| _matched_row_id_state_size += current_size - previous_size; |
| } else { |
| _matched_row_id_state_size -= previous_size - current_size; |
| } |
| } |
| return Status::OK(); |
| } |
| |
| Status VIcebergMergeSink::close(Status close_status) { |
| SCOPED_TIMER(_close_timer); |
| |
| if (!close_status.ok()) { |
| LOG(WARNING) << fmt::format("VIcebergMergeSink close with error: {}", |
| close_status.to_string()); |
| if (_table_writer) { |
| static_cast<void>(_table_writer->close(close_status)); |
| } |
| if (_delete_writer) { |
| static_cast<void>(_delete_writer->close(close_status)); |
| } |
| return close_status; |
| } |
| |
| Status table_status = Status::OK(); |
| Status delete_status = Status::OK(); |
| if (_table_writer) { |
| table_status = _table_writer->close(close_status); |
| } |
| if (_delete_writer) { |
| delete_status = _delete_writer->close(close_status); |
| } |
| |
| if (_written_rows_counter != nullptr) { |
| COUNTER_SET(_written_rows_counter, static_cast<int64_t>(_row_count)); |
| } |
| if (_insert_rows_counter != nullptr) { |
| COUNTER_SET(_insert_rows_counter, static_cast<int64_t>(_insert_row_count)); |
| } |
| if (_delete_rows_counter != nullptr) { |
| COUNTER_SET(_delete_rows_counter, static_cast<int64_t>(_delete_row_count)); |
| } |
| |
| Status result_status = table_status.ok() ? delete_status : table_status; |
| if (_table_writer) { |
| _table_writer->finish_deferred_file_cleanup(result_status); |
| } |
| if (_delete_writer) { |
| _delete_writer->finish_deferred_file_cleanup(result_status); |
| } |
| return result_status; |
| } |
| |
| Status VIcebergMergeSink::_build_inner_sinks() { |
| if (!_t_sink.__isset.iceberg_merge_sink) { |
| return Status::InternalError("Missing iceberg merge sink config"); |
| } |
| |
| const auto& merge_sink = _t_sink.iceberg_merge_sink; |
| // Missing means an old FE plan, which predates SQL MERGE cardinality validation. |
| _require_merge_cardinality_check = merge_sink.__isset.require_merge_cardinality_check && |
| merge_sink.require_merge_cardinality_check; |
| |
| TIcebergTableSink table_sink; |
| if (merge_sink.__isset.db_name) { |
| table_sink.__set_db_name(merge_sink.db_name); |
| } |
| if (merge_sink.__isset.tb_name) { |
| table_sink.__set_tb_name(merge_sink.tb_name); |
| } |
| if (merge_sink.__isset.schema_json) { |
| table_sink.__set_schema_json(merge_sink.schema_json); |
| } |
| if (merge_sink.__isset.partition_specs_json) { |
| table_sink.__set_partition_specs_json(merge_sink.partition_specs_json); |
| } |
| if (merge_sink.__isset.partition_spec_id) { |
| table_sink.__set_partition_spec_id(merge_sink.partition_spec_id); |
| } |
| if (merge_sink.__isset.sort_fields) { |
| table_sink.__set_sort_fields(merge_sink.sort_fields); |
| } |
| if (merge_sink.__isset.file_format) { |
| table_sink.__set_file_format(merge_sink.file_format); |
| } |
| if (merge_sink.__isset.compression_type) { |
| table_sink.__set_compression_type(merge_sink.compression_type); |
| } |
| if (merge_sink.__isset.output_path) { |
| table_sink.__set_output_path(merge_sink.output_path); |
| } |
| if (merge_sink.__isset.original_output_path) { |
| table_sink.__set_original_output_path(merge_sink.original_output_path); |
| } |
| if (merge_sink.__isset.hadoop_config) { |
| table_sink.__set_hadoop_config(merge_sink.hadoop_config); |
| } |
| if (merge_sink.__isset.file_type) { |
| table_sink.__set_file_type(merge_sink.file_type); |
| } |
| if (merge_sink.__isset.broker_addresses) { |
| table_sink.__set_broker_addresses(merge_sink.broker_addresses); |
| } |
| if (merge_sink.__isset.collect_column_stats) { |
| table_sink.__set_collect_column_stats(merge_sink.collect_column_stats); |
| } |
| _table_sink.__set_type(TDataSinkType::ICEBERG_TABLE_SINK); |
| _table_sink.__set_iceberg_table_sink(table_sink); |
| |
| TIcebergDeleteSink delete_sink; |
| if (merge_sink.__isset.db_name) { |
| delete_sink.__set_db_name(merge_sink.db_name); |
| } |
| if (merge_sink.__isset.tb_name) { |
| delete_sink.__set_tb_name(merge_sink.tb_name); |
| } |
| if (merge_sink.__isset.delete_type) { |
| delete_sink.__set_delete_type(merge_sink.delete_type); |
| } |
| if (merge_sink.__isset.file_format) { |
| delete_sink.__set_file_format(merge_sink.file_format); |
| } |
| if (merge_sink.__isset.compression_type) { |
| delete_sink.__set_compress_type(merge_sink.compression_type); |
| } |
| if (merge_sink.__isset.output_path) { |
| delete_sink.__set_output_path(merge_sink.output_path); |
| } |
| if (merge_sink.__isset.table_location) { |
| delete_sink.__set_table_location(merge_sink.table_location); |
| } |
| if (merge_sink.__isset.hadoop_config) { |
| delete_sink.__set_hadoop_config(merge_sink.hadoop_config); |
| } |
| if (merge_sink.__isset.file_type) { |
| delete_sink.__set_file_type(merge_sink.file_type); |
| } |
| if (merge_sink.__isset.partition_spec_id_for_delete) { |
| delete_sink.__set_partition_spec_id(merge_sink.partition_spec_id_for_delete); |
| } |
| if (merge_sink.__isset.partition_data_json_for_delete) { |
| delete_sink.__set_partition_data_json(merge_sink.partition_data_json_for_delete); |
| } |
| if (merge_sink.__isset.broker_addresses) { |
| delete_sink.__set_broker_addresses(merge_sink.broker_addresses); |
| } |
| if (merge_sink.__isset.format_version) { |
| delete_sink.__set_format_version(merge_sink.format_version); |
| } |
| if (merge_sink.__isset.rewritable_delete_file_sets) { |
| delete_sink.__set_rewritable_delete_file_sets(merge_sink.rewritable_delete_file_sets); |
| } |
| _delete_sink.__set_type(TDataSinkType::ICEBERG_DELETE_SINK); |
| _delete_sink.__set_iceberg_delete_sink(delete_sink); |
| |
| return Status::OK(); |
| } |
| |
| Status VIcebergMergeSink::_prepare_output_layout() { |
| if (_vec_output_expr_ctxs.empty()) { |
| return Status::InternalError("Iceberg merge sink has empty output expressions"); |
| } |
| |
| std::string row_id_name = doris::to_lower(BeConsts::ICEBERG_ROWID_COL); |
| std::string op_name = doris::to_lower(kOperationColumnName); |
| |
| _operation_idx = -1; |
| _row_id_idx = -1; |
| for (size_t i = 0; i < _vec_output_expr_ctxs.size(); ++i) { |
| std::string expr_name = doris::to_lower(_vec_output_expr_ctxs[i]->expr_name()); |
| if (_operation_idx < 0 && expr_name == op_name) { |
| _operation_idx = static_cast<int>(i); |
| } else if (_row_id_idx < 0 && expr_name == row_id_name) { |
| _row_id_idx = static_cast<int>(i); |
| } |
| } |
| |
| if (_operation_idx < 0) { |
| return Status::InternalError("Iceberg merge sink missing operation column"); |
| } |
| if (_row_id_idx < 0) { |
| return Status::InternalError("Iceberg merge sink missing row_id column"); |
| } |
| |
| _data_column_indices.clear(); |
| _table_output_expr_ctxs.clear(); |
| for (size_t i = 0; i < _vec_output_expr_ctxs.size(); ++i) { |
| if (static_cast<int>(i) == _operation_idx || static_cast<int>(i) == _row_id_idx) { |
| continue; |
| } |
| _data_column_indices.push_back(static_cast<int>(i)); |
| _table_output_expr_ctxs.emplace_back(_vec_output_expr_ctxs[i]); |
| } |
| |
| return Status::OK(); |
| } |
| |
| } // namespace doris |
