diff --git a/be/src/exec/operator/iceberg_sorter_reserve_memory.h b/be/src/exec/operator/iceberg_sorter_reserve_memory.h new file mode 100644 index 00000000000000..1ad1454b4af5f3 --- /dev/null +++ b/be/src/exec/operator/iceberg_sorter_reserve_memory.h @@ -0,0 +1,45 @@ +// 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. + +#pragma once + +#include +#include +#include + +namespace doris { + +struct IcebergSorterReserveMemory { + size_t retained_growth = 0; + size_t transient_workspace = 0; +}; + +inline size_t bounded_iceberg_reserve_size( + const std::vector& per_partition_reservations) { + size_t retained_growth = 0; + size_t transient_workspace = 0; + for (const auto& reservation : per_partition_reservations) { + retained_growth = std::min(std::numeric_limits::max() - retained_growth, + reservation.retained_growth) + + retained_growth; + transient_workspace = std::max(transient_workspace, reservation.transient_workspace); + } + return std::min(std::numeric_limits::max() - retained_growth, transient_workspace) + + retained_growth; +} + +} // namespace doris diff --git a/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp b/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp index cf7c8e6e1a1538..5d0d1cb00a7916 100644 --- a/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp +++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp @@ -55,26 +55,34 @@ size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state if (!_writer) { return 0; } - auto current_writer = _writer->current_writer(); - auto* sort_writer = dynamic_cast(current_writer.get()); - if (!sort_writer) { - return 0; + std::vector per_partition_reservations; + auto active_writers = _writer->active_writers(); + per_partition_reservations.reserve(active_writers->size()); + for (const auto& writer : *active_writers) { + if (auto* sort_writer = dynamic_cast(writer.get())) { + auto reservation = sort_writer->get_reserve_mem_size_components(state, eos); + per_partition_reservations.push_back( + {.retained_growth = reservation.retained_growth, + .transient_workspace = reservation.transient_workspace}); + } } - - return sort_writer->get_reserve_mem_size(state, eos); + // Column growth remains in every touched sorter, while sorting workspace is reused by serial dispatch. + return bounded_iceberg_reserve_size(per_partition_reservations); } size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState* state) const { if (!_writer) { return 0; } - auto current_writer = _writer->current_writer(); - auto* sort_writer = dynamic_cast(current_writer.get()); - if (!sort_writer) { - return 0; + size_t revocable_size = 0; + // Retain the published container while the async writer may replace the current snapshot. + auto active_writers = _writer->active_writers(); + for (const auto& writer : *active_writers) { + if (auto* sort_writer = dynamic_cast(writer.get())) { + revocable_size += sort_writer->data_size(); + } } - - return sort_writer->data_size(); + return revocable_size; } Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) { @@ -82,20 +90,25 @@ Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) { if (!_writer) { return Status::OK(); } - auto current_writer = _writer->current_writer(); - auto* sort_writer = dynamic_cast(current_writer.get()); - if (!sort_writer) { - return Status::OK(); + std::shared_ptr largest_writer; + size_t largest_size = 0; + // Retain the published container while the async writer may replace the current snapshot. + auto active_writers = _writer->active_writers(); + for (const auto& writer : *active_writers) { + if (auto* sort_writer = dynamic_cast(writer.get())) { + size_t size = sort_writer->data_size(); + if (size > largest_size) { + largest_size = size; + largest_writer = writer; + } + } } - - auto exception_catch_func = [current_writer, sort_writer]() { - auto status = [&]() { - RETURN_IF_CATCH_EXCEPTION({ return sort_writer->trigger_spill(); }); - }(); - return status; - }; - - return exception_catch_func(); + if (largest_writer != nullptr) { + // Repeated revocation drains the largest partition first without launching O(P) spill jobs at once. + auto* sort_writer = dynamic_cast(largest_writer.get()); + RETURN_IF_CATCH_EXCEPTION({ RETURN_IF_ERROR(sort_writer->trigger_spill()); }); + } + return Status::OK(); } SpillIcebergTableSinkOperatorX::SpillIcebergTableSinkOperatorX( diff --git a/be/src/exec/operator/spill_iceberg_table_sink_operator.h b/be/src/exec/operator/spill_iceberg_table_sink_operator.h index 6da926ae20fb91..5ffdd7505599ea 100644 --- a/be/src/exec/operator/spill_iceberg_table_sink_operator.h +++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.h @@ -18,7 +18,9 @@ #pragma once #include +#include +#include "exec/operator/iceberg_sorter_reserve_memory.h" #include "exec/operator/operator.h" #include "exec/sink/writer/iceberg/viceberg_table_writer.h" @@ -87,4 +89,4 @@ class SpillIcebergTableSinkOperatorX final ObjectPool* _pool = nullptr; }; -} // namespace doris \ No newline at end of file +} // namespace doris diff --git a/be/src/exec/pipeline/pipeline_fragment_context.cpp b/be/src/exec/pipeline/pipeline_fragment_context.cpp index def2cf473aea77..5f3bb8d1070158 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.cpp +++ b/be/src/exec/pipeline/pipeline_fragment_context.cpp @@ -467,6 +467,8 @@ Status PipelineFragmentContext::_build_pipeline_tasks_for_instance( _params.query_options, _query_ctx->query_globals, _exec_env, _query_ctx.get()); { // Initialize runtime state for this task + task_runtime_state->set_iceberg_commit_data_budget( + _runtime_state->iceberg_commit_data_budget()); task_runtime_state->set_query_mem_tracker(_query_ctx->query_mem_tracker()); task_runtime_state->set_task_execution_context(shared_from_this()); @@ -2514,16 +2516,14 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r } } } - if (auto icd = req.runtime_state->iceberg_commit_datas(); !icd.empty()) { + req.runtime_state->append_iceberg_commit_datas(¶ms.iceberg_commit_datas); + if (!params.iceberg_commit_datas.empty()) { params.__isset.iceberg_commit_datas = true; - params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(), icd.begin(), - icd.end()); } else if (!req.runtime_states.empty()) { for (auto* rs : req.runtime_states) { - if (auto rs_icd = rs->iceberg_commit_datas(); !rs_icd.empty()) { + rs->append_iceberg_commit_datas(¶ms.iceberg_commit_datas); + if (!params.iceberg_commit_datas.empty()) { params.__isset.iceberg_commit_datas = true; - params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(), - rs_icd.begin(), rs_icd.end()); } } } diff --git a/be/src/exec/sink/viceberg_delete_sink.cpp b/be/src/exec/sink/viceberg_delete_sink.cpp index 172fdd28177c62..a3dcb1df22f657 100644 --- a/be/src/exec/sink/viceberg_delete_sink.cpp +++ b/be/src/exec/sink/viceberg_delete_sink.cpp @@ -283,8 +283,12 @@ Status VIcebergDeleteSink::close(Status close_status) { _delete_file_count); if (_state != nullptr) { - for (const auto& commit_data : _commit_data_list) { - _state->add_iceberg_commit_datas(commit_data); + for (auto& commit_data : _commit_data_list) { + Status report_status = _state->add_iceberg_commit_datas(std::move(commit_data)); + if (!report_status.ok()) { + _cleanup_created_files(); + return report_status; + } } } diff --git a/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp b/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp index ba7644daec751f..7faaebe9525e88 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp @@ -69,6 +69,7 @@ Status VIcebergPartitionWriter::open(RuntimeState* state, RuntimeProfile* profil io::FileWriterOptions file_writer_options = {.used_by_s3_committer = false}; RETURN_IF_ERROR(_fs->create_file(file_description.path, &_file_writer, &file_writer_options)); + Status open_status; switch (_file_format_type) { case TFileFormatType::FORMAT_PARQUET: { TParquetCompressionType::type parquet_compression_type; @@ -92,9 +93,13 @@ Status VIcebergPartitionWriter::open(RuntimeState* state, RuntimeProfile* profil break; } default: { - return Status::InternalError("Unsupported compress type {} with parquet", - to_string(_compress_type)); + open_status = Status::InternalError("Unsupported compress type {} with parquet", + to_string(_compress_type)); + break; + } } + if (!open_status.ok()) { + break; } ParquetFileOptions parquet_options = {.compression_type = parquet_compression_type, .parquet_version = TParquetVersion::PARQUET_1_0, @@ -103,19 +108,28 @@ Status VIcebergPartitionWriter::open(RuntimeState* state, RuntimeProfile* profil _file_format_transformer = std::make_unique( state, _file_writer.get(), _write_output_expr_ctxs, _write_column_names, false, parquet_options, _iceberg_schema_json, &_schema); - return _file_format_transformer->open(); + open_status = _file_format_transformer->open(); + break; } case TFileFormatType::FORMAT_ORC: { _file_format_transformer = std::make_unique( state, _file_writer.get(), _write_output_expr_ctxs, "", _write_column_names, false, _compress_type, &_schema, _fs); - return _file_format_transformer->open(); + open_status = _file_format_transformer->open(); + break; } default: { - return Status::InternalError("Unsupported file format type {}", - to_string(_file_format_type)); + open_status = Status::InternalError("Unsupported file format type {}", + to_string(_file_format_type)); + break; } } + if (!open_status.ok()) { + // A transformer failure happens after object creation, so abort multipart state before deleting the path. + WARN_IF_ERROR(_file_writer->abort(), "failed to abort Iceberg file after open error"); + WARN_IF_ERROR(_fs->delete_file(_path), "failed to delete Iceberg file after open error"); + } + return open_status; } Status VIcebergPartitionWriter::close(const Status& status) { @@ -147,7 +161,12 @@ Status VIcebergPartitionWriter::close(const Status& status) { } return commit_status; } - _state->add_iceberg_commit_datas(commit_data); + Status report_status = _state->add_iceberg_commit_datas(std::move(commit_data)); + if (!report_status.ok()) { + // A closed object that cannot be reported can never be committed, so remove it immediately. + WARN_IF_ERROR(_fs->delete_file(_path), "failed to delete unreportable Iceberg file"); + return report_status; + } if (_closed_file_callback) { _closed_file_callback(_fs, _path); } diff --git a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp index 6081166777fc28..444aec8933ae6a 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp @@ -84,6 +84,13 @@ size_t VIcebergSortWriter::get_reserve_mem_size(RuntimeState* state, bool eos) c return _sorter == nullptr ? 0 : _sorter->get_reserve_mem_size(state, eos); } +SorterReserveMemory VIcebergSortWriter::get_reserve_mem_size_components(RuntimeState* state, + bool eos) const { + std::lock_guard lock(_sorter_mutex); + return _sorter == nullptr ? SorterReserveMemory {} + : _sorter->get_reserve_mem_size_components(state, eos); +} + Status VIcebergSortWriter::trigger_spill() { std::lock_guard lock(_sorter_mutex); if (_closed || _sorter == nullptr) { @@ -102,80 +109,35 @@ Status VIcebergSortWriter::close(const Status& status) { } Status VIcebergSortWriter::_close_locked(const Status& status) { - // Track the actual internal status of operations performed during close. - // This is important because if intermediate operations (like do_sort()) fail, - // we need to propagate the actual error status to the underlying partition writer's - // close() call, rather than the original status parameter which could be OK. - Status internal_status = Status::OK(); - // Track the close status of the underlying partition writer. - // If _iceberg_partition_writer->close() fails (e.g., Parquet file flush error), - // we must propagate this error to the caller to avoid silent data loss. - Status close_status = Status::OK(); - - // Defer ensures the underlying partition writer is always closed and - // spill streams are cleaned up, regardless of whether intermediate operations succeed. - // Uses internal_status to propagate any errors that occurred during close operations. - Defer defer {[&]() { - // If any intermediate operation failed, pass that error to the partition writer; - // otherwise, pass the original status from the caller. - close_status = - _iceberg_partition_writer->close(internal_status.ok() ? status : internal_status); - if (!close_status.ok()) { - LOG(WARNING) << fmt::format("_iceberg_partition_writer close failed, reason: {}", - close_status.to_string()); - } - _cleanup_spill_streams(); - }}; - - // If the original status is already an error or the query is cancelled, - // skip all close operations and propagate the original error - if (!status.ok() || _runtime_state->is_cancelled()) { - return status; - } - - // If sorter was never initialized (e.g., no data was written), nothing to do - if (_sorter == nullptr) { - return Status::OK(); - } - - // Check if there is any remaining data in the sorter (either unsorted or already sorted blocks) - if (!_sorter->merge_sort_state()->unsorted_block()->empty() || - !_sorter->merge_sort_state()->get_sorted_block().empty()) { - if (_sorted_spill_files.empty()) { - // No spill has occurred, all data is in memory. - // Sort the remaining data, prepare for reading, and write to file. - internal_status = _sorter->do_sort(); - if (!internal_status.ok()) { - return internal_status; - } - internal_status = _sorter->prepare_for_read(false); - if (!internal_status.ok()) { - return internal_status; + Status internal_status = status; + if (status.ok() && !_runtime_state->is_cancelled()) { + internal_status = Status::OK(); + if (_sorter != nullptr && (!_sorter->merge_sort_state()->unsorted_block()->empty() || + !_sorter->merge_sort_state()->get_sorted_block().empty())) { + if (_sorted_spill_files.empty()) { + internal_status = _sorter->do_sort(); + if (internal_status.ok()) { + internal_status = _sorter->prepare_for_read(false); + } + if (internal_status.ok()) { + internal_status = _write_sorted_data(); + } + } else { + internal_status = _do_spill(); } - internal_status = _write_sorted_data(); - return internal_status; } - - // Some data has already been spilled to disk. - // Spill the remaining in-memory data to a new spill stream. - internal_status = _do_spill(); - if (!internal_status.ok()) { - return internal_status; + if (internal_status.ok() && !_sorted_spill_files.empty()) { + internal_status = _combine_files_output(); } } - // Merge all spilled streams using multi-way merge sort and output final sorted data to files - if (!_sorted_spill_files.empty()) { - internal_status = _combine_files_output(); - if (!internal_status.ok()) { - return internal_status; - } + // Form the return value only after the underlying close runs; a deferred assignment is too late. + Status close_status = + _iceberg_partition_writer->close(internal_status.ok() ? status : internal_status); + _cleanup_spill_streams(); + if (!internal_status.ok()) { + return internal_status; } - - // Return close_status if internal operations succeeded but the underlying - // partition writer's close() failed (e.g., file flush error). - // This prevents silent data loss where the caller thinks the write succeeded - // but the file was not properly closed. return close_status; } diff --git a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h index e1e512f0a0cf79..37659eeca4bc89 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h +++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h @@ -105,6 +105,8 @@ class VIcebergSortWriter : public IPartitionWriterBase { size_t get_reserve_mem_size(RuntimeState* state, bool eos) const; + SorterReserveMemory get_reserve_mem_size_components(RuntimeState* state, bool eos) const; + // Called by the memory management system to trigger spilling data to disk Status trigger_spill(); diff --git a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp index 80db0876a9ba40..24c0269cf18901 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp @@ -44,6 +44,7 @@ VIcebergTableWriter::VIcebergTableWriter(const TDataSink& t_sink, std::shared_ptr fin_dep) : AsyncResultWriter(output_expr_ctxs, dep, fin_dep), _t_sink(t_sink) { DCHECK(_t_sink.__isset.iceberg_table_sink); + _active_writers.store(std::make_shared()); } Status VIcebergTableWriter::open(RuntimeState* state, RuntimeProfile* profile) { @@ -242,7 +243,8 @@ Status VIcebergTableWriter::_process_row_lineage_columns(Block& block) { Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { RETURN_IF_ERROR(_process_row_lineage_columns(output_block)); - std::unordered_map, IColumn::Filter> writer_positions; + std::unordered_map, IColumn::Permutation> + writer_positions; _row_count += output_block.rows(); // Case 1: Full static partition - all data goes to a single partition @@ -259,6 +261,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } _partitions_to_writers.insert({_static_partition_path, writer}); RETURN_IF_ERROR(writer->open(_state, _operator_profile, _row_desc)); + _publish_active_writers(); } else { if (writer_iter->second->written_len() > _target_file_size_bytes) { std::string file_name(writer_iter->second->file_name()); @@ -268,6 +271,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { RETURN_IF_ERROR(writer_iter->second->close(Status::OK())); } _partitions_to_writers.erase(writer_iter); + _publish_active_writers(); try { writer = _create_partition_writer(nullptr, -1, &file_name, file_name_index + 1); @@ -276,6 +280,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } _partitions_to_writers.insert({_static_partition_path, writer}); RETURN_IF_ERROR(writer->open(_state, _operator_profile, _row_desc)); + _publish_active_writers(); } else { writer = writer_iter->second; } @@ -284,7 +289,6 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { SCOPED_RAW_TIMER(&_partition_writers_write_ns); output_block.erase(_non_write_columns_indices); RETURN_IF_ERROR(writer->write(output_block)); - _current_writer.store(writer); return Status::OK(); } @@ -302,6 +306,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } _partitions_to_writers.insert({"", writer}); RETURN_IF_ERROR(writer->open(_state, _operator_profile, _row_desc)); + _publish_active_writers(); } else { if (writer_iter->second->written_len() > _target_file_size_bytes) { std::string file_name(writer_iter->second->file_name()); @@ -311,6 +316,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { RETURN_IF_ERROR(writer_iter->second->close(Status::OK())); } _partitions_to_writers.erase(writer_iter); + _publish_active_writers(); try { writer = _create_partition_writer(nullptr, -1, &file_name, file_name_index + 1); @@ -319,6 +325,7 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } _partitions_to_writers.insert({"", writer}); RETURN_IF_ERROR(writer->open(_state, _operator_profile, _row_desc)); + _publish_active_writers(); } else { writer = writer_iter->second; } @@ -327,7 +334,6 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { SCOPED_RAW_TIMER(&_partition_writers_write_ns); output_block.erase(_non_write_columns_indices); RETURN_IF_ERROR(writer->write(output_block)); - _current_writer.store(writer); return Status::OK(); } @@ -376,10 +382,8 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { auto writer = _create_partition_writer(&transformed_block, position, file_name, file_name_index); RETURN_IF_ERROR(writer->open(_state, _operator_profile, _row_desc)); - IColumn::Filter filter(output_block.rows(), 0); - filter[position] = 1; - writer_positions.insert({writer, std::move(filter)}); _partitions_to_writers.insert({partition_name, writer}); + _publish_active_writers(); writer_ptr = writer; } catch (doris::Exception& e) { return e.to_status(); @@ -388,8 +392,8 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { }; auto writer_iter = _partitions_to_writers.find(partition_name); + std::shared_ptr writer; if (writer_iter == _partitions_to_writers.end()) { - std::shared_ptr writer; if (_partitions_to_writers.size() + 1 > config::table_sink_partition_write_max_partition_nums_per_writer) { return Status::InternalError( @@ -398,7 +402,6 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } RETURN_IF_ERROR(create_and_open_writer(partition_name, i, nullptr, 0, writer)); } else { - std::shared_ptr writer; if (writer_iter->second->written_len() > _target_file_size_bytes) { std::string file_name(writer_iter->second->file_name()); int file_name_index = writer_iter->second->file_name_index(); @@ -408,53 +411,53 @@ Status VIcebergTableWriter::_write_prepared_block(Block& output_block) { } writer_positions.erase(writer_iter->second); _partitions_to_writers.erase(writer_iter); + _publish_active_writers(); RETURN_IF_ERROR(create_and_open_writer(partition_name, i, &file_name, file_name_index + 1, writer)); } else { writer = writer_iter->second; } - auto writer_pos_iter = writer_positions.find(writer); - if (writer_pos_iter == writer_positions.end()) { - IColumn::Filter filter(output_block.rows(), 0); - filter[i] = 1; - writer_positions.insert({writer, std::move(filter)}); - } else { - writer_pos_iter->second[i] = 1; - } + } + auto writer_pos_iter = writer_positions.find(writer); + if (writer_pos_iter == writer_positions.end()) { + IColumn::Permutation rows {static_cast(i)}; + writer_positions.insert({writer, std::move(rows)}); + } else { + writer_pos_iter->second.push_back(static_cast(i)); } } } SCOPED_RAW_TIMER(&_partition_writers_write_ns); output_block.erase(_non_write_columns_indices); for (auto it = writer_positions.begin(); it != writer_positions.end(); ++it) { - Block filtered_block; - RETURN_IF_ERROR(_filter_block(output_block, &it->second, &filtered_block)); - RETURN_IF_ERROR(it->first->write(filtered_block)); - _current_writer.store(it->first); + Block selected_block; + RETURN_IF_ERROR(_select_block(output_block, it->second, &selected_block)); + RETURN_IF_ERROR(it->first->write(selected_block)); } return Status::OK(); } -Status VIcebergTableWriter::_filter_block(doris::Block& block, const IColumn::Filter* filter, +Status VIcebergTableWriter::_select_block(doris::Block& block, const IColumn::Permutation& rows, doris::Block* output_block) { const ColumnsWithTypeAndName& columns_with_type_and_name = block.get_columns_with_type_and_name(); ColumnsWithTypeAndName result_columns; + result_columns.reserve(columns_with_type_and_name.size()); for (const auto& col : columns_with_type_and_name) { - result_columns.emplace_back(col.column->clone_resized(col.column->size()), col.type, - col.name); + // Across all partitions the permutations contain exactly one entry per input row, avoiding O(P*C*R). + result_columns.emplace_back(col.column->permute(rows, rows.size()), col.type, col.name); } *output_block = {std::move(result_columns)}; + return Status::OK(); +} - std::vector columns_to_filter; - int column_to_keep = output_block->columns(); - columns_to_filter.resize(column_to_keep); - for (uint32_t i = 0; i < column_to_keep; ++i) { - columns_to_filter[i] = i; +void VIcebergTableWriter::_publish_active_writers() { + auto snapshot = std::make_shared(); + snapshot->reserve(_partitions_to_writers.size()); + for (const auto& entry : _partitions_to_writers) { + snapshot->push_back(entry.second); } - - Block::filter_block_internal(output_block, columns_to_filter, *filter); - return Status::OK(); + _active_writers.store(std::move(snapshot)); } Status VIcebergTableWriter::close(Status status) { @@ -474,6 +477,7 @@ Status VIcebergTableWriter::close(Status status) { } } _partitions_to_writers.clear(); + _publish_active_writers(); } if (status.ok()) { SCOPED_TIMER(_operator_profile->total_time_counter()); diff --git a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h index 070019d85db9be..4929f976b71e41 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h @@ -72,20 +72,14 @@ class VIcebergTableWriter final : public AsyncResultWriter { TIcebergWriteType::type write_type() const { return _write_type; } - // Getter for the current partition writer. - // Used by SpillIcebergTableSinkLocalState to access the current writer for - // memory management operations (get_reserve_mem_size, revocable_mem_size, etc.). - // Returns a snapshot by value: the async writer thread updates _current_writer - // concurrently with the spill/revoke path, so callers must hold their own copy - // while operating on it instead of dereferencing the underlying member directly. - std::shared_ptr current_writer() const { return _current_writer.load(); } + using ActiveWriterSnapshot = std::vector>; + std::shared_ptr active_writers() const { return _active_writers.load(); } private: - // The currently active partition writer (may be VIcebergPartitionWriter or VIcebergSortWriter). - // Updated during write() to track which writer received the most recent data. - // Wrapped in atomic_shared_ptr because revoke_memory / get_revocable_mem_size run on - // a different thread than the async writer that assigns to it. - doris::atomic_shared_ptr _current_writer; + friend class VIcebergTableWriterTest; + + // The spill thread needs a stable view of every partition sorter, while the async writer owns the map. + doris::atomic_shared_ptr _active_writers; class IcebergPartitionColumn { public: IcebergPartitionColumn(const iceberg::PartitionField& field, @@ -140,8 +134,9 @@ class VIcebergTableWriter final : public AsyncResultWriter { std::string _compute_file_name(); - Status _filter_block(doris::Block& block, const IColumn::Filter* filter, + Status _select_block(doris::Block& block, const IColumn::Permutation& rows, doris::Block* output_block); + void _publish_active_writers(); Status _write_prepared_block(Block& output_block); Status _process_row_lineage_columns(Block& block); diff --git a/be/src/exec/sort/sorter.cpp b/be/src/exec/sort/sorter.cpp index 2d9304adfa2f8e..c64d07b219a6cf 100644 --- a/be/src/exec/sort/sorter.cpp +++ b/be/src/exec/sort/sorter.cpp @@ -202,7 +202,12 @@ bool FullSorter::has_enough_capacity(Block* input_block, Block* unsorted_block) } size_t FullSorter::get_reserve_mem_size(RuntimeState* state, bool eos) const { - size_t size_to_reserve = 0; + return get_reserve_mem_size_components(state, eos).total(); +} + +SorterReserveMemory FullSorter::get_reserve_mem_size_components(RuntimeState* state, + bool eos) const { + SorterReserveMemory reserve; const auto rows = _state->unsorted_block()->rows(); if (rows != 0) { const auto bytes = _state->unsorted_block()->bytes(); @@ -213,24 +218,24 @@ size_t FullSorter::get_reserve_mem_size(RuntimeState* state, bool eos) const { auto new_rows = rows + state->batch_size(); // If the new size is greater than 85% of allocalted bytes, it maybe need to realloc. if ((new_block_bytes * 100 / allocated_bytes) >= 85) { - size_to_reserve += (size_t)(allocated_bytes * 1.15); + reserve.retained_growth += (size_t)(allocated_bytes * 1.15); } auto sort = new_rows > _buffered_block_size || new_block_bytes > _buffered_block_bytes; if (sort) { // new column is created when doing sort, reserve average size of one column // for estimation - size_to_reserve += new_block_bytes / _state->unsorted_block()->columns(); + reserve.transient_workspace += new_block_bytes / _state->unsorted_block()->columns(); // helping data structures used during sorting - size_to_reserve += new_rows * sizeof(IColumn::Permutation::value_type); + reserve.transient_workspace += new_rows * sizeof(IColumn::Permutation::value_type); auto sort_columns_count = _ordering_expr_ctxs.size(); if (1 != sort_columns_count) { - size_to_reserve += new_rows * sizeof(EqualRangeIterator); + reserve.transient_workspace += new_rows * sizeof(EqualRangeIterator); } } } - return size_to_reserve; + return reserve; } Status FullSorter::append_block(Block* block) { diff --git a/be/src/exec/sort/sorter.h b/be/src/exec/sort/sorter.h index 1651247eecc1ab..5c748f86a7f858 100644 --- a/be/src/exec/sort/sorter.h +++ b/be/src/exec/sort/sorter.h @@ -39,6 +39,13 @@ #include "runtime/runtime_state.h" namespace doris { + +struct SorterReserveMemory { + size_t retained_growth = 0; + size_t transient_workspace = 0; + + size_t total() const { return retained_growth + transient_workspace; } +}; class ObjectPool; class RowDescriptor; } // namespace doris @@ -194,6 +201,8 @@ class FullSorter final : public Sorter { size_t get_reserve_mem_size(RuntimeState* state, bool eos) const override; + SorterReserveMemory get_reserve_mem_size_components(RuntimeState* state, bool eos) const; + Status merge_sort_read_for_spill(RuntimeState* state, doris::Block* block, int batch_size, bool* eos) override; void reset() override; diff --git a/be/src/io/fs/azure_obj_storage_client.cpp b/be/src/io/fs/azure_obj_storage_client.cpp index 9702c87b3b304b..daa215e3c58d9e 100644 --- a/be/src/io/fs/azure_obj_storage_client.cpp +++ b/be/src/io/fs/azure_obj_storage_client.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include "common/exception.h" @@ -46,8 +47,8 @@ #include "cpp/obj_retry_strategy.h" #include "io/fs/obj_storage_client.h" #include "util/bvar_helper.h" -#include "util/coding.h" #include "util/s3_util.h" +#include "util/uuid_generator.h" using namespace Azure::Storage::Blobs; @@ -64,10 +65,13 @@ std::string to_lower_ascii(std::string_view input) { return lowered; } -auto base64_encode_part_num(int part_num) { - uint8_t buf[4]; - doris::encode_fixed32_le(buf, static_cast(part_num)); - return Aws::Utils::HashingUtils::Base64Encode({buf, sizeof(buf)}); +std::string azure_block_id(const doris::io::ObjectStoragePathOptions& opts, int part_num) { + DCHECK(opts.upload_id.has_value()); + // Azure requires every block ID for one blob to have the same decoded length. + std::string raw_id = fmt::format("{}:{:010}", *opts.upload_id, part_num); + Aws::Utils::ByteBuffer bytes(reinterpret_cast(raw_id.data()), + raw_id.size()); + return Aws::Utils::HashingUtils::Base64Encode(bytes); } // Rate limiting is applied by RateLimitedObjStorageClient, the decorator that @@ -194,11 +198,13 @@ struct AzureBatchDeleter { std::vector> deferred_resps; }; -// Azure would do nothing ObjectStorageUploadResponse AzureObjStorageClient::create_multipart_upload( const ObjectStoragePathOptions& opts) { + std::stringstream upload_id; + upload_id << UUIDGenerator::instance()->next_uuid(); return ObjectStorageUploadResponse { .resp = ObjectStorageResponse::OK(), + .upload_id = upload_id.str(), }; } @@ -217,17 +223,20 @@ ObjectStorageUploadResponse AzureObjStorageClient::upload_part(const ObjectStora std::string_view stream, int part_num) { auto client = _client->GetBlockBlobClient(opts.key); + std::string block_id = azure_block_id(opts, part_num); auto resp = do_azure_client_call( [&]() { Azure::Core::IO::MemoryBodyStream memory_body( reinterpret_cast(stream.data()), stream.size()); // The blockId must be base64 encoded SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency); - client.StageBlock(base64_encode_part_num(part_num), memory_body); + client.StageBlock(block_id, memory_body); }, opts, _tls_debug_context); return ObjectStorageUploadResponse { .resp = resp, + // Hive defers completion to FE, so the exact staged ID must cross that boundary. + .etag = block_id, }; } @@ -238,7 +247,7 @@ ObjectStorageResponse AzureObjStorageClient::complete_multipart_upload( std::vector string_block_ids; std::ranges::transform( completed_parts, std::back_inserter(string_block_ids), - [](const ObjectCompleteMultiPart& i) { return base64_encode_part_num(i.part_num); }); + [&opts](const ObjectCompleteMultiPart& i) { return azure_block_id(opts, i.part_num); }); return do_azure_client_call( [&]() { SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency); @@ -247,6 +256,41 @@ ObjectStorageResponse AzureObjStorageClient::complete_multipart_upload( opts, _tls_debug_context); } +ObjectStorageResponse AzureObjStorageClient::abort_multipart_upload( + const ObjectStoragePathOptions& opts) { + auto client = _client->GetBlockBlobClient(opts.key); + auto response = do_azure_client_call( + [&]() { + GetBlockListOptions get_options; + get_options.ListType = Models::BlockListType::All; + auto block_list = client.GetBlockList(get_options); + const bool has_committed_blob = azure_block_list_has_committed_blob( + block_list.Value.CommittedBlocks.size(), block_list.Value.ETag.HasValue()); + if (!has_committed_blob) { + // Uncommitted blocks are invisible and expire without deleting a racing commit. + return; + } + if (block_list.Value.CommittedBlocks.empty()) { + // Azure cannot selectively discard staged blocks without replacing Put Blob content. + return; + } + std::vector committed_ids; + committed_ids.reserve(block_list.Value.CommittedBlocks.size()); + std::ranges::transform(block_list.Value.CommittedBlocks, + std::back_inserter(committed_ids), + [](const Models::BlobBlock& block) { return block.Name; }); + CommitBlockListOptions commit_options; + commit_options.AccessConditions.IfMatch = block_list.Value.ETag; + // Recommitting only the old IDs discards this writer's unique staged blocks. + client.CommitBlockList(committed_ids, commit_options); + }, + opts, _tls_debug_context); + // Azure creates no server-side object until the first block is staged, so absence is clean. + return response.http_code == static_cast(Azure::Core::Http::HttpStatusCode::NotFound) + ? ObjectStorageResponse::OK() + : response; +} + ObjectStorageHeadResponse AzureObjStorageClient::head_object(const ObjectStoragePathOptions& opts) { Models::BlobProperties properties {}; auto resp = do_azure_client_call( diff --git a/be/src/io/fs/azure_obj_storage_client.h b/be/src/io/fs/azure_obj_storage_client.h index 7d1cecc502e44d..fe3cf01417aa3b 100644 --- a/be/src/io/fs/azure_obj_storage_client.h +++ b/be/src/io/fs/azure_obj_storage_client.h @@ -33,6 +33,10 @@ class ObjClientHolder; bool is_azure_tls_ca_error_message(std::string_view message); std::string build_azure_tls_debug_suffix(std::string_view error_message, std::string_view tls_debug_context); +inline bool azure_block_list_has_committed_blob(size_t committed_block_count, bool has_etag) { + // Put Blob creates committed content without block IDs; its ETag is the only safe discriminator. + return committed_block_count > 0 || has_etag; +} class AzureObjStorageClient final : public ObjStorageClient { public: @@ -49,6 +53,7 @@ class AzureObjStorageClient final : public ObjStorageClient { ObjectStorageResponse complete_multipart_upload( const ObjectStoragePathOptions& opts, const std::vector& completed_parts) override; + ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions& opts) override; ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) override; ObjectStorageResponse get_object(const ObjectStoragePathOptions& opts, void* buffer, size_t offset, size_t bytes_read, diff --git a/be/src/io/fs/file_writer.h b/be/src/io/fs/file_writer.h index 9402fdef18303c..08754ec3689a0f 100644 --- a/be/src/io/fs/file_writer.h +++ b/be/src/io/fs/file_writer.h @@ -73,6 +73,9 @@ class FileWriter { // If there is no data appended, an empty file will be persisted. virtual Status close(bool non_block = false) = 0; + // Abandon an unpublished file. Remote writers should cancel multipart state instead of completing it. + virtual Status abort() { return close(); } + // Non-blocking probe for a previous close(true). // OK means close finished successfully. NeedSendAgain means close is still running. // Other errors mean close finished with error or the writer does not support this API. diff --git a/be/src/io/fs/obj_storage_client.h b/be/src/io/fs/obj_storage_client.h index fa239ca3282e2a..db326a931719f9 100644 --- a/be/src/io/fs/obj_storage_client.h +++ b/be/src/io/fs/obj_storage_client.h @@ -44,7 +44,7 @@ struct ObjectStoragePathOptions { std::string bucket = std::string(); // blob container in azure std::string key = std::string(); // blob name in azure std::string prefix = std::string(); // for batch delete and recursive delete - std::optional upload_id = std::nullopt; // only used for S3 upload + std::optional upload_id = std::nullopt; // provider-specific upload token }; struct ObjectCompleteMultiPart { @@ -86,7 +86,7 @@ struct ObjectStorageHeadResponse : ObjectStorageResponse { class ObjStorageClient { public: virtual ~ObjStorageClient() = default; - // Create a multi-part upload request. On AWS-compatible systems, it will return an upload ID, but not on Azure. + // Create a multi-part upload request. The returned provider token identifies this upload's parts. // The input parameters should include the bucket and key for the object storage. virtual ObjectStorageUploadResponse create_multipart_upload( const ObjectStoragePathOptions& opts) = 0; @@ -106,6 +106,10 @@ class ObjStorageClient { virtual ObjectStorageResponse complete_multipart_upload( const ObjectStoragePathOptions& opts, const std::vector& completed_parts) = 0; + virtual ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions&) { + return {.status = {.code = ErrorCode::NOT_IMPLEMENTED_ERROR, + .msg = "abort multipart upload is not supported"}}; + } // According to the passed bucket and key, it will access whether the corresponding file exists in the object storage. // If it exists, it will return the corresponding file size virtual ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) = 0; diff --git a/be/src/io/fs/rate_limited_obj_storage_client.cpp b/be/src/io/fs/rate_limited_obj_storage_client.cpp index 1b8730847162df..218c39cb4b19ec 100644 --- a/be/src/io/fs/rate_limited_obj_storage_client.cpp +++ b/be/src/io/fs/rate_limited_obj_storage_client.cpp @@ -73,6 +73,12 @@ ObjectStorageResponse RateLimitedObjStorageClient::complete_multipart_upload( return _inner->complete_multipart_upload(opts, completed_parts); } +ObjectStorageResponse RateLimitedObjStorageClient::abort_multipart_upload( + const ObjectStoragePathOptions& opts) { + // Cleanup must reach the provider even when a hard PUT limit caused the upload failure. + return _inner->abort_multipart_upload(opts); +} + ObjectStorageHeadResponse RateLimitedObjStorageClient::head_object( const ObjectStoragePathOptions& opts) { S3RateLimitGuard guard(S3RateLimitType::GET, 0); diff --git a/be/src/io/fs/rate_limited_obj_storage_client.h b/be/src/io/fs/rate_limited_obj_storage_client.h index 00725d7edcb299..dc6fb1503c375d 100644 --- a/be/src/io/fs/rate_limited_obj_storage_client.h +++ b/be/src/io/fs/rate_limited_obj_storage_client.h @@ -50,6 +50,7 @@ class RateLimitedObjStorageClient final : public ObjStorageClient { ObjectStorageResponse complete_multipart_upload( const ObjectStoragePathOptions& opts, const std::vector& completed_parts) override; + ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions& opts) override; ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) override; ObjectStorageResponse get_object(const ObjectStoragePathOptions& opts, void* buffer, size_t offset, size_t bytes_read, diff --git a/be/src/io/fs/s3_file_writer.cpp b/be/src/io/fs/s3_file_writer.cpp index f8b836607a14a6..43663afa6dfdf8 100644 --- a/be/src/io/fs/s3_file_writer.cpp +++ b/be/src/io/fs/s3_file_writer.cpp @@ -78,6 +78,8 @@ S3FileWriter::~S3FileWriter() { // For thread safety std::ignore = _async_close_pack->future.get(); _async_close_pack = nullptr; + } else if (state() == State::OPENED) { + WARN_IF_ERROR(abort(), "failed to abort unfinished S3 writer"); } else { // Consider one situation where the file writer is destructed after it submit at least one async task // without calling close(), then there exists one occasion where the async task is executed right after @@ -85,13 +87,42 @@ S3FileWriter::~S3FileWriter() { _wait_until_finish(fmt::format("wait s3 file {} upload to be finished", _obj_storage_path_opts.path.native())); } - // We won't do S3 abort operation in BE, we let s3 service do it own. if (state() == State::OPENED && !_failed) { s3_bytes_written_total << _bytes_appended; } s3_file_being_written << -1; } +Status S3FileWriter::abort() { + if (state() == State::CLOSED) { + return Status::OK(); + } + if (state() == State::ASYNC_CLOSING) { + return Status::InternalError("cannot abort an asynchronously closing S3 writer"); + } + RETURN_IF_ERROR(_abort_impl()); + _state = State::CLOSED; + return Status::OK(); +} + +Status S3FileWriter::_abort_impl() { + _wait_until_finish( + fmt::format("wait s3 file {} before abort", _obj_storage_path_opts.path.native())); + _pending_buf.reset(); + if (_multipart_upload_started) { + const auto& client = _obj_client->get(); + if (client == nullptr) { + return Status::InternalError("invalid obj storage client"); + } + auto response = client->abort_multipart_upload(_obj_storage_path_opts); + if (response.status.code != ErrorCode::OK) { + return {response.status.code, std::move(response.status.msg)}; + } + } + // Once abort returns, no destructor or retry may complete the abandoned upload. + return Status::OK(); +} + Status S3FileWriter::_create_multi_upload_request() { LOG(INFO) << "create_multi_upload_request " << _obj_storage_path_opts.path.native(); const auto& client = _obj_client->get(); @@ -100,6 +131,8 @@ Status S3FileWriter::_create_multi_upload_request() { } auto resp = client->create_multipart_upload(_obj_storage_path_opts); if (resp.resp.status.code == ErrorCode::OK) { + // Some providers identify staged uploads by block IDs instead of a server-issued upload ID. + _multipart_upload_started = true; _obj_storage_path_opts.upload_id = resp.upload_id; } return {resp.resp.status.code, std::move(resp.resp.status.msg)}; @@ -162,6 +195,10 @@ Status S3FileWriter::close(bool non_block) { s3_file_writer_async_close_queuing << -1; s3_file_writer_async_close_processing << 1; _st = _close_impl(); + if (!_st.ok()) { + // A failed completion must not leave server-side multipart state behind. + WARN_IF_ERROR(_abort_impl(), "failed to abort incomplete S3 upload"); + } _async_close_pack->promise.set_value(_st); s3_file_writer_async_close_processing << -1; }); @@ -172,12 +209,18 @@ Status S3FileWriter::close(bool non_block) { << _obj_storage_path_opts.path.native() << ", fallback to sync close, status=" << submit_status; _st = _close_impl(); + if (!_st.ok()) { + WARN_IF_ERROR(_abort_impl(), "failed to abort incomplete S3 upload"); + } _async_close_pack->promise.set_value(_st); return _st; } return Status::OK(); } _st = _close_impl(); + if (!_st.ok()) { + WARN_IF_ERROR(_abort_impl(), "failed to abort incomplete S3 upload"); + } _state = State::CLOSED; if (!non_block && _st.ok()) { _record_close_latency(); diff --git a/be/src/io/fs/s3_file_writer.h b/be/src/io/fs/s3_file_writer.h index 5a8075e03cf404..eb40772eea095e 100644 --- a/be/src/io/fs/s3_file_writer.h +++ b/be/src/io/fs/s3_file_writer.h @@ -71,11 +71,12 @@ class S3FileWriter final : public FileWriter { } Status close(bool non_block = false) override; + Status abort() override; Status try_finish_close() override; private: + Status _abort_impl(); Status _close_impl(); - Status _abort(); [[nodiscard]] std::string _dump_completed_part() const; void _wait_until_finish(std::string_view task_name); Status _complete(); @@ -118,6 +119,7 @@ class S3FileWriter final : public FileWriter { std::shared_ptr _obj_client; std::optional _first_append_timestamp; bool _close_latency_recorded = false; + bool _multipart_upload_started = false; }; } // namespace io diff --git a/be/src/io/fs/s3_obj_storage_client.cpp b/be/src/io/fs/s3_obj_storage_client.cpp index 0c0b0370f8097f..54ba6b687e9790 100644 --- a/be/src/io/fs/s3_obj_storage_client.cpp +++ b/be/src/io/fs/s3_obj_storage_client.cpp @@ -275,6 +275,25 @@ ObjectStorageResponse S3ObjStorageClient::complete_multipart_upload( return ObjectStorageResponse::OK(); } +ObjectStorageResponse S3ObjStorageClient::abort_multipart_upload( + const ObjectStoragePathOptions& opts) { + AbortMultipartUploadRequest request; + request.WithBucket(opts.bucket).WithKey(opts.key).WithUploadId(*opts.upload_id); + auto outcome = SYNC_POINT_HOOK_RETURN_VALUE(_client->AbortMultipartUpload(request), + "s3_file_writer::abort_multi_part", + std::cref(request).get()); + if (!outcome.IsSuccess()) { + record_s3_request_failed(outcome.GetError()); + auto status = s3fs_error(outcome.GetError(), + fmt::format("failed to AbortMultipartUpload: {}, upload_id={}", + opts.path.native(), *opts.upload_id)); + return {convert_to_obj_response(std::move(status)), + static_cast(outcome.GetError().GetResponseCode()), + outcome.GetError().GetRequestId()}; + } + return ObjectStorageResponse::OK(); +} + ObjectStorageHeadResponse S3ObjStorageClient::head_object(const ObjectStoragePathOptions& opts) { Aws::S3::Model::HeadObjectRequest request; request.WithBucket(opts.bucket).WithKey(opts.key); diff --git a/be/src/io/fs/s3_obj_storage_client.h b/be/src/io/fs/s3_obj_storage_client.h index 45294226594d81..10bcf6b2e9495b 100644 --- a/be/src/io/fs/s3_obj_storage_client.h +++ b/be/src/io/fs/s3_obj_storage_client.h @@ -43,6 +43,7 @@ class S3ObjStorageClient final : public ObjStorageClient { ObjectStorageResponse complete_multipart_upload( const ObjectStoragePathOptions& opts, const std::vector& completed_parts) override; + ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions& opts) override; ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) override; ObjectStorageResponse get_object(const ObjectStoragePathOptions& opts, void* buffer, size_t offset, size_t bytes_read, diff --git a/be/src/runtime/runtime_state.cpp b/be/src/runtime/runtime_state.cpp index 5dfd027d42d4ce..fde7ca3bc2d0ad 100644 --- a/be/src/runtime/runtime_state.cpp +++ b/be/src/runtime/runtime_state.cpp @@ -52,12 +52,37 @@ #include "runtime/thread_context.h" #include "storage/id_manager.h" #include "storage/storage_engine.h" +#include "util/thrift_util.h" #include "util/timezone_utils.h" #include "util/uid_util.h" namespace doris { using namespace ErrorCode; +Status RuntimeState::add_iceberg_commit_datas(TIcebergCommitData iceberg_commit_data) { + ThriftSerializer serializer(false, 256); + uint32_t serialized_size = 0; + uint8_t* buffer = nullptr; + RETURN_IF_ERROR(serializer.serialize(&iceberg_commit_data, &serialized_size, &buffer)); + + constexpr size_t report_envelope_headroom = 1024 * 1024; + const size_t thrift_limit = static_cast(std::max(config::thrift_max_message_size, 0)); + const size_t commit_data_limit = + thrift_limit > report_envelope_headroom ? thrift_limit - report_envelope_headroom : 0; + std::lock_guard budget_lock(_iceberg_commit_data_budget->mutex); + // Parallel task states share this budget because FE receives their vectors in one fragment report. + if (_iceberg_commit_data_budget->serialized_bytes + serialized_size + sizeof(uint32_t) > + commit_data_limit) { + return Status::InternalError( + "Iceberg commit metadata exceeds the Thrift report limit; reduce output file " + "count"); + } + std::lock_guard data_lock(_iceberg_commit_datas_mutex); + _iceberg_commit_data_budget->serialized_bytes += serialized_size + sizeof(uint32_t); + _iceberg_commit_datas.emplace_back(std::move(iceberg_commit_data)); + return Status::OK(); +} + RuntimeState::RuntimeState(const TPlanFragmentExecParams& fragment_exec_params, const TQueryOptions& query_options, const TQueryGlobals& query_globals, ExecEnv* exec_env, QueryContext* ctx, diff --git a/be/src/runtime/runtime_state.h b/be/src/runtime/runtime_state.h index 0551c2e533689c..faa3ce1a7ef040 100644 --- a/be/src/runtime/runtime_state.h +++ b/be/src/runtime/runtime_state.h @@ -74,6 +74,14 @@ class RuntimeFilterConsumer; class RuntimeFilterProducer; class TaskExecutionContext; +class IcebergCommitDataBudget { + friend class RuntimeState; + +private: + std::mutex mutex; + size_t serialized_bytes = 0; +}; + // A collection of items that are part of the global state of a // query and shared across all execution nodes of that query. class RuntimeState { @@ -523,14 +531,19 @@ class RuntimeState { _hive_partition_updates.emplace_back(hive_partition_update); } - std::vector iceberg_commit_datas() const { + void append_iceberg_commit_datas(std::vector* output) const { std::lock_guard lock(_iceberg_commit_datas_mutex); - return _iceberg_commit_datas; + output->insert(output->end(), _iceberg_commit_datas.begin(), _iceberg_commit_datas.end()); } - void add_iceberg_commit_datas(const TIcebergCommitData& iceberg_commit_data) { - std::lock_guard lock(_iceberg_commit_datas_mutex); - _iceberg_commit_datas.emplace_back(iceberg_commit_data); + Status add_iceberg_commit_datas(TIcebergCommitData iceberg_commit_data); + + void set_iceberg_commit_data_budget(std::shared_ptr budget) { + _iceberg_commit_data_budget = std::move(budget); + } + + const std::shared_ptr& iceberg_commit_data_budget() const { + return _iceberg_commit_data_budget; } std::vector mc_commit_datas() const { @@ -976,6 +989,8 @@ class RuntimeState { mutable std::mutex _iceberg_commit_datas_mutex; std::vector _iceberg_commit_datas; + std::shared_ptr _iceberg_commit_data_budget = + std::make_shared(); mutable std::mutex _mc_commit_datas_mutex; std::vector _mc_commit_datas; diff --git a/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp b/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp new file mode 100644 index 00000000000000..5d1fa473205f0b --- /dev/null +++ b/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp @@ -0,0 +1,39 @@ +// 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 + +#include "exec/operator/iceberg_sorter_reserve_memory.h" + +namespace doris { + +TEST(SpillIcebergTableSinkOperatorTest, BoundsManyPartitionReservationToOneInputBlock) { + std::vector per_partition_reservations( + 128, {.retained_growth = 0, .transient_workspace = 8 * 1024 * 1024}); + + EXPECT_EQ(8 * 1024 * 1024, bounded_iceberg_reserve_size(per_partition_reservations)); +} + +TEST(SpillIcebergTableSinkOperatorTest, AccumulatesRetainedGrowthAcrossTouchedPartitions) { + std::vector per_partition_reservations { + {.retained_growth = 3 * 1024 * 1024, .transient_workspace = 7 * 1024 * 1024}, + {.retained_growth = 4 * 1024 * 1024, .transient_workspace = 5 * 1024 * 1024}}; + + EXPECT_EQ(14 * 1024 * 1024, bounded_iceberg_reserve_size(per_partition_reservations)); +} + +} // namespace doris diff --git a/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp b/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp index d453177cf25044..18af0af7cb23bd 100644 --- a/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp +++ b/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp @@ -20,6 +20,8 @@ #include #include "exec/sink/writer/iceberg/viceberg_partition_writer.h" +#include "exec/sink/writer/iceberg/viceberg_sort_writer.h" +#include "testutil/mock/mock_runtime_state.h" namespace doris { @@ -27,13 +29,18 @@ namespace { class FakeFileFormatTransformer final : public VFileFormatTransformer { public: - explicit FakeFileFormatTransformer(const VExprContextSPtrs& output_exprs) - : VFileFormatTransformer(nullptr, output_exprs, false) {} + explicit FakeFileFormatTransformer(const VExprContextSPtrs& output_exprs, + Status close_status = Status::OK()) + : VFileFormatTransformer(nullptr, output_exprs, false), + _close_status(std::move(close_status)) {} Status open() override { return Status::OK(); } Status write(const Block&) override { return Status::OK(); } - Status close() override { return Status::OK(); } + Status close() override { return _close_status; } int64_t written_len() override { return 64; } + +private: + Status _close_status; }; TDataSink make_table_sink(std::optional collect_column_stats) { @@ -64,9 +71,10 @@ class VIcebergPartitionWriterTest : public testing::Test { } static void install_fake_transformer(VIcebergPartitionWriter* writer, - const VExprContextSPtrs& output_exprs) { + const VExprContextSPtrs& output_exprs, + Status close_status = Status::OK()) { writer->_file_format_transformer = - std::make_unique(output_exprs); + std::make_unique(output_exprs, std::move(close_status)); } static Status build_commit_data(VIcebergPartitionWriter* writer, @@ -104,4 +112,23 @@ TEST_F(VIcebergPartitionWriterTest, MissingPolicyKeepsCollectionEnabledForRollin EXPECT_TRUE(collect_column_stats(*writer)); } +TEST_F(VIcebergPartitionWriterTest, SortWriterPropagatesUnderlyingCloseFailure) { + VExprContextSPtrs output_exprs; + iceberg::Schema schema(std::vector {}); + std::string schema_json; + std::map hadoop_conf; + auto partition_writer = std::shared_ptr( + make_writer(make_table_sink(false), output_exprs, schema, &schema_json, hadoop_conf)); + install_fake_transformer(partition_writer.get(), output_exprs, + Status::IOError("injected close failure")); + VIcebergSortWriter sort_writer(partition_writer, TSortInfo(), 1024); + MockRuntimeState state; + sort_writer._runtime_state = &state; + + Status status = sort_writer.close(Status::OK()); + + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("injected close failure"), std::string::npos); +} + } // namespace doris diff --git a/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp b/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp new file mode 100644 index 00000000000000..e38da59cfbc176 --- /dev/null +++ b/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp @@ -0,0 +1,144 @@ +// 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 + +#include +#include + +#include "core/column/column_vector.h" +#include "core/data_type/data_type_number.h" +#include "exec/sink/writer/iceberg/viceberg_table_writer.h" +#include "exec/sink/writer/iceberg/vpartition_writer_base.h" + +namespace doris { + +namespace { + +class FakePartitionWriter final : public IPartitionWriterBase { +public: + explicit FakePartitionWriter(std::atomic* destroyed = nullptr) : _destroyed(destroyed) {} + ~FakePartitionWriter() override { + if (_destroyed != nullptr) { + ++(*_destroyed); + } + } + Status open(RuntimeState*, RuntimeProfile*, const RowDescriptor*) override { + return Status::OK(); + } + Status write(Block&) override { return Status::OK(); } + Status close(const Status&) override { return Status::OK(); } + const std::string& file_name() const override { return _name; } + int file_name_index() const override { return 0; } + size_t written_len() const override { return 0; } + +private: + std::string _name = "fake"; + std::atomic* _destroyed; +}; + +TDataSink make_sink() { + TDataSink sink; + sink.__set_type(TDataSinkType::ICEBERG_TABLE_SINK); + sink.__set_iceberg_table_sink(TIcebergTableSink()); + return sink; +} + +} // namespace + +class VIcebergTableWriterTest : public testing::Test { +protected: + static Status select_block(VIcebergTableWriter* writer, Block& input, + const IColumn::Permutation& rows, Block* selected) { + return writer->_select_block(input, rows, selected); + } + + static void add_writer(VIcebergTableWriter* writer, std::string partition) { + writer->_partitions_to_writers.emplace(std::move(partition), + std::make_shared()); + } + + static void add_writer(VIcebergTableWriter* writer, std::string partition, + std::shared_ptr partition_writer) { + writer->_partitions_to_writers.emplace(std::move(partition), std::move(partition_writer)); + } + + static void clear_writers(VIcebergTableWriter* writer) { + writer->_partitions_to_writers.clear(); + } + + static void publish_active_writers(VIcebergTableWriter* writer) { + writer->_publish_active_writers(); + } +}; + +TEST_F(VIcebergTableWriterTest, SelectBlockUsesRowPermutation) { + VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr); + auto values = ColumnInt32::create(); + values->insert_value(10); + values->insert_value(20); + values->insert_value(30); + Block input; + input.insert({std::move(values), std::make_shared(), "value"}); + IColumn::Permutation rows {2, 0}; + Block selected; + + ASSERT_TRUE(select_block(&writer, input, rows, &selected).ok()); + + const auto& result = assert_cast(*selected.get_by_position(0).column); + ASSERT_EQ(result.size(), 2); + EXPECT_EQ(result.get_element(0), 30); + EXPECT_EQ(result.get_element(1), 10); +} + +TEST_F(VIcebergTableWriterTest, ActiveWriterSnapshotContainsEveryOpenPartition) { + VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr); + add_writer(&writer, "p=1"); + add_writer(&writer, "p=2"); + + publish_active_writers(&writer); + + ASSERT_NE(writer.active_writers(), nullptr); + EXPECT_EQ(writer.active_writers()->size(), 2); +} + +TEST_F(VIcebergTableWriterTest, LoadedSnapshotRetainsWritersDuringConcurrentPublication) { + VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr); + std::atomic destroyed = 0; + add_writer(&writer, "p=1", std::make_shared(&destroyed)); + publish_active_writers(&writer); + std::promise snapshot_loaded; + std::promise replacement_published; + + auto reader = std::async(std::launch::async, [&]() { + auto snapshot = writer.active_writers(); + snapshot_loaded.set_value(); + replacement_published.get_future().wait(); + EXPECT_EQ(1, snapshot->size()); + EXPECT_EQ("fake", snapshot->front()->file_name()); + }); + + snapshot_loaded.get_future().wait(); + clear_writers(&writer); + publish_active_writers(&writer); + EXPECT_EQ(0, destroyed.load()); + replacement_published.set_value(); + reader.get(); + EXPECT_EQ(1, destroyed.load()); +} + +} // namespace doris diff --git a/be/test/io/fs/azure_obj_storage_client_test.cpp b/be/test/io/fs/azure_obj_storage_client_test.cpp index 7591b4bf2ea997..88e81ea7fe6ed0 100644 --- a/be/test/io/fs/azure_obj_storage_client_test.cpp +++ b/be/test/io/fs/azure_obj_storage_client_test.cpp @@ -19,6 +19,8 @@ #include +#include + #include "io/fs/file_system.h" #include "io/fs/obj_storage_client.h" #include "util/s3_util.h" @@ -32,6 +34,12 @@ namespace doris { +TEST(AzureObjStorageClientAbortHelperTest, preserves_committed_put_blob_without_block_list) { + EXPECT_FALSE(io::azure_block_list_has_committed_blob(0, false)); + EXPECT_TRUE(io::azure_block_list_has_committed_blob(0, true)); + EXPECT_TRUE(io::azure_block_list_has_committed_blob(1, false)); +} + #ifdef USE_AZURE using namespace Azure::Storage::Blobs; @@ -156,6 +164,58 @@ TEST_F(AzureObjStorageClientTest, delete_objects_recursively) { EXPECT_EQ(response.status.code, ErrorCode::OK); EXPECT_EQ(files.size(), 0); } + +TEST_F(AzureObjStorageClientTest, abort_multipart_upload_discards_staged_blocks) { + io::ObjectStoragePathOptions opts; + auto create_response = + AzureObjStorageClientTest::obj_storage_client->create_multipart_upload(opts); + ASSERT_EQ(create_response.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(create_response.upload_id.has_value()); + opts.key = "AzureObjStorageClientTest/abort_multipart_upload_" + *create_response.upload_id; + opts.upload_id = create_response.upload_id; + + auto upload_response = + AzureObjStorageClientTest::obj_storage_client->upload_part(opts, "staged", 1); + ASSERT_EQ(upload_response.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(upload_response.etag.has_value()); + EXPECT_FALSE(upload_response.etag->empty()); + auto abort_response = + AzureObjStorageClientTest::obj_storage_client->abort_multipart_upload(opts); + ASSERT_EQ(abort_response.status.code, ErrorCode::OK); + + auto head_response = AzureObjStorageClientTest::obj_storage_client->head_object(opts); + EXPECT_EQ(head_response.resp.status.code, ErrorCode::NOT_FOUND); +} + +TEST_F(AzureObjStorageClientTest, abort_multipart_upload_preserves_existing_put_blob) { + io::ObjectStoragePathOptions opts; + auto create_response = + AzureObjStorageClientTest::obj_storage_client->create_multipart_upload(opts); + ASSERT_EQ(create_response.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(create_response.upload_id.has_value()); + opts.key = "AzureObjStorageClientTest/abort_preserves_put_blob_" + *create_response.upload_id; + opts.upload_id = create_response.upload_id; + + auto put_response = AzureObjStorageClientTest::obj_storage_client->put_object(opts, "original"); + ASSERT_EQ(put_response.status.code, ErrorCode::OK); + auto upload_response = + AzureObjStorageClientTest::obj_storage_client->upload_part(opts, "replacement", 1); + ASSERT_EQ(upload_response.resp.status.code, ErrorCode::OK); + + auto abort_response = + AzureObjStorageClientTest::obj_storage_client->abort_multipart_upload(opts); + ASSERT_EQ(abort_response.status.code, ErrorCode::OK); + std::array contents {}; + size_t size_return = 0; + auto get_response = AzureObjStorageClientTest::obj_storage_client->get_object( + opts, contents.data(), 0, contents.size(), &size_return); + ASSERT_EQ(get_response.status.code, ErrorCode::OK); + EXPECT_EQ(size_return, contents.size()); + EXPECT_EQ(std::string_view(contents.data(), contents.size()), "original"); + + EXPECT_EQ(AzureObjStorageClientTest::obj_storage_client->delete_object(opts).status.code, + ErrorCode::OK); +} #else class AzureObjStorageClientTest : public testing::Test { diff --git a/be/test/io/fs/rate_limited_obj_storage_client_test.cpp b/be/test/io/fs/rate_limited_obj_storage_client_test.cpp index 657b139c0a4fcd..6120347beb0f7c 100644 --- a/be/test/io/fs/rate_limited_obj_storage_client_test.cpp +++ b/be/test/io/fs/rate_limited_obj_storage_client_test.cpp @@ -62,6 +62,11 @@ class FakeObjStorageClient : public ObjStorageClient { ++calls; return ObjectStorageResponse::OK(); } + ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions& opts) override { + ++calls; + ++abort_multipart_upload_calls; + return ObjectStorageResponse::OK(); + } ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) override { ++calls; return {}; @@ -106,6 +111,7 @@ class FakeObjStorageClient : public ObjStorageClient { int create_multipart_upload_calls = 0; int create_multipart_upload_provider_calls = 0; int create_multipart_upload_provider_calls_per_logical_call = 1; + int abort_multipart_upload_calls = 0; int delete_objects_recursively_calls = 0; int delete_objects_recursively_provider_calls = 0; int delete_objects_recursively_provider_calls_per_logical_call = 1; @@ -368,6 +374,23 @@ TEST(RateLimitedObjStorageClientTest, multipart_control_apis_map_to_put_qps_with EXPECT_EQ(-1, put_bytes->add(1)); } +TEST(RateLimitedObjStorageClientTest, abortBypassesAnExhaustedPutLimit) { + RateLimiterConfigGuard guard; + config::enable_s3_rate_limiter = true; + auto& manager = S3RateLimiterManager::instance(); + manager.qps_limiter(S3RateLimitType::PUT) + ->reset(kNoThrottleBytesPerSecond, kNoThrottleBytesPerSecond, 1); + manager.bytes_limiter(S3RateLimitType::PUT)->reset(0, 0, 0); + + auto fake = std::make_shared(); + RateLimitedObjStorageClient client(fake); + ObjectStoragePathOptions opts {.bucket = "b", .key = "k", .upload_id = "upload"}; + + EXPECT_EQ(0, client.create_multipart_upload(opts).resp.status.code); + EXPECT_EQ(0, client.abort_multipart_upload(opts).status.code); + EXPECT_EQ(1, fake->abort_multipart_upload_calls); +} + TEST(RateLimitedObjStorageClientTest, delete_apis_map_to_put_qps_without_bytes) { RateLimiterConfigGuard guard; config::enable_s3_rate_limiter = true; diff --git a/be/test/io/fs/s3_file_writer_test.cpp b/be/test/io/fs/s3_file_writer_test.cpp index 3937d6e38561fe..b00f6c491c72d4 100644 --- a/be/test/io/fs/s3_file_writer_test.cpp +++ b/be/test/io/fs/s3_file_writer_test.cpp @@ -316,6 +316,21 @@ class S3FileWriterTest : public testing::Test { } }; +TEST_F(S3FileWriterTest, abort_cleans_up_multipart_upload) { + mock_client = std::make_shared(); + doris::io::FileWriterOptions options; + + io::FileWriterPtr writer; + ASSERT_TRUE(s3_fs->create_file("abort_multipart", &writer, &options).ok()); + std::string data(config::s3_write_buffer_size, 'a'); + ASSERT_TRUE(writer->append(Slice(data)).ok()); + ASSERT_FALSE(static_cast(writer.get())->upload_id().empty()); + + ASSERT_TRUE(writer->abort().ok()); + EXPECT_EQ(writer->state(), io::FileWriter::State::CLOSED); + EXPECT_TRUE(mock_client->contents().empty()); +} + TEST_F(S3FileWriterTest, multi_part_io_error) { mock_client = std::make_shared(); doris::io::FileWriterOptions state; @@ -1154,6 +1169,14 @@ class SimpleMockObjStorageClient : public io::ObjStorageClient { return default_response; } + ObjectStorageResponse abort_multipart_upload(const ObjectStoragePathOptions& opts) override { + std::lock_guard lock(_mutex); + abort_multipart_count++; + last_opts = opts; + parts.clear(); + return default_response; + } + ObjectStorageHeadResponse head_object(const ObjectStoragePathOptions& opts) override { std::lock_guard lock(_mutex); return {.resp = ObjectStorageResponse::OK(), @@ -1228,6 +1251,7 @@ class SimpleMockObjStorageClient : public io::ObjStorageClient { int put_object_count = 0; int upload_part_count = 0; int complete_multipart_count = 0; + int abort_multipart_count = 0; // Structures to store input parameters for each call struct UploadPartParams { @@ -1266,6 +1290,7 @@ class SimpleMockObjStorageClient : public io::ObjStorageClient { put_object_count = 0; upload_part_count = 0; complete_multipart_count = 0; + abort_multipart_count = 0; create_multipart_params.clear(); put_object_params.clear(); @@ -1307,6 +1332,19 @@ create_s3_client(const std::string& path) { return {mock_client, s3_file_writer}; } +TEST_F(S3FileWriterTest, abortsProviderMultipartWithoutAnUploadId) { + auto [client, writer] = create_s3_client("provider_without_upload_id"); + client->default_upload_response.upload_id.reset(); + std::string data(config::s3_write_buffer_size, 'a'); + + ASSERT_TRUE(writer->append(Slice(data)).ok()); + ASSERT_TRUE(writer->abort().ok()); + + EXPECT_EQ(1, client->create_multipart_count); + EXPECT_EQ(1, client->abort_multipart_count); + EXPECT_EQ(FileWriter::State::CLOSED, writer->state()); +} + /** * Generate test data for S3FileWriter boundary tests. * Returns a vector of sizes that we'll use to generate data on demand. diff --git a/be/test/runtime/runtime_state_block_budget_test.cpp b/be/test/runtime/runtime_state_block_budget_test.cpp index 22ebc5ebf8a0ee..24afe920d197c5 100644 --- a/be/test/runtime/runtime_state_block_budget_test.cpp +++ b/be/test/runtime/runtime_state_block_budget_test.cpp @@ -24,6 +24,41 @@ namespace doris { +TEST(RuntimeStateIcebergCommitDataTest, RejectsMetadataBeforeItCanExceedTheThriftLimit) { + RuntimeState state; + const int32_t saved_limit = config::thrift_max_message_size; + config::thrift_max_message_size = 128; + TIcebergCommitData commit_data; + commit_data.__set_file_path(std::string(256, 'x')); + + Status status = state.add_iceberg_commit_datas(commit_data); + + config::thrift_max_message_size = saved_limit; + EXPECT_FALSE(status.ok()); + std::vector collected; + state.append_iceberg_commit_datas(&collected); + EXPECT_TRUE(collected.empty()); +} + +TEST(RuntimeStateIcebergCommitDataTest, SharesTheReportBudgetAcrossParallelTasks) { + RuntimeState first; + RuntimeState second; + auto budget = std::make_shared(); + first.set_iceberg_commit_data_budget(budget); + second.set_iceberg_commit_data_budget(budget); + const int32_t saved_limit = config::thrift_max_message_size; + config::thrift_max_message_size = 1024 * 1024 + 512; + TIcebergCommitData commit_data; + commit_data.__set_file_path(std::string(300, 'x')); + + Status first_status = first.add_iceberg_commit_datas(commit_data); + Status second_status = second.add_iceberg_commit_datas(commit_data); + + config::thrift_max_message_size = saved_limit; + EXPECT_TRUE(first_status.ok()) << first_status; + EXPECT_FALSE(second_status.ok()); +} + // --------------------------------------------------------------------------- // RuntimeState::batch_size() // --------------------------------------------------------------------------- diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadata.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadata.java index 62e98dc2a0a048..ab87d325bf937e 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadata.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadata.java @@ -2128,8 +2128,8 @@ private long parseTimestampMillis(ConnectorSession session, ConnectorTimeTravelS * Threads a resolved MVCC / time-travel pin onto the handle BEFORE the scan reads it (the generic * {@code PluginDrivenScanNode} calls this via {@code applyMvccSnapshotPin}). Reads the typed * {@code snapshotId}/{@code schemaId} and the {@code iceberg.scan.ref} property; an empty-table / query-begin - * latest pin ({@code snapshotId<0} and no ref) returns the handle UNCHANGED (read latest — a - * {@code useSnapshot(-1)} would be a non-existent snapshot; mirrors paimon's {@code -1} guard). + * latest pin ({@code snapshotId<0} and no ref) remains distinguishable from no pin while scans still + * read latest (a {@code useSnapshot(-1)} would be a non-existent snapshot). */ @Override public ConnectorTableHandle applySnapshot(ConnectorSession session, @@ -2140,9 +2140,6 @@ public ConnectorTableHandle applySnapshot(ConnectorSession session, } String ref = snapshot.getProperties().get(REF_PROPERTY); long snapshotId = snapshot.getSnapshotId(); - if (snapshotId < 0 && ref == null) { - return iceHandle; - } return iceHandle.withSnapshot(snapshotId, ref, snapshot.getSchemaId()); } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorTransaction.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorTransaction.java index 666e59a8741f8e..bd3b67d8052450 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorTransaction.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorTransaction.java @@ -148,8 +148,8 @@ public class IcebergConnectorTransaction implements ConnectorTransaction, Rewrit private boolean staticPartitionOverwrite; private Map staticPartitionValues = Collections.emptyMap(); private String branchName; - // The current snapshot pinned at begin time for a DELETE/MERGE (null for INSERT/OVERWRITE). Consumed by - // the commit validation suite (validateFromSnapshot). + // The snapshot pinned at begin time for DELETE/MERGE and OVERWRITE (null for INSERT). Consumed by the + // commit validation suite (validateFromSnapshot). private Long baseSnapshotId; // Session zone for human-readable TIMESTAMP partition value parsing (DV-T04-f). private ZoneId zone = ZoneOffset.UTC; @@ -294,7 +294,6 @@ private void applyBeginGuards(IcebergWriteContext ctx, String tableName) { } } else { // INSERT / OVERWRITE (append path). - this.baseSnapshotId = null; if (ctx.getBranchName().isPresent()) { this.branchName = ctx.getBranchName().get(); SnapshotRef branchRef = table.refs().get(branchName); @@ -304,10 +303,49 @@ private void applyBeginGuards(IcebergWriteContext ctx, String tableName) { throw new IllegalArgumentException(branchName + " is a tag, not a branch. Tags cannot be targets for producing snapshots"); } + this.baseSnapshotId = op == WriteOperation.OVERWRITE + ? resolveOverwriteBaseSnapshot(ctx, branchRef.snapshotId(), tableName) : null; } else { this.branchName = null; + this.baseSnapshotId = op == WriteOperation.OVERWRITE + ? resolveOverwriteBaseSnapshot(ctx, getSnapshotIdIfPresent(table), tableName) : null; + } + } + } + + private Long resolveOverwriteBaseSnapshot(IcebergWriteContext ctx, Long targetHead, String tableName) { + if (!ctx.isReadSnapshotPinned()) { + return targetHead; + } + long readSnapshotId = ctx.getReadSnapshotId(); + if (readSnapshotId < 0) { + // An explicit empty read must conflict with any snapshot created before beginWrite. + if (targetHead != null) { + throw new DorisConnectorException("Iceberg table " + tableName + + " changed after the statement read an empty snapshot"); + } + return null; + } + if (targetHead == null || !isAncestorOfTarget(readSnapshotId, targetHead)) { + throw new DorisConnectorException("Read snapshot " + readSnapshotId + + " is not an ancestor of the target branch for Iceberg table " + tableName); + } + return readSnapshotId; + } + + private boolean isAncestorOfTarget(long ancestorId, long targetHeadId) { + Long snapshotId = targetHeadId; + while (snapshotId != null) { + if (snapshotId == ancestorId) { + return true; + } + Snapshot snapshot = table.snapshot(snapshotId); + if (snapshot == null) { + return false; } + snapshotId = snapshot.parentId(); } + return false; } @Override @@ -564,7 +602,13 @@ private void commitReplaceTxn(List pendingResults) { if (branchName != null) { overwriteFiles = overwriteFiles.toBranch(branchName); } - try (CloseableIterable fileScanTasks = table.newScan().planFiles()) { + // Clearing a table must fail if any data or delete landed after the statement's base snapshot. + overwriteFiles = validateOverwrite(overwriteFiles, Expressions.alwaysTrue()); + TableScan overwriteScan = table.newScan(); + if (branchName != null) { + overwriteScan = overwriteScan.useRef(branchName); + } + try (CloseableIterable fileScanTasks = overwriteScan.planFiles()) { OverwriteFiles finalOverwriteFiles = overwriteFiles; fileScanTasks.forEach(f -> finalOverwriteFiles.deleteFile(f.file())); } catch (IOException e) { @@ -579,6 +623,11 @@ private void commitReplaceTxn(List pendingResults) { if (branchName != null) { appendPartitionOp = appendPartitionOp.toBranch(branchName); } + // Partition replacement must preserve concurrent files instead of deleting or reviving them silently. + if (baseSnapshotId != null) { + appendPartitionOp = appendPartitionOp.validateFromSnapshot(baseSnapshotId); + } + appendPartitionOp = appendPartitionOp.validateNoConflictingData().validateNoConflictingDeletes(); for (WriteResult result : pendingResults) { Preconditions.checkState(result.referencedDataFiles().length == 0, "Should have no referenced data files."); @@ -603,6 +652,7 @@ private void commitStaticPartitionOverwrite(List pendingResults) { overwriteFiles = overwriteFiles.toBranch(branchName); } overwriteFiles = overwriteFiles.overwriteByRowFilter(partitionFilter); + overwriteFiles = validateOverwrite(overwriteFiles, partitionFilter); for (WriteResult result : pendingResults) { Preconditions.checkState(result.referencedDataFiles().length == 0, @@ -612,6 +662,14 @@ private void commitStaticPartitionOverwrite(List pendingResults) { overwriteFiles.commit(); } + private OverwriteFiles validateOverwrite(OverwriteFiles overwriteFiles, Expression conflictFilter) { + overwriteFiles = overwriteFiles.conflictDetectionFilter(conflictFilter); + if (baseSnapshotId != null) { + overwriteFiles = overwriteFiles.validateFromSnapshot(baseSnapshotId); + } + return overwriteFiles.validateNoConflictingData().validateNoConflictingDeletes(); + } + /** * Build an iceberg {@link Expression} from the static partition key-value pairs. Identity partitions * require the SOURCE column name (not the partition field name) in the expression. diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergProcedureOps.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergProcedureOps.java index a1daff65192e09..09501fc48a0a7e 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergProcedureOps.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergProcedureOps.java @@ -43,8 +43,7 @@ import java.util.stream.Collectors; /** - * Executes iceberg's {@code ALTER TABLE EXECUTE} procedures (the 9 legacy - * {@code datasource/iceberg/action/*} actions) behind the {@link ConnectorProcedureOps} SPI. + * Executes iceberg's {@code ALTER TABLE EXECUTE} procedures behind the {@link ConnectorProcedureOps} SPI. * *

Mirrors {@link IcebergWritePlanProvider}: a fresh instance per call over the lazily-built live * catalog, threading the same {@code properties} / {@link IcebergCatalogOps} / {@link ConnectorContext} @@ -52,14 +51,11 @@ * runs in the connector; argument validation is connector-local (the engine cannot reach * {@code org.apache.doris.common.NamedArguments} across the import gate).

* - *

T03 dispatch skeleton. {@link #getSupportedProcedures()} exports the factory's name list and + *

{@link #getSupportedProcedures()} exports the factory's name list and * {@link #execute} routes through {@link IcebergExecuteActionFactory} → {@link BaseIcebergAction}: validate * arguments, load the SDK table inside {@code context.executeAuthenticated}, run the body and wrap the - * single row. The 9 procedure bodies (the factory's switch cases) are ported in T04 (the 8 pure-SDK - * procedures) / T05–T06 ({@code rewrite_data_files}); until then a known name reaches the factory's faithful - * "Unsupported Iceberg procedure" rejection. Inert pre-cutover regardless: iceberg tables are not - * {@code PluginDrivenExternalTable} until P6.6, so {@code ExecuteActionCommand} still routes them to the - * legacy fe-core actions and never reaches this class.

+ * single row. {@code rewrite_data_files} is planned as a distributed INSERT-SELECT operation and therefore + * bypasses the single-call action factory.

*/ public class IcebergProcedureOps implements ConnectorProcedureOps { diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java index 42db23c93ac1c4..9d8b07a9d85b65 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergScanPlanProvider.java @@ -898,7 +898,7 @@ private List doPlanPositionDeletesSystemTableScan(IcebergTab Table metadataTable, List columns, Optional filter, ConnectorSession session) { BatchScan scan = metadataTable.newBatchScan(); - if (handle.hasSnapshotPin()) { + if (handle.hasSnapshotSelection()) { if (handle.getRef() != null) { scan = scan.useRef(handle.getRef()); } else { @@ -1120,6 +1120,9 @@ private TableScan buildScan(Table table, IcebergTableHandle handle, Optional - *
  • {@code snapshotId} ({@code -1} = none) — {@code FOR VERSION AS OF } / {@code FOR TIME AS OF}.
    • + *
  • {@code snapshotId} ({@code -1} = none or an explicitly empty-table pin) — + * {@code FOR VERSION AS OF } / {@code FOR TIME AS OF}.
    • *
  • {@code ref} ({@code null} = none) — a tag/branch name; the scan pins by REF ({@code useRef}) so a * later commit to the tag/branch is honored (legacy parity).
    • *
  • {@code schemaId} ({@code -1} = latest) — the schema version AS OF the pin, so the field-id dictionary @@ -53,7 +54,7 @@ public class IcebergTableHandle implements ConnectorTableHandle { private static final long serialVersionUID = 1L; - /** Sentinel for "no snapshot / latest schema" — mirrors legacy {@code IcebergUtils.UNKNOWN_SNAPSHOT_ID}. */ + /** Numeric sentinel shared by no pin and an explicit empty pin; {@link #snapshotPinned} distinguishes them. */ private static final long NO_PIN = -1L; private final String dbName; @@ -61,6 +62,7 @@ public class IcebergTableHandle implements ConnectorTableHandle { private final long snapshotId; private final String ref; private final long schemaId; + private final boolean snapshotPinned; /** * Bare system-table name (no {@code "$"}), lower-cased by the caller @@ -97,16 +99,18 @@ public class IcebergTableHandle implements ConnectorTableHandle { private final boolean topnLazyMaterialize; public IcebergTableHandle(String dbName, String tableName) { - this(dbName, tableName, NO_PIN, null, NO_PIN, null, null, false); + this(dbName, tableName, NO_PIN, null, NO_PIN, false, null, null, false); } private IcebergTableHandle(String dbName, String tableName, long snapshotId, String ref, long schemaId, - String sysTableName, Set rewriteFileScope, boolean topnLazyMaterialize) { + boolean snapshotPinned, String sysTableName, Set rewriteFileScope, + boolean topnLazyMaterialize) { this.dbName = dbName; this.tableName = tableName; this.snapshotId = snapshotId; this.ref = ref; this.schemaId = schemaId; + this.snapshotPinned = snapshotPinned; this.sysTableName = sysTableName; this.rewriteFileScope = rewriteFileScope; this.topnLazyMaterialize = topnLazyMaterialize; @@ -121,7 +125,8 @@ private IcebergTableHandle(String dbName, String tableName, long snapshotId, Str */ public static IcebergTableHandle forSystemTable(String dbName, String tableName, String sysName, long snapshotId, String ref, long schemaId) { - return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, sysName, null, false); + return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, + snapshotId >= 0 || ref != null, sysName, null, false); } public String getDbName() { @@ -132,7 +137,7 @@ public String getTableName() { return tableName; } - /** The pinned snapshot id, or {@code -1} when there is no snapshot-id pin. */ + /** The pinned snapshot id, or {@code -1} when no snapshot exists or no snapshot-id pin is present. */ public long getSnapshotId() { return snapshotId; } @@ -157,8 +162,13 @@ public boolean isSystemTable() { return sysTableName != null; } - /** Whether this handle carries an explicit MVCC / time-travel pin (a snapshot id or a tag/branch ref). */ + /** Whether this handle carries an explicit MVCC pin, including an empty-table query-begin pin. */ public boolean hasSnapshotPin() { + return snapshotPinned; + } + + /** Whether the pin selects an Iceberg snapshot/ref rather than the explicit empty-table state. */ + public boolean hasSnapshotSelection() { return snapshotId >= 0 || ref != null; } @@ -183,7 +193,7 @@ public IcebergTableHandle withSnapshot(long snapshotId, String ref, long schemaI // sysTableName, rewriteFileScope and topnLazyMaterialize are preserved: threading a resolved // time-travel pin in must not degrade a sys handle (t$snapshots) into a normal data-table handle, // drop a rewrite scope, or drop the lazy-materialization signal. - return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, sysTableName, + return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, true, sysTableName, rewriteFileScope, topnLazyMaterialize); } @@ -196,7 +206,7 @@ public IcebergTableHandle withSnapshot(long snapshotId, String ref, long schemaI * The other carriers (snapshot/ref/schema/sys) are preserved. */ public IcebergTableHandle withRewriteFileScope(Set rawDataFilePaths) { - return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, sysTableName, + return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, snapshotPinned, sysTableName, ImmutableSet.copyOf(rawDataFilePaths), topnLazyMaterialize); } @@ -205,7 +215,7 @@ public IcebergTableHandle withRewriteFileScope(Set rawDataFilePaths) { * {@link #topnLazyMaterialize}). The other carriers (snapshot/ref/schema/sys/rewriteScope) are preserved. */ public IcebergTableHandle withTopnLazyMaterialize(boolean topnLazyMaterialize) { - return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, sysTableName, + return new IcebergTableHandle(dbName, tableName, snapshotId, ref, schemaId, snapshotPinned, sysTableName, rewriteFileScope, topnLazyMaterialize); } @@ -220,6 +230,7 @@ public boolean equals(Object o) { IcebergTableHandle that = (IcebergTableHandle) o; return snapshotId == that.snapshotId && schemaId == that.schemaId + && snapshotPinned == that.snapshotPinned && topnLazyMaterialize == that.topnLazyMaterialize && Objects.equals(dbName, that.dbName) && Objects.equals(tableName, that.tableName) @@ -230,7 +241,8 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(dbName, tableName, snapshotId, ref, schemaId, sysTableName, rewriteFileScope, + return Objects.hash(dbName, tableName, snapshotId, ref, schemaId, snapshotPinned, sysTableName, + rewriteFileScope, topnLazyMaterialize); } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriteContext.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriteContext.java index 18d1445817713e..05878f9a1b6560 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriteContext.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWriteContext.java @@ -42,20 +42,28 @@ final class IcebergWriteContext { private final Map staticPartitionValues; private final Optional branchName; private final long readSnapshotId; + private final boolean readSnapshotPinned; IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, Map staticPartitionValues, Optional branchName) { - this(writeOperation, overwrite, staticPartitionValues, branchName, -1L); + this(writeOperation, overwrite, staticPartitionValues, branchName, -1L, false); } IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, Map staticPartitionValues, Optional branchName, long readSnapshotId) { + this(writeOperation, overwrite, staticPartitionValues, branchName, readSnapshotId, true); + } + + IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, + Map staticPartitionValues, Optional branchName, long readSnapshotId, + boolean readSnapshotPinned) { this.writeOperation = writeOperation; this.overwrite = overwrite; this.staticPartitionValues = staticPartitionValues == null ? Collections.emptyMap() : new HashMap<>(staticPartitionValues); this.branchName = branchName == null ? Optional.empty() : branchName; this.readSnapshotId = readSnapshotId; + this.readSnapshotPinned = readSnapshotPinned; } WriteOperation getWriteOperation() { @@ -81,7 +89,8 @@ Optional getBranchName() { /** * The statement's READ snapshot id (the MVCC pin the scan used, S_read), threaded from the write - * handle in {@code planWrite}; {@code -1} = no pin (the legacy fresh-current behavior). The + * handle in {@code planWrite}; {@code -1} means either no pin or an explicitly empty read, as + * distinguished by {@link #isReadSnapshotPinned()}. The * RowDelta path anchors {@code baseSnapshotId} at this snapshot so the commit-time removeDeletes * (option D) and the scan-time deletes BE unions into the new DV share one snapshot — see * {@link IcebergConnectorTransaction} [SHOULD-2] / Fix B. @@ -89,4 +98,8 @@ Optional getBranchName() { long getReadSnapshotId() { return readSnapshotId; } + + boolean isReadSnapshotPinned() { + return readSnapshotPinned; + } } diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java index d9410c3dc93621..aa9b12a759910a 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java @@ -369,13 +369,15 @@ private IcebergWriteContext buildWriteContext(ConnectorWriteHandle handle) { // Carry it on the op-context so beginWrite anchors the RowDelta baseSnapshotId at S_read, keeping // the commit-time removeDeletes (option D) and BE's scan-time DV union on one snapshot. -1 (no pin) // preserves the legacy begin-time current snapshot. - long readSnapshotId = handle.getTableHandle() instanceof IcebergTableHandle - ? ((IcebergTableHandle) handle.getTableHandle()).getSnapshotId() : -1L; + IcebergTableHandle icebergHandle = handle.getTableHandle() instanceof IcebergTableHandle + ? (IcebergTableHandle) handle.getTableHandle() : null; + long readSnapshotId = icebergHandle == null ? -1L : icebergHandle.getSnapshotId(); + boolean readSnapshotPinned = icebergHandle != null && icebergHandle.hasSnapshotPin(); // Branch-targeted INSERT (INSERT INTO tbl@branch): the branch is threaded from the generic insert // command context onto the write handle; beginWrite validates it against the table refs and points // the commit at the branch. Empty for a default-ref write. return new IcebergWriteContext(op, handle.isOverwrite(), handle.getStaticPartitionSpec(), - handle.getBranchName(), readSnapshotId); + handle.getBranchName(), readSnapshotId, readSnapshotPinned); } private TIcebergTableSink buildSink(Table table, IcebergTableHandle tableHandle, diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactory.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactory.java index 6a901feb349a74..52e531ab338b25 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactory.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactory.java @@ -36,7 +36,7 @@ * {@code Optional} / nereids {@code Expression}. * *

    T03 scaffolding. The {@code createAction} switch carries only the faithful default rejection; - * the 9 procedure cases (their bodies) are ported in T04 ({@code rewrite_data_files} in T05/T06). The + * the procedure cases (their bodies) are ported in T04 ({@code rewrite_data_files} in T05/T06). The * {@link #getSupportedActions()} registry — exported to {@code getSupportedProcedures()} and embedded in * the rejection message — is complete and final. */ @@ -52,6 +52,7 @@ public class IcebergExecuteActionFactory { public static final String REWRITE_DATA_FILES = "rewrite_data_files"; public static final String PUBLISH_CHANGES = "publish_changes"; public static final String REWRITE_MANIFESTS = "rewrite_manifests"; + public static final String REMOVE_ORPHAN_FILES = "remove_orphan_files"; /** * Create an iceberg procedure body for {@code actionType}. @@ -83,6 +84,8 @@ public static BaseIcebergAction createAction(String actionType, Map properties, List partitionNames, + ConnectorPredicate whereCondition) { + super("remove_orphan_files", properties, partitionNames, whereCondition); + } + + @Override + protected void registerIcebergArguments() { + namedArguments.registerRequiredArgument(OLDER_THAN, "Creation time cutoff in milliseconds", + ArgumentParsers.nonNegativeLong(OLDER_THAN)); + namedArguments.registerOptionalArgument(LOCATION, "Prefix within the table location", + null, ArgumentParsers.nonEmptyString(LOCATION)); + namedArguments.registerOptionalArgument(DRY_RUN, "Only count orphan files", true, + ArgumentParsers.booleanValue(DRY_RUN)); + } + + @Override + protected void validateIcebergAction() { + validateNoPartitions(); + validateNoWhereCondition(); + String location = namedArguments.getString(LOCATION); + if (location != null) { + try { + normalizeLocation(location); + } catch (IllegalArgumentException e) { + throw new DorisConnectorException("Invalid location URI: " + location, e); + } + } + } + + @Override + protected List executeAction(Table table, ConnectorSession session) { + if (!(table.io() instanceof SupportsPrefixOperations)) { + throw new DorisConnectorException("remove_orphan_files requires FileIO prefix listing support"); + } + if (!PropertyUtil.propertyAsBoolean(table.properties(), TableProperties.GC_ENABLED, + TableProperties.GC_ENABLED_DEFAULT)) { + // A GC-disabled table may share files with another table, so no destructive scan is safe. + throw new DorisConnectorException("Cannot remove orphan files: Iceberg GC is disabled"); + } + List scanLocations = resolveScanLocations(table); + + try { + ReachableIndex reachable = new ReachableIndex(collectReachableFiles(table)); + long orphanCount = 0; + long deletedCount = 0; + long olderThan = namedArguments.getLong(OLDER_THAN); + // The SQL procedure needs a retention fence because concurrent uploads are not reachable until commit. + if (olderThan > System.currentTimeMillis() - MIN_RETENTION_MS) { + throw new DorisConnectorException( + "older_than must retain at least 24 hours of files"); + } + boolean dryRun = namedArguments.getBoolean(DRY_RUN); + Set visitedFiles = new HashSet<>(); + for (String scanLocation : scanLocations) { + // Object stores use raw prefix matching, so the separator excludes sibling prefixes. + String listingPrefix = scanLocation.endsWith("/") ? scanLocation : scanLocation + "/"; + for (FileInfo file : ((SupportsPrefixOperations) table.io()).listPrefix(listingPrefix)) { + if (visitedFiles.add(file.location()) && file.createdAtMillis() < olderThan + && !isReachable(file.location(), reachable)) { + orphanCount++; + if (!dryRun) { + table.io().deleteFile(file.location()); + deletedCount++; + } + } + } + } + return Lists.newArrayList(String.valueOf(orphanCount), String.valueOf(deletedCount)); + } catch (Exception e) { + throw new DorisConnectorException("Failed to remove orphan files: " + e.getMessage(), e); + } + } + + private List resolveScanLocations(Table table) { + String tableRoot = normalizeLocation(table.location()); + String dataRoot = normalizeLocation(resolveDataLocation(table, tableRoot)); + Set ownedRoots = new LinkedHashSet<>(); + ownedRoots.add(tableRoot); + ownedRoots.add(dataRoot); + + String requested = namedArguments.getString(LOCATION); + if (requested != null) { + String normalized = normalizeLocation(requested); + boolean owned = ownedRoots.stream().anyMatch(root -> isWithin(normalized, root)); + if (!owned) { + throw new DorisConnectorException( + "location must be within an Iceberg table-owned metadata or data location"); + } + return Lists.newArrayList(normalized); + } + + List roots = new ArrayList<>(); + for (String candidate : ownedRoots) { + // Avoid listing a nested default data directory twice when the table root already covers it. + if (ownedRoots.stream().noneMatch(other -> !other.equals(candidate) && isWithin(candidate, other))) { + roots.add(candidate); + } + } + return roots; + } + + private String resolveDataLocation(Table table, String tableRoot) { + Map properties = table.properties(); + String dataLocation = nonEmpty(properties.get(TableProperties.WRITE_DATA_LOCATION)); + if (dataLocation == null && Boolean.parseBoolean(properties.get(TableProperties.OBJECT_STORE_ENABLED))) { + dataLocation = nonEmpty(properties.get(TableProperties.OBJECT_STORE_PATH)); + } + if (dataLocation == null) { + dataLocation = nonEmpty(properties.get(TableProperties.WRITE_FOLDER_STORAGE_LOCATION)); + } + return dataLocation == null ? tableRoot + "/data" : dataLocation; + } + + private String nonEmpty(String location) { + return location == null || location.isEmpty() ? null : location; + } + + private boolean isWithin(String location, String root) { + FileIdentity child = FileIdentity.of(location); + FileIdentity parent = FileIdentity.of(root); + String pathPrefix = parent.path.endsWith("/") ? parent.path : parent.path + "/"; + return child.scheme.equals(parent.scheme) && child.authority.equals(parent.authority) + && (child.path.equals(parent.path) || child.path.startsWith(pathPrefix)); + } + + private Set collectReachableFiles(Table table) throws IOException { + Set reachable = new HashSet<>(ReachableFileUtil.metadataFileLocations(table, true)); + // Hadoop tables consult this live pointer even though it is not part of the metadata log. + reachable.add(ReachableFileUtil.versionHintLocation(table)); + Set scannedDataManifests = new HashSet<>(); + Set scannedDeleteManifests = new HashSet<>(); + reachable.addAll(ReachableFileUtil.manifestListLocations(table)); + reachable.addAll(ReachableFileUtil.statisticsFilesLocations(table)); + for (Snapshot snapshot : table.snapshots()) { + for (ManifestFile manifest : snapshot.allManifests(table.io())) { + reachable.add(manifest.path()); + if (manifest.content() == ManifestContent.DATA) { + // Snapshots inherit manifests, so read each path once to keep work linear. + if (scannedDataManifests.add(manifest.path())) { + try (ManifestReader dataFiles = + ManifestFiles.read(manifest, table.io(), table.specs())) { + dataFiles.forEach(dataFile -> reachable.add(dataFile.location())); + } + } + } else if (scannedDeleteManifests.add(manifest.path())) { + // A retained delete file may not apply to any current data task, so read it directly. + try (ManifestReader deletes = + ManifestFiles.readDeleteManifest(manifest, table.io(), table.specs())) { + deletes.forEach(delete -> reachable.add(delete.location())); + } + } + } + } + return reachable; + } + + private static boolean isReachable(String candidate, ReachableIndex reachable) { + FileIdentity candidateIdentity = FileIdentity.of(candidate); + if (reachable.identities.contains(candidateIdentity)) { + return true; + } + if (reachable.paths.contains(candidateIdentity.path)) { + // A path collision across unknown providers/authorities cannot be classified safely. + throw new DorisConnectorException( + "Cannot determine whether listed and reachable file locations are equivalent"); + } + return false; + } + + static boolean sameFileIdentity(String first, String second) { + return FileIdentity.of(first).equals(FileIdentity.of(second)); + } + + static void verifyNoPrefixMismatch(String candidate, Set reachable) { + FileIdentity candidateIdentity = FileIdentity.of(candidate); + for (String retained : reachable) { + FileIdentity retainedIdentity = FileIdentity.of(retained); + // Matching paths with different providers/authorities are ambiguous; deletion must fail closed. + if (candidateIdentity.path.equals(retainedIdentity.path) + && !candidateIdentity.equals(retainedIdentity)) { + throw new DorisConnectorException( + "Cannot determine whether listed and reachable file locations are equivalent"); + } + } + } + + private static final class FileIdentity { + private final String scheme; + private final String authority; + private final String path; + + private FileIdentity(String scheme, String authority, String path) { + this.scheme = scheme; + this.authority = authority; + this.path = path; + } + + private static FileIdentity of(String location) { + URI uri = URI.create(location).normalize(); + String scheme = uri.getScheme(); + scheme = scheme == null ? "" : scheme.toLowerCase(Locale.ROOT); + if (scheme.equals("s3a") || scheme.equals("s3n")) { + scheme = "s3"; + } + String authority = uri.getAuthority(); + authority = authority == null ? "" : authority.toLowerCase(Locale.ROOT); + String path = uri.getPath(); + return new FileIdentity(scheme, authority, path == null ? "" : path); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof FileIdentity)) { + return false; + } + FileIdentity that = (FileIdentity) other; + return scheme.equals(that.scheme) && authority.equals(that.authority) + && path.equals(that.path); + } + + @Override + public int hashCode() { + return Objects.hash(scheme, authority, path); + } + } + + private static final class ReachableIndex { + private final Set identities = new HashSet<>(); + private final Set paths = new HashSet<>(); + + private ReachableIndex(Set locations) { + for (String location : locations) { + FileIdentity identity = FileIdentity.of(location); + identities.add(identity); + paths.add(identity.path); + } + } + } + + private String normalizeLocation(String location) { + String normalized = URI.create(location).normalize().toString(); + return normalized.length() > 1 && normalized.endsWith("/") + ? normalized.substring(0, normalized.length() - 1) : normalized; + } + + @Override + protected List getResultSchema() { + return Lists.newArrayList( + new ConnectorColumn("orphan_file_count", ConnectorType.of("BIGINT"), + "Number of old unreachable files", false, null), + new ConnectorColumn("deleted_file_count", ConnectorType.of("BIGINT"), + "Number of files deleted", false, null)); + } +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadataMvccTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadataMvccTest.java index 35748356013e52..283c434905f596 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadataMvccTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorMetadataMvccTest.java @@ -335,16 +335,17 @@ public void applySnapshotThreadsRef() { } @Test - public void applySnapshotLatestPinLeavesHandleUnchanged() { + public void applySnapshotPreservesExplicitEmptyPinWhileScanningLatest() { Fixture f = fixture(); IcebergConnectorMetadata md = metadataFor(f.table, new RecordingIcebergCatalogOps()); ConnectorTableHandle bare = handle(); - // null snapshot and an empty-table (-1, no ref) pin must both read latest (handle unchanged) — a - // useSnapshot(-1) would be a non-existent snapshot. + // Null means no pin, but an empty-table pin must remain observable to the write OCC path even though + // the scan still reads latest (useSnapshot(-1) would be a non-existent snapshot). Assertions.assertSame(bare, md.applySnapshot(null, bare, null)); IcebergTableHandle afterMinusOne = (IcebergTableHandle) md.applySnapshot(null, bare, ConnectorMvccSnapshot.builder().snapshotId(-1L).build()); - Assertions.assertFalse(afterMinusOne.hasSnapshotPin()); + Assertions.assertTrue(afterMinusOne.hasSnapshotPin()); + Assertions.assertEquals(-1L, afterMinusOne.getSnapshotId()); } // --------------------------------------------------------------------- diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTransactionTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTransactionTest.java index 0ed5a711ae065a..73ad91ee82b916 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTransactionTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTransactionTest.java @@ -126,6 +126,11 @@ private static IcebergWriteContext overwriteCtx() { return new IcebergWriteContext(WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.empty()); } + private static IcebergWriteContext overwriteCtxPinned(long readSnapshotId) { + return new IcebergWriteContext(WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.empty(), + readSnapshotId); + } + private static IcebergWriteContext overwriteStaticCtx(Map staticValues) { return new IcebergWriteContext(WriteOperation.OVERWRITE, true, staticValues, Optional.empty()); } @@ -577,6 +582,106 @@ public void overwriteDynamicReplacesPartitions() { Assertions.assertEquals("1", snap.summary().get("added-data-files")); } + @Test + public void overwriteDynamicRejectsConcurrentDataInReplacedPartition() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + PartitionSpec spec = PartitionSpec.builderFor(PART_SCHEMA).identity("region").build(); + Table table = catalog.createTable(id, PART_SCHEMA, spec, props("write.format.default", "parquet")); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/seed.parquet", 1L, "region=us")).commit(); + + IcebergConnectorTransaction txn = txnFor(opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", overwriteCtx()); + Table concurrent = catalog.loadTable(id); + concurrent.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/concurrent.parquet", 1L, "region=us")).commit(); + txn.addCommitData(commitBytes( + dataFileItem("s3://b/db1/t1/region=us/replacement.parquet", 2L, 1024L, + Collections.singletonList("us")))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit, + "dynamic overwrite must not silently replace data committed after its base snapshot"); + } + + @Test + public void overwriteDynamicRejectsDataCommittedBetweenScanAndBegin() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + PartitionSpec spec = PartitionSpec.builderFor(PART_SCHEMA).identity("region").build(); + Table table = catalog.createTable(id, PART_SCHEMA, spec, props("write.format.default", "parquet")); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/seed.parquet", 1L, "region=us")).commit(); + long readSnapshotId = table.currentSnapshot().snapshotId(); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/between-scan-and-begin.parquet", 1L, "region=us")).commit(); + + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", overwriteCtxPinned(readSnapshotId)); + txn.addCommitData(commitBytes( + dataFileItem("s3://b/db1/t1/region=us/replacement.parquet", 2L, 1024L, + Collections.singletonList("us")))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + + @Test + public void overwriteRejectsFirstSnapshotCommittedAfterEmptyRead() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table table = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), + props("write.format.default", "parquet")); + table.newAppend().appendFile(dataFile(table.spec(), + "s3://b/db1/t1/between-scan-and-begin.parquet", 1L)).commit(); + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + + Assertions.assertThrows(DorisConnectorException.class, + () -> txn.beginWrite(SESSION, "db1", "t1", overwriteCtxPinned(-1L))); + } + + @Test + public void overwriteRejectsFirstSnapshotCommittedAfterBeginFromEmptyRead() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table table = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), + props("write.format.default", "parquet")); + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", overwriteCtxPinned(-1L)); + + table.newAppend().appendFile(dataFile(table.spec(), + "s3://b/db1/t1/after-begin.parquet", 1L)).commit(); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + + @Test + public void overwriteBranchUsesTheSnapshotReadFromThatBranch() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + PartitionSpec spec = PartitionSpec.builderFor(PART_SCHEMA).identity("region").build(); + Table table = catalog.createTable(id, PART_SCHEMA, spec, props("write.format.default", "parquet")); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/seed.parquet", 1L, "region=us")).commit(); + long readSnapshotId = table.currentSnapshot().snapshotId(); + table.manageSnapshots().createBranch("b1", readSnapshotId).commit(); + table.newAppend().toBranch("b1").appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/between-scan-and-begin.parquet", 1L, "region=us")).commit(); + + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", + new IcebergWriteContext(WriteOperation.OVERWRITE, true, Collections.emptyMap(), + Optional.of("b1"), readSnapshotId)); + txn.addCommitData(commitBytes( + dataFileItem("s3://b/db1/t1/region=us/replacement.parquet", 2L, 1024L, + Collections.singletonList("us")))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + @Test public void overwriteEmptyUnpartitionedClearsTable() { InMemoryCatalog catalog = freshCatalog(); @@ -597,6 +702,24 @@ public void overwriteEmptyUnpartitionedClearsTable() { "the existing data file must be removed (table cleared)"); } + @Test + public void overwriteEmptyUnpartitionedRejectsDataCommittedBetweenScanAndBegin() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table table = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), + props("write.format.default", "parquet")); + table.newAppend().appendFile(dataFile(table.spec(), "s3://b/db1/t1/seed.parquet", 1L)).commit(); + long readSnapshotId = table.currentSnapshot().snapshotId(); + table.newAppend().appendFile(dataFile(table.spec(), + "s3://b/db1/t1/between-scan-and-begin.parquet", 1L)).commit(); + + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", overwriteCtxPinned(readSnapshotId)); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + @Test public void overwriteStaticPartitionUsesRowFilter() { InMemoryCatalog catalog = freshCatalog(); @@ -616,6 +739,52 @@ public void overwriteStaticPartitionUsesRowFilter() { Assertions.assertEquals("1", snap.summary().get("added-data-files")); } + @Test + public void overwriteStaticRejectsConcurrentDataInTargetPartition() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + PartitionSpec spec = PartitionSpec.builderFor(PART_SCHEMA).identity("region").build(); + Table table = catalog.createTable(id, PART_SCHEMA, spec, props("write.format.default", "parquet")); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/seed.parquet", 1L, "region=us")).commit(); + + IcebergConnectorTransaction txn = txnFor(opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", overwriteStaticCtx(Collections.singletonMap("region", "us"))); + Table concurrent = catalog.loadTable(id); + concurrent.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/concurrent.parquet", 1L, "region=us")).commit(); + txn.addCommitData(commitBytes( + dataFileItem("s3://b/db1/t1/region=us/replacement.parquet", 2L, 1024L, + Collections.singletonList("us")))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit, + "static overwrite must reject concurrent data matching its target partition filter"); + } + + @Test + public void overwriteStaticRejectsDataCommittedBetweenScanAndBegin() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + PartitionSpec spec = PartitionSpec.builderFor(PART_SCHEMA).identity("region").build(); + Table table = catalog.createTable(id, PART_SCHEMA, spec, props("write.format.default", "parquet")); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/seed.parquet", 1L, "region=us")).commit(); + long readSnapshotId = table.currentSnapshot().snapshotId(); + table.newAppend().appendFile(partitionedDataFile(spec, + "s3://b/db1/t1/region=us/between-scan-and-begin.parquet", 1L, "region=us")).commit(); + + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", + new IcebergWriteContext(WriteOperation.OVERWRITE, true, + Collections.singletonMap("region", "us"), Optional.empty(), readSnapshotId)); + txn.addCommitData(commitBytes( + dataFileItem("s3://b/db1/t1/region=us/replacement.parquet", 2L, 1024L, + Collections.singletonList("us")))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + @Test public void deleteWritesRowDeltaDeleteFiles() { InMemoryCatalog catalog = freshCatalog(); diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergProcedureOpsTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergProcedureOpsTest.java index d2fa6cf9221b77..32958429302064 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergProcedureOpsTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergProcedureOpsTest.java @@ -86,7 +86,8 @@ public void getSupportedProceduresExportsFactoryNamesInLegacyOrder() { "expire_snapshots", "rewrite_data_files", "publish_changes", - "rewrite_manifests"), + "rewrite_manifests", + "remove_orphan_files"), newOps().getSupportedProcedures()); } @@ -120,7 +121,8 @@ public void executeRejectsUnknownProcedureWithLegacyMessage() { Assertions.assertEquals( "Unsupported Iceberg procedure: no_such_proc. Supported procedures: rollback_to_snapshot, " + "rollback_to_timestamp, set_current_snapshot, cherrypick_snapshot, fast_forward, " - + "expire_snapshots, rewrite_data_files, publish_changes, rewrite_manifests", + + "expire_snapshots, rewrite_data_files, publish_changes, rewrite_manifests, " + + "remove_orphan_files", e.getMessage()); } diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergTableHandleTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergTableHandleTest.java index f5e928afaaa50c..f01683a8831258 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergTableHandleTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergTableHandleTest.java @@ -57,6 +57,16 @@ public void withSnapshotPinsByIdAndCarriesSchemaId() { Assertions.assertEquals("t1", pinned.getTableName()); } + @Test + public void explicitEmptySnapshotIsDistinctFromNoPin() { + IcebergTableHandle bare = new IcebergTableHandle("db1", "t1"); + IcebergTableHandle empty = bare.withSnapshot(-1L, null, -1L); + + Assertions.assertTrue(empty.hasSnapshotPin()); + Assertions.assertFalse(empty.hasSnapshotSelection()); + Assertions.assertNotEquals(bare, empty); + } + @Test public void withSnapshotPinsByRef() { IcebergTableHandle pinned = new IcebergTableHandle("db1", "t1").withSnapshot(7L, "b1", 2L); diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactoryTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactoryTest.java index 1e0da5abff5955..d5fdf67695e187 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactoryTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergExecuteActionFactoryTest.java @@ -19,24 +19,23 @@ import org.apache.doris.connector.api.DorisConnectorException; +import com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import java.util.Collections; /** - * Pins the connector port of legacy {@code IcebergExecuteActionFactory} (the name registry + dispatch). + * Pins the Iceberg procedure name registry and dispatch. * *

    WHY this matters: the supported-name list is exported to {@code getSupportedProcedures()} and - * embedded in the unknown-procedure error, so its membership and order must match legacy byte-for-byte - * (T08 byte-parity). The {@code table} parameter is dropped (it was always dead in legacy). The 9 switch - * cases are added in T04 (the procedure bodies); T03 fixes the registry + the faithful unknown-procedure - * rejection.

    + * embedded in the unknown-procedure error, so membership, ordering, and executable action mappings must stay + * synchronized.

    */ public class IcebergExecuteActionFactoryTest { @Test - public void getSupportedActionsReturnsNineNamesInLegacyOrder() { + public void getSupportedActionsIncludesOrphanCleanup() { Assertions.assertArrayEquals( new String[] { "rollback_to_snapshot", @@ -48,6 +47,7 @@ public void getSupportedActionsReturnsNineNamesInLegacyOrder() { "rewrite_data_files", "publish_changes", "rewrite_manifests", + "remove_orphan_files", }, IcebergExecuteActionFactory.getSupportedActions()); } @@ -60,15 +60,32 @@ public void createActionRejectsUnknownProcedureWithLegacyMessage() { Assertions.assertEquals( "Unsupported Iceberg procedure: no_such_proc. Supported procedures: rollback_to_snapshot, " + "rollback_to_timestamp, set_current_snapshot, cherrypick_snapshot, fast_forward, " - + "expire_snapshots, rewrite_data_files, publish_changes, rewrite_manifests", + + "expire_snapshots, rewrite_data_files, publish_changes, rewrite_manifests, " + + "remove_orphan_files", e.getMessage()); } + @Test + public void createRemoveOrphanFilesAction() { + BaseIcebergAction action = IcebergExecuteActionFactory.createAction( + "remove_orphan_files", Collections.singletonMap("older_than", "1"), + Collections.emptyList(), null); + Assertions.assertInstanceOf(IcebergRemoveOrphanFilesAction.class, action); + } + + @Test + public void removeOrphanFilesRejectsInvalidLocationUri() { + BaseIcebergAction action = IcebergExecuteActionFactory.createAction( + "remove_orphan_files", ImmutableMap.of("older_than", "1", "location", "://"), + Collections.emptyList(), null); + Assertions.assertThrows(DorisConnectorException.class, action::validate); + } + /** * CANARY for the dormant {@code rewrite_data_files} gap: it is advertised in {@link - * IcebergExecuteActionFactory#getSupportedActions()} (9 names) but has NO {@code createAction} switch - * case yet (8 cases), so it falls through to the faithful unknown-procedure rejection. This pins that - * dormant state and goes RED exactly when the T05/T06 body is wired in. + * IcebergExecuteActionFactory#getSupportedActions()} but has NO {@code createAction} switch + * case because it is dispatched through the distributed rewrite planner, so it falls through to the + * unknown-procedure rejection in this single-call factory. */ @Test public void rewriteDataFilesIsAdvertisedButNotYetExecutable() { diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergRemoveOrphanFilesActionTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergRemoveOrphanFilesActionTest.java new file mode 100644 index 00000000000000..05d3dc1ae4e048 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergRemoveOrphanFilesActionTest.java @@ -0,0 +1,308 @@ +// 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. + +package org.apache.doris.connector.iceberg.action; + +import org.apache.doris.connector.api.DorisConnectorException; +import org.apache.doris.connector.api.procedure.ConnectorProcedureResult; + +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.BaseTable; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.HasTableOperations; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.ReachableFileUtil; +import org.apache.iceberg.StaticTableOperations; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.FileInfo; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.io.SupportsPrefixOperations; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.attribute.FileTime; +import java.time.Duration; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +public class IcebergRemoveOrphanFilesActionTest { + private static final long MIN_RETENTION_MS = Duration.ofHours(24).toMillis(); + + @Test + public void gcDisabledPreventsDeletion(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("table"), + Collections.singletonMap(TableProperties.GC_ENABLED, "false")); + Path orphan = createOldFile(temp.resolve("table/data/orphan.parquet")); + IcebergRemoveOrphanFilesAction action = action(System.currentTimeMillis() - MIN_RETENTION_MS, false); + action.validate(); + + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(table, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(orphan)); + } + + @Test + public void recentCutoffCannotRaceAnUncommittedWriter(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("table"), Collections.emptyMap()); + Path uncommitted = createOldFile(temp.resolve("table/data/uncommitted.parquet")); + IcebergRemoveOrphanFilesAction action = action(System.currentTimeMillis(), false); + action.validate(); + + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(table, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(uncommitted)); + } + + @Test + public void keepsVersionHintWhileDeletingAnOldOrphan(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("table"), Collections.emptyMap()); + Path orphan = createOldFile(temp.resolve("table/data/orphan.parquet")); + Path versionHint = Path.of(java.net.URI.create(ReachableFileUtil.versionHintLocation(table))); + Files.setLastModifiedTime(versionHint, FileTime.fromMillis(1)); + IcebergRemoveOrphanFilesAction action = action(System.currentTimeMillis() - MIN_RETENTION_MS, false); + action.validate(); + + ConnectorProcedureResult result = action.execute(table, ActionTestTables.session("UTC")); + + Assertions.assertEquals("1", result.getRows().get(0).get(0)); + Assertions.assertEquals("1", result.getRows().get(0).get(1)); + Assertions.assertFalse(Files.exists(orphan)); + Assertions.assertTrue(Files.exists(versionHint)); + } + + @Test + public void treatsS3SchemeAliasesAsTheSameFile() { + Assertions.assertTrue(IcebergRemoveOrphanFilesAction.sameFileIdentity( + "s3://bucket/path/data.parquet", "s3a://bucket/path/data.parquet")); + Assertions.assertTrue(IcebergRemoveOrphanFilesAction.sameFileIdentity( + "s3n://bucket/path/data.parquet", "s3://BUCKET/path/data.parquet")); + } + + @Test + public void rejectsUnresolvedPrefixMismatches() { + Assertions.assertThrows(DorisConnectorException.class, + () -> IcebergRemoveOrphanFilesAction.verifyNoPrefixMismatch( + "s3://first/path/data.parquet", + Collections.singleton("s3://second/path/data.parquet"))); + } + + @Test + public void readsEachSharedDataManifestOnlyOnce(@TempDir Path temp) throws Exception { + Map properties = new HashMap<>(); + properties.put(TableProperties.MANIFEST_MERGE_ENABLED, "false"); + Table table = createTable(temp.resolve("table"), properties); + appendDataFile(table, createOldFile(temp.resolve("table/data/first.parquet"))); + appendDataFile(table, createOldFile(temp.resolve("table/data/second.parquet"))); + + Set dataManifestPaths = new HashSet<>(); + int[] manifestReferenceCount = {0}; + table.snapshots().forEach(snapshot -> snapshot.dataManifests(table.io()) + .forEach(manifest -> { + manifestReferenceCount[0]++; + dataManifestPaths.add(manifest.path()); + })); + Assertions.assertTrue(manifestReferenceCount[0] > dataManifestPaths.size()); + RecordingFileIO recordingFileIO = new RecordingFileIO(table.io(), dataManifestPaths); + Table recordingTable = new BaseTable( + new StaticTableOperations(((HasTableOperations) table).operations().current(), recordingFileIO), + table.name()); + IcebergRemoveOrphanFilesAction action = action( + System.currentTimeMillis() - MIN_RETENTION_MS, true); + action.validate(); + + action.execute(recordingTable, ActionTestTables.session("UTC")); + + Assertions.assertEquals(dataManifestPaths.size(), recordingFileIO.manifestOpenCount()); + dataManifestPaths.forEach(path -> Assertions.assertEquals(1, + recordingFileIO.openCounts.getOrDefault(path, 0), path)); + } + + @Test + public void scansConfiguredDataRootOutsideTableLocationByDefault(@TempDir Path temp) throws Exception { + Path dataRoot = temp.resolve("owned-data"); + Table table = createTable(temp.resolve("metadata"), + Collections.singletonMap(TableProperties.WRITE_DATA_LOCATION, + dataRoot.toUri().toString())); + Path orphan = createOldFile(dataRoot.resolve("orphan.parquet")); + IcebergRemoveOrphanFilesAction action = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + action.validate(); + + ConnectorProcedureResult result = action.execute(table, ActionTestTables.session("UTC")); + + Assertions.assertEquals("1", result.getRows().get(0).get(0)); + Assertions.assertEquals("1", result.getRows().get(0).get(1)); + Assertions.assertFalse(Files.exists(orphan)); + } + + @Test + public void allowsExplicitConfiguredDataRootButRejectsArbitraryRoot(@TempDir Path temp) throws Exception { + Path dataRoot = temp.resolve("owned-data"); + Table table = createTable(temp.resolve("metadata"), + Collections.singletonMap(TableProperties.WRITE_DATA_LOCATION, + dataRoot.toUri().toString())); + Path orphan = createOldFile(dataRoot.resolve("orphan.parquet")); + + IcebergRemoveOrphanFilesAction configured = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, dataRoot.toUri().toString()); + configured.validate(); + configured.execute(table, ActionTestTables.session("UTC")); + Assertions.assertFalse(Files.exists(orphan)); + + IcebergRemoveOrphanFilesAction arbitrary = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + temp.resolve("unowned").toUri().toString()); + arbitrary.validate(); + Assertions.assertThrows(DorisConnectorException.class, + () -> arbitrary.execute(table, ActionTestTables.session("UTC"))); + } + + @Test + public void scansObjectStoreAndFolderStorageFallbackRoots(@TempDir Path temp) throws Exception { + Path objectRoot = temp.resolve("object-data"); + Map objectProperties = new HashMap<>(); + objectProperties.put(TableProperties.OBJECT_STORE_ENABLED, "true"); + objectProperties.put(TableProperties.OBJECT_STORE_PATH, objectRoot.toUri().toString()); + Table objectTable = createTable(temp.resolve("object-metadata"), objectProperties); + Path objectOrphan = createOldFile(objectRoot.resolve("orphan.parquet")); + + Path folderRoot = temp.resolve("folder-data"); + Table folderTable = createTable(temp.resolve("folder-metadata"), + Collections.singletonMap(TableProperties.WRITE_FOLDER_STORAGE_LOCATION, + folderRoot.toUri().toString())); + Path folderOrphan = createOldFile(folderRoot.resolve("orphan.parquet")); + + IcebergRemoveOrphanFilesAction objectAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + objectAction.validate(); + objectAction.execute(objectTable, ActionTestTables.session("UTC")); + IcebergRemoveOrphanFilesAction folderAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + folderAction.validate(); + folderAction.execute(folderTable, ActionTestTables.session("UTC")); + + Assertions.assertFalse(Files.exists(objectOrphan)); + Assertions.assertFalse(Files.exists(folderOrphan)); + } + + private static IcebergRemoveOrphanFilesAction action(long olderThan, boolean dryRun) { + return action(olderThan, dryRun, null); + } + + private static IcebergRemoveOrphanFilesAction action(long olderThan, boolean dryRun, String location) { + Map properties = new HashMap<>(); + properties.put(IcebergRemoveOrphanFilesAction.OLDER_THAN, String.valueOf(olderThan)); + properties.put(IcebergRemoveOrphanFilesAction.DRY_RUN, String.valueOf(dryRun)); + if (location != null) { + properties.put(IcebergRemoveOrphanFilesAction.LOCATION, location); + } + return new IcebergRemoveOrphanFilesAction(properties, Collections.emptyList(), null); + } + + private static Table createTable(Path location, Map properties) { + HadoopTables tables = new HadoopTables(new Configuration()); + return tables.create(ActionTestTables.SCHEMA, PartitionSpec.unpartitioned(), properties, + location.toUri().toString()); + } + + private static Path createOldFile(Path path) throws Exception { + Files.createDirectories(path.getParent()); + Files.write(path, new byte[] {1}); + Files.setLastModifiedTime(path, FileTime.fromMillis(1)); + return path; + } + + private static void appendDataFile(Table table, Path path) { + DataFile dataFile = DataFiles.builder(table.spec()) + .withPath(path.toUri().toString()) + .withFileSizeInBytes(1) + .withRecordCount(1) + .build(); + table.newFastAppend().appendFile(dataFile).commit(); + } + + private static final class RecordingFileIO implements SupportsPrefixOperations { + private final FileIO delegate; + private final SupportsPrefixOperations prefixDelegate; + private final Set manifestPaths; + private final Map openCounts = new HashMap<>(); + + private RecordingFileIO(FileIO delegate, Set manifestPaths) { + this.delegate = delegate; + this.prefixDelegate = (SupportsPrefixOperations) delegate; + this.manifestPaths = manifestPaths; + } + + private void record(String path) { + if (manifestPaths.contains(path)) { + openCounts.merge(path, 1, Integer::sum); + } + } + + private int manifestOpenCount() { + return openCounts.values().stream().mapToInt(Integer::intValue).sum(); + } + + @Override + public InputFile newInputFile(String path) { + record(path); + return delegate.newInputFile(path); + } + + @Override + public InputFile newInputFile(String path, long length) { + record(path); + return delegate.newInputFile(path, length); + } + + @Override + public OutputFile newOutputFile(String path) { + return delegate.newOutputFile(path); + } + + @Override + public void deleteFile(String path) { + delegate.deleteFile(path); + } + + @Override + public Map properties() { + return delegate.properties(); + } + + @Override + public Iterable listPrefix(String prefix) { + return prefixDelegate.listPrefix(prefix); + } + + @Override + public void deletePrefix(String prefix) { + prefixDelegate.deletePrefix(prefix); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/AbstractInsertExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/AbstractInsertExecutor.java index bdcb6e29448e1e..381b38e9a8bd52 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/AbstractInsertExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/AbstractInsertExecutor.java @@ -140,6 +140,10 @@ public void unregisterListener(InsertExecutorListener listener) { listeners.remove(listener); } + protected void handleAfterCompleteFailure(Exception e) throws Exception { + throw e; + } + public Coordinator getCoordinator() { return coordinator; } @@ -259,8 +263,9 @@ private void checkStrictModeAndFilterRatio() throws Exception { * execute insert txn for insert into select command. */ public void executeSingleInsert(StmtExecutor executor) throws Exception { - beforeExec(); try { + // Pre-execution work may register external resources, so it must share the transaction cleanup scope. + beforeExec(); executor.updateProfile(false); execImpl(executor); checkStrictModeAndFilterRatio(); @@ -269,7 +274,11 @@ public void executeSingleInsert(StmtExecutor executor) throws Exception { } onComplete(); for (InsertExecutorListener listener : listeners) { - listener.afterComplete(this, executor, jobId); + try { + listener.afterComplete(this, executor, jobId); + } catch (Exception e) { + handleAfterCompleteFailure(e); + } } } catch (Throwable t) { onFail(t); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BaseExternalTableInsertExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BaseExternalTableInsertExecutor.java index 797c865b7dbf8c..217c524a499f06 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BaseExternalTableInsertExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/BaseExternalTableInsertExecutor.java @@ -126,8 +126,13 @@ protected void onComplete() throws UserException { txnStatus = TransactionStatus.COMMITTED; long t2 = System.currentTimeMillis(); - // Handle post-commit operations (e.g., cache refresh) - doAfterCommit(); + try { + doAfterCommit(); + } catch (Exception e) { + // Cache refresh cannot undo a durable remote commit, so it must not make clients retry the write. + LOG.warn("Post-commit refresh failed for table {}. Data was committed successfully.", + table.getName(), e); + } long t3 = System.currentTimeMillis(); LOG.info("Transaction commit breakdown: doBeforeCommit={}ms, commit={}ms, doAfterCommit={}ms, total={}ms", t1 - t0, t2 - t1, t3 - t2, t3 - t0); @@ -149,6 +154,16 @@ protected void doAfterCommit() throws DdlException { true); } + @Override + protected void handleAfterCompleteFailure(Exception e) throws Exception { + if (txnStatus != TransactionStatus.COMMITTED) { + super.handleAfterCompleteFailure(e); + return; + } + // A post-commit listener cannot undo remote data, so failing the statement would invite duplicate retries. + LOG.warn("Post-commit listener failed for table {}. Data was committed successfully.", table.getName(), e); + } + @Override protected void finalizeSink(PlanFragment fragment, DataSink sink, PhysicalSink physicalSink) { try { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/OlapInsertExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/OlapInsertExecutorTest.java index 6c8537c5d62bba..96fb43ffbae48e 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/OlapInsertExecutorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/OlapInsertExecutorTest.java @@ -181,6 +181,33 @@ void testOnFailAbortsUncommittedTransaction() throws Exception { } } + @Test + void testBeforeExecFailureUsesTheNormalAbortAndCleanupPath() throws Exception { + ConnectContext ctx = createExecutorContext(); + Coordinator coordinator = createCoordinator(); + GlobalTransactionMgrIface txnMgr = Mockito.mock(GlobalTransactionMgrIface.class); + TransactionState txnState = Mockito.mock(TransactionState.class); + LoadManager loadManager = Mockito.mock(LoadManager.class); + Env currentEnv = createCurrentEnv(loadManager); + StmtExecutor stmtExecutor = createStmtExecutor(); + + try (MockedStatic envFactoryMock = Mockito.mockStatic(EnvFactory.class); + MockedStatic envMock = Mockito.mockStatic(Env.class)) { + prepareFactoryMocks(envFactoryMock, envMock, coordinator, txnMgr, txnState, currentEnv); + ctx.setEnv(currentEnv); + + OlapInsertExecutor executor = createExecutorWithBeforeExecFailure(ctx); + executor.txnId = 10004L; + + Assertions.assertDoesNotThrow(() -> executor.executeSingleInsert(stmtExecutor)); + Assertions.assertEquals(MysqlStateType.ERR, ctx.getState().getStateType()); + Assertions.assertTrue(ctx.getState().getErrorMessage().contains("beforeExec failure")); + Mockito.verify(txnMgr).abortTransaction(1L, 10004L, "beforeExec failure"); + Mockito.verify(coordinator).close(); + Mockito.verify(stmtExecutor).updateProfile(true); + } + } + // Build a fresh context per case so insertResult and QueryState do not leak between tests. private ConnectContext createExecutorContext() { ConnectContext ctx = new ConnectContext(); @@ -256,6 +283,25 @@ private OlapInsertExecutor createExecutor(ConnectContext ctx) { Optional.empty(), false, 0L); } + private OlapInsertExecutor createExecutorWithBeforeExecFailure(ConnectContext ctx) { + Database database = Mockito.mock(Database.class); + Mockito.when(database.getFullName()).thenReturn("test_db"); + Mockito.when(database.getId()).thenReturn(1L); + + OlapTable table = Mockito.mock(OlapTable.class); + Mockito.when(table.getDatabase()).thenReturn(database); + Mockito.when(table.getName()).thenReturn("test_tbl"); + Mockito.when(table.getId()).thenReturn(2L); + + return new OlapInsertExecutor(ctx, table, "label_test", Mockito.mock(NereidsPlanner.class), + Optional.empty(), false, 0L) { + @Override + protected void beforeExec() { + throw new RuntimeException("beforeExec failure"); + } + }; + } + // Redirect coordinator creation and transaction access to mocks so the test stays deterministic. private void prepareFactoryMocks(MockedStatic envFactoryMock, MockedStatic envMock, Coordinator coordinator, GlobalTransactionMgrIface txnMgr, TransactionState txnState, Env currentEnv) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutorTest.java index 9cc5f9561a35f6..f05ea67823fb89 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutorTest.java @@ -33,6 +33,7 @@ import org.apache.doris.planner.PluginDrivenTableSink; import org.apache.doris.thrift.TDataSink; import org.apache.doris.transaction.PluginDrivenTransactionManager; +import org.apache.doris.transaction.TransactionStatus; import org.apache.doris.transaction.TransactionType; import org.junit.jupiter.api.Assertions; @@ -215,6 +216,17 @@ public void doBeforeCommitKeepsCoordinatorRowCountWhenTransactionReportsNoCount( "a -1 (no count) transaction must leave the coordinator-counted loadedRows untouched"); } + @Test + public void postCommitListenerFailureDoesNotTurnACommittedWriteIntoAnError() { + PluginDrivenInsertExecutor exec = newUnconstructedExecutor(); + Deencapsulation.setField(exec, "txnStatus", TransactionStatus.COMMITTED); + Deencapsulation.setField(exec, "table", Mockito.mock(PluginDrivenExternalTable.class)); + + Assertions.assertDoesNotThrow(() -> Deencapsulation.invoke(exec, + "handleAfterCompleteFailure", new RuntimeException("listener failure")), + "a listener cannot roll back or fail a connector write after its remote commit is durable"); + } + /** * Creates a {@link PluginDrivenInsertExecutor} without running its constructor. See the class * javadoc: the constructor builds a Coordinator that needs a live planner/EnvFactory. diff --git a/fe/fe-filesystem/fe-filesystem-azure/src/main/java/org/apache/doris/filesystem/azure/AzureObjStorage.java b/fe/fe-filesystem/fe-filesystem-azure/src/main/java/org/apache/doris/filesystem/azure/AzureObjStorage.java index e97a8284136b05..197b6c64a7851e 100644 --- a/fe/fe-filesystem/fe-filesystem-azure/src/main/java/org/apache/doris/filesystem/azure/AzureObjStorage.java +++ b/fe/fe-filesystem/fe-filesystem-azure/src/main/java/org/apache/doris/filesystem/azure/AzureObjStorage.java @@ -250,7 +250,9 @@ public void completeMultipartUpload(String remotePath, String uploadId, List sorted = new ArrayList<>(parts); sorted.sort((a, b) -> Integer.compare(a.partNumber(), b.partNumber())); for (UploadPartResult part : sorted) { - blockIds.add(toBlockId(part.partNumber())); + // New BEs carry their exact UUID-prefixed ID; the fallback completes uploads from older BEs. + blockIds.add(part.etag() == null || part.etag().isEmpty() + ? toBlockId(part.partNumber()) : part.etag()); } blockBlobClient.commitBlockList(blockIds); } catch (BlobStorageException e) { diff --git a/fe/fe-filesystem/fe-filesystem-azure/src/test/java/org/apache/doris/filesystem/azure/AzureObjStorageExtensionTest.java b/fe/fe-filesystem/fe-filesystem-azure/src/test/java/org/apache/doris/filesystem/azure/AzureObjStorageExtensionTest.java index eca48e0c4bf867..b6a0bc6f973abb 100644 --- a/fe/fe-filesystem/fe-filesystem-azure/src/test/java/org/apache/doris/filesystem/azure/AzureObjStorageExtensionTest.java +++ b/fe/fe-filesystem/fe-filesystem-azure/src/test/java/org/apache/doris/filesystem/azure/AzureObjStorageExtensionTest.java @@ -17,6 +17,7 @@ package org.apache.doris.filesystem.azure; +import org.apache.doris.filesystem.UploadPartResult; import org.apache.doris.filesystem.spi.RemoteObjects; import com.azure.storage.blob.BlobClient; @@ -378,6 +379,48 @@ void deleteObjectsByKeys_attachesPerKeyExceptionsAsSuppressed() throws Exception // F20 — abortMultipartUpload safe-noop / commit-empty behaviour // ------------------------------------------------------------------ + @Test + void completeMultipartUpload_usesExactBlockIdsReportedByBe() throws Exception { + com.azure.storage.blob.specialized.BlockBlobClient blockClient = + Mockito.mock(com.azure.storage.blob.specialized.BlockBlobClient.class); + BlobClient blobClient = Mockito.mock(BlobClient.class); + Mockito.when(blobClient.getBlockBlobClient()).thenReturn(blockClient); + BlobContainerClient containerClient = Mockito.mock(BlobContainerClient.class); + Mockito.when(containerClient.getBlobClient("stage/blob")).thenReturn(blobClient); + BlobServiceClient serviceClient = Mockito.mock(BlobServiceClient.class); + Mockito.when(serviceClient.getBlobContainerClient("mycontainer")).thenReturn(containerClient); + TestableAzureObjStorage storage = new TestableAzureObjStorage(buildBasicProps(), serviceClient); + + String firstBlockId = "YmUtZ2VuZXJhdGVkLXVwbG9hZC1pZDowMDAwMDAwMDAx"; + String secondBlockId = "YmUtZ2VuZXJhdGVkLXVwbG9hZC1pZDowMDAwMDAwMDAy"; + storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "be-generated-upload-id", + Arrays.asList(new UploadPartResult(2, secondBlockId), + new UploadPartResult(1, firstBlockId))); + + Mockito.verify(blockClient).commitBlockList(Arrays.asList(firstBlockId, secondBlockId)); + } + + @Test + void completeMultipartUpload_fallsBackForOlderBeWithoutBlockIds() throws Exception { + com.azure.storage.blob.specialized.BlockBlobClient blockClient = + Mockito.mock(com.azure.storage.blob.specialized.BlockBlobClient.class); + BlobClient blobClient = Mockito.mock(BlobClient.class); + Mockito.when(blobClient.getBlockBlobClient()).thenReturn(blockClient); + BlobContainerClient containerClient = Mockito.mock(BlobContainerClient.class); + Mockito.when(containerClient.getBlobClient("stage/blob")).thenReturn(blobClient); + BlobServiceClient serviceClient = Mockito.mock(BlobServiceClient.class); + Mockito.when(serviceClient.getBlobContainerClient("mycontainer")).thenReturn(containerClient); + TestableAzureObjStorage storage = new TestableAzureObjStorage(buildBasicProps(), serviceClient); + + storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "legacy-upload-id", Collections.singletonList(new UploadPartResult(1, ""))); + + Mockito.verify(blockClient).commitBlockList(Collections.singletonList("AQAAAA==")); + } + @Test void abortMultipartUpload_safeNoopWhenCommittedBlobExists() throws Exception { com.azure.storage.blob.models.BlobProperties props =