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..3026922219e2aa --- /dev/null +++ b/be/src/exec/operator/iceberg_sorter_reserve_memory.h @@ -0,0 +1,74 @@ +// 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 { + +class Block; + +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; +} + +inline size_t iceberg_reserve_size( + const std::vector& per_partition_reservations, + size_t incoming_block_bytes) { + size_t sorter_reserve = bounded_iceberg_reserve_size(per_partition_reservations); + // The incoming block creates cold partition writers before they can appear in the published snapshot. + return std::min(std::numeric_limits::max() - sorter_reserve, incoming_block_bytes) + + sorter_reserve; +} + +size_t iceberg_cold_writer_reserve_size(const Block& block, size_t writer_workspace_bytes); + +inline size_t iceberg_spill_merge_workspace(size_t spill_file_count, size_t spill_buffer_bytes, + size_t merge_limit_bytes) { + if (spill_file_count == 0 || spill_buffer_bytes == 0) { + return 0; + } + const size_t max_fan_in = std::max(2, merge_limit_bytes / spill_buffer_bytes); + const size_t input_count = std::min(spill_file_count, max_fan_in); + const size_t max_size = std::numeric_limits::max(); + const size_t input_bytes = input_count > max_size / spill_buffer_bytes + ? max_size + : input_count * spill_buffer_bytes; + // VSortedRunMerger materializes one block per input cursor plus the block being emitted. + return input_bytes > max_size - spill_buffer_bytes ? max_size + : input_bytes + spill_buffer_bytes; +} + +} // namespace doris diff --git a/be/src/exec/operator/operator.h b/be/src/exec/operator/operator.h index becffbb171ee55..d60c8fe23209f4 100644 --- a/be/src/exec/operator/operator.h +++ b/be/src/exec/operator/operator.h @@ -631,6 +631,10 @@ class DataSinkOperatorXBase : public OperatorBase { [[nodiscard]] virtual size_t get_reserve_mem_size(RuntimeState* state, bool eos) { return state->minimum_operator_memory_required_bytes(); } + [[nodiscard]] virtual size_t get_reserve_mem_size(RuntimeState* state, bool eos, + const Block* block) { + return get_reserve_mem_size(state, eos); + } bool is_blockable(RuntimeState* state) const override { return state->get_sink_local_state()->is_blockable(); } 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..266dc654f4315f 100644 --- a/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp +++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp @@ -18,18 +18,35 @@ #include "exec/operator/spill_iceberg_table_sink_operator.h" #include "common/status.h" +#include "core/block/block.h" #include "exec/operator/iceberg_table_sink_operator.h" #include "exec/sink/writer/iceberg/viceberg_sort_writer.h" #include "exec/sink/writer/iceberg/viceberg_table_writer.h" namespace doris { +size_t iceberg_cold_writer_reserve_size(const Block& block, size_t writer_workspace_bytes) { + const size_t block_bytes = block.allocated_bytes(); + const size_t row_index_bytes = + std::min(std::numeric_limits::max() / sizeof(size_t), block.rows()) * + sizeof(size_t); + const size_t selected_and_retained_bytes = + std::min(std::numeric_limits::max() / 2, block_bytes) * 2; + size_t reserve = std::min(std::numeric_limits::max() - writer_workspace_bytes, + selected_and_retained_bytes) + + writer_workspace_bytes; + // Cold dispatch may allocate a selected block and a retained sorter copy before publication. + return std::min(std::numeric_limits::max() - reserve, row_index_bytes) + reserve; +} + SpillIcebergTableSinkLocalState::SpillIcebergTableSinkLocalState(DataSinkOperatorXBase* parent, RuntimeState* state) : Base(parent, state) {} Status SpillIcebergTableSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& info) { RETURN_IF_ERROR(Base::init(state, info)); + // Admission samples async sorter state, so the next block must wait until the prior append publishes it. + _writer->wait_for_processing_before_next_sink(); SCOPED_TIMER(exec_time_counter()); SCOPED_TIMER(_init_timer); @@ -51,30 +68,44 @@ bool SpillIcebergTableSinkLocalState::is_blockable() const { return true; } -size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state, bool eos) { +size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state, bool eos, + const Block* block) { 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. + // The final queued item may contain rows and also owns the reservation used by async finish(). + const size_t incoming_reserve = + block == nullptr ? state->minimum_operator_memory_required_bytes() + : iceberg_cold_writer_reserve_size( + *block, state->minimum_operator_memory_required_bytes()); + return iceberg_reserve_size(per_partition_reservations, incoming_reserve); } 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 +113,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( @@ -125,9 +161,10 @@ Status SpillIcebergTableSinkOperatorX::sink_impl(RuntimeState* state, Block* in_ return local_state.sink(state, in_block, eos); } -size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos) { +size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos, + const Block* block) { auto& local_state = get_local_state(state); - return local_state.get_reserve_mem_size(state, eos); + return local_state.get_reserve_mem_size(state, eos, block); } size_t SpillIcebergTableSinkOperatorX::revocable_mem_size(RuntimeState* state) const { 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..bd981531896c6c 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" @@ -41,7 +43,7 @@ class SpillIcebergTableSinkLocalState final Status open(RuntimeState* state) override; bool is_blockable() const override; - [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos); + [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos, const Block* block); Status revoke_memory(RuntimeState* state); size_t get_revocable_mem_size(RuntimeState* state) const; @@ -65,7 +67,7 @@ class SpillIcebergTableSinkOperatorX final Status sink_impl(RuntimeState* state, Block* in_block, bool eos) override; - size_t get_reserve_mem_size(RuntimeState* state, bool eos) override; + size_t get_reserve_mem_size(RuntimeState* state, bool eos, const Block* block) override; size_t revocable_mem_size(RuntimeState* state) const override; @@ -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 91991ce3aa248d..c401d16cf5ba58 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.cpp +++ b/be/src/exec/pipeline/pipeline_fragment_context.cpp @@ -26,6 +26,7 @@ #include #include + // IWYU pragma: no_include #include #include @@ -117,6 +118,7 @@ #include "exec/operator/union_source_operator.h" #include "exec/pipeline/dependency.h" #include "exec/pipeline/pipeline_task.h" +#include "exec/pipeline/report_exec_status_size.h" #include "exec/pipeline/task_scheduler.h" #include "exec/runtime_filter/runtime_filter_mgr.h" #include "exec/sort/topn_sorter.h" @@ -467,6 +469,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_external_file_report_state( + _runtime_state->external_file_report_state()); task_runtime_state->set_query_mem_tracker(_query_ctx->query_mem_tracker()); task_runtime_state->set_task_execution_context(shared_from_this()); @@ -2347,6 +2351,15 @@ std::string PipelineFragmentContext::_to_http_path(const std::string& file_name) return url.str(); } +void PipelineFragmentContext::_append_external_file_commit_data( + const ReportStatusRequest& req, TReportExecStatusParams* params) const { + // External-file cleanup remains BE-owned until the final report transfers commit metadata. + req.runtime_state->append_external_file_commit_data(params, req.done); + for (auto* rs : req.runtime_states) { + rs->append_external_file_commit_data(params, req.done); + } +} + void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& req) { DBUG_EXECUTE_IF("FragmentMgr::coordinator_callback.report_delay", { int random_seconds = req.status.is() ? 8 : 2; @@ -2358,6 +2371,11 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r DCHECK(req.status.ok() || req.done); // if !status.ok() => done if (req.coord_addr.hostname == "external") { // External query (flink/spark read tablets) not need to report to FE. + if (req.done) { + // Without a coordinator acknowledgement no external-write file may escape rollback. + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); + } return; } int callback_retries = 10; @@ -2378,6 +2396,10 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r static_cast(req.cancel_fn(Status::InternalError( "query_id: {}, couldn't get a client for {}, reason is {}", uid.to_string(), PrintThriftNetworkAddress(req.coord_addr), coord_status.to_string()))); + if (req.done) { + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); + } return; } @@ -2501,45 +2523,7 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r } } } - if (auto hpu = req.runtime_state->hive_partition_updates(); !hpu.empty()) { - params.__isset.hive_partition_updates = true; - params.hive_partition_updates.insert(params.hive_partition_updates.end(), hpu.begin(), - hpu.end()); - } else if (!req.runtime_states.empty()) { - for (auto* rs : req.runtime_states) { - if (auto rs_hpu = rs->hive_partition_updates(); !rs_hpu.empty()) { - params.__isset.hive_partition_updates = true; - params.hive_partition_updates.insert(params.hive_partition_updates.end(), - rs_hpu.begin(), rs_hpu.end()); - } - } - } - if (auto icd = req.runtime_state->iceberg_commit_datas(); !icd.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()) { - params.__isset.iceberg_commit_datas = true; - params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(), - rs_icd.begin(), rs_icd.end()); - } - } - } - - if (auto mcd = req.runtime_state->mc_commit_datas(); !mcd.empty()) { - params.__isset.mc_commit_datas = true; - params.mc_commit_datas.insert(params.mc_commit_datas.end(), mcd.begin(), mcd.end()); - } else if (!req.runtime_states.empty()) { - for (auto* rs : req.runtime_states) { - if (auto rs_mcd = rs->mc_commit_datas(); !rs_mcd.empty()) { - params.__isset.mc_commit_datas = true; - params.mc_commit_datas.insert(params.mc_commit_datas.end(), rs_mcd.begin(), - rs_mcd.end()); - } - } - } + _append_external_file_commit_data(req, ¶ms); req.runtime_state->get_unreported_errors(&(params.error_log)); params.__isset.error_log = (!params.error_log.empty()); @@ -2548,8 +2532,20 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r params.__set_backend_id(_exec_env->cluster_info()->backend_id); } + Status report_size_status = validate_report_exec_status_size( + params, req.runtime_state->coordinator_thrift_message_limit()); + if (!report_size_status.ok()) { + if (req.done) { + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); + } + req.cancel_fn(report_size_status); + return; + } + TReportExecStatusResult res; Status rpc_status; + bool report_outcome_ambiguous = false; VLOG_DEBUG << "reportExecStatus params is " << apache::thrift::ThriftDebugString(params).c_str(); @@ -2562,12 +2558,19 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r try { (*coord)->reportExecStatus(res, params); } catch (apache::thrift::transport::TTransportException& e) { + report_outcome_ambiguous = true; LOG(WARNING) << "Retrying ReportExecStatus. query id: " << print_id(req.query_id) << ", instance id: " << print_id(req.fragment_instance_id) << " to " << req.coord_addr << ", err: " << e.what(); rpc_status = coord->reopen(); if (!rpc_status.ok()) { + // The first request may have been consumed; keep files until metadata or orphan cleanup wins. + report_outcome_ambiguous = true; + if (req.done) { + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::AMBIGUOUS); + } req.cancel_fn(rpc_status); return; } @@ -2576,14 +2579,38 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r rpc_status = Status::create(res.status); } catch (apache::thrift::TException& e) { + report_outcome_ambiguous = true; rpc_status = Status::InternalError("ReportExecStatus() to {} failed: {}", PrintThriftNetworkAddress(req.coord_addr), e.what()); } + const bool requires_external_file_ack = params.__isset.iceberg_commit_datas; + if (rpc_status.ok() && requires_external_file_ack && + (!res.__isset.external_file_commit_data_accepted || + !res.external_file_commit_data_accepted)) { + rpc_status = Status::InternalError( + "Coordinator did not accept ownership of the external-file report"); + } + if (!rpc_status.ok()) { + if (req.done && !report_outcome_ambiguous) { + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); + } else if (req.done) { + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::AMBIGUOUS); + } LOG_INFO("Going to cancel query {} since report exec status got rpc failed: {}", print_id(req.query_id), rpc_status.to_string()); req.cancel_fn(rpc_status); + } else if (req.done && req.status.ok()) { + // Files remain rollback-owned until the coordinator has acknowledged the final metadata report. + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::ACKNOWLEDGED); + } else if (req.done) { + // An acknowledged error report confirms that FE will not publish this write's files. + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); } } @@ -2630,13 +2657,20 @@ Status PipelineFragmentContext::send_report(bool done) { .first_error_msg = first_error_msg, .cancel_fn = [this](const Status& reason) { cancel(reason); }}; auto ctx = std::dynamic_pointer_cast(shared_from_this()); - return _exec_env->fragment_mgr()->get_thread_pool()->submit_func([this, req, ctx]() { - SCOPED_ATTACH_TASK(ctx->get_query_ctx()->query_mem_tracker()); - _coordinator_callback(req); - if (!req.done) { - ctx->refresh_next_report_time(); - } - }); + Status submit_status = + _exec_env->fragment_mgr()->get_thread_pool()->submit_func([this, req, ctx]() { + SCOPED_ATTACH_TASK(ctx->get_query_ctx()->query_mem_tracker()); + _coordinator_callback(req); + if (!req.done) { + ctx->refresh_next_report_time(); + } + }); + if (!submit_status.ok() && req.done) { + // A rejected final callback can never transfer ownership to the coordinator. + req.runtime_state->finalize_external_file_report_cleanup( + ExternalFileReportOutcome::REJECTED); + } + return submit_status; } size_t PipelineFragmentContext::get_revocable_size(bool* has_running_task) const { diff --git a/be/src/exec/pipeline/pipeline_fragment_context.h b/be/src/exec/pipeline/pipeline_fragment_context.h index 1a63426738bef8..7243b0214d9fa0 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.h +++ b/be/src/exec/pipeline/pipeline_fragment_context.h @@ -153,6 +153,8 @@ class PipelineFragmentContext : public TaskExecutionContext { private: void _coordinator_callback(const ReportStatusRequest& req); + void _append_external_file_commit_data(const ReportStatusRequest& req, + TReportExecStatusParams* params) const; std::string _to_http_path(const std::string& file_name) const; void _release_resource(); diff --git a/be/src/exec/pipeline/pipeline_task.cpp b/be/src/exec/pipeline/pipeline_task.cpp index 064e34fe2ca07c..fe1cbba6cbec92 100644 --- a/be/src/exec/pipeline/pipeline_task.cpp +++ b/be/src/exec/pipeline/pipeline_task.cpp @@ -664,7 +664,8 @@ Status PipelineTask::execute(bool* done) { ->task_controller() ->is_enable_reserve_memory() && workload_group && !(_wake_up_early || _dry_run)) { - const auto sink_reserve_size = _sink->get_reserve_mem_size(_state, _eos); + const auto sink_reserve_size = + _sink->get_reserve_mem_size(_state, _eos, _block.get()); if (sink_reserve_size > 0 && _should_trigger_revoking(sink_reserve_size)) { LOG(INFO) << fmt::format( diff --git a/be/src/exec/pipeline/report_exec_status_size.h b/be/src/exec/pipeline/report_exec_status_size.h new file mode 100644 index 00000000000000..b5920963bfbcd4 --- /dev/null +++ b/be/src/exec/pipeline/report_exec_status_size.h @@ -0,0 +1,42 @@ +// 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 "common/status.h" +#include "util/thrift_util.h" + +namespace doris { + +inline Status validate_report_exec_status_size(const TReportExecStatusParams& params, + size_t thrift_limit) { + ThriftSerializer serializer(false, 256); + uint32_t serialized_size = 0; + uint8_t* buffer = nullptr; + RETURN_IF_ERROR(serializer.serialize(¶ms, &serialized_size, &buffer)); + // Include the args field header and RPC method/version/sequence envelope around the params. + constexpr size_t rpc_envelope_bytes = 64; + if (thrift_limit < rpc_envelope_bytes || serialized_size > thrift_limit - rpc_envelope_bytes) { + return Status::InternalError( + "ReportExecStatus exceeds the coordinator Thrift message limit"); + } + return Status::OK(); +} + +} // namespace doris diff --git a/be/src/exec/sink/viceberg_delete_sink.cpp b/be/src/exec/sink/viceberg_delete_sink.cpp index 172fdd28177c62..92d26b560f7160 100644 --- a/be/src/exec/sink/viceberg_delete_sink.cpp +++ b/be/src/exec/sink/viceberg_delete_sink.cpp @@ -34,6 +34,7 @@ #include "core/data_type/data_type_string.h" #include "core/data_type/data_type_struct.h" #include "exec/common/endian.h" +#include "exec/sink/writer/iceberg/iceberg_writer_compatibility.h" #include "exprs/vexpr.h" #include "format/table/deletion_vector.h" #include "format/table/iceberg_delete_file_reader_helper.h" @@ -203,6 +204,8 @@ Status VIcebergDeleteSink::init_properties(ObjectPool* pool) { Status VIcebergDeleteSink::open(RuntimeState* state, RuntimeProfile* profile) { _state = state; + RETURN_IF_ERROR(validate_iceberg_external_file_report_ack(state->query_options())); + // Initialize counters _written_rows_counter = ADD_COUNTER(profile, "RowsWritten", TUnit::UNIT); _send_data_timer = ADD_TIMER(profile, "SendDataTime"); @@ -283,13 +286,17 @@ 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; + } } } if (!_defer_file_cleanup_until_outer_close) { - _created_files.clear(); + _transfer_created_files_to_report_cleanup(); } return Status::OK(); @@ -299,11 +306,24 @@ void VIcebergDeleteSink::finish_deferred_file_cleanup(Status outer_status) { if (!outer_status.ok()) { _cleanup_created_files(); } else { - _created_files.clear(); + _transfer_created_files_to_report_cleanup(); } _defer_file_cleanup_until_outer_close = false; } +void VIcebergDeleteSink::_transfer_created_files_to_report_cleanup() { + DCHECK(_state != nullptr); + for (auto& created_file : _created_files) { + _state->add_rejected_external_file_report_cleanup( + [cleanup_fs = std::move(created_file.first), + cleanup_path = std::move(created_file.second)] { + WARN_IF_ERROR(cleanup_fs->delete_file(cleanup_path), + "failed to delete an Iceberg delete file after report failure"); + }); + } + _created_files.clear(); +} + void VIcebergDeleteSink::_cleanup_created_files() { for (const auto& [fs, path] : _created_files) { Status delete_status = fs->delete_file(path); diff --git a/be/src/exec/sink/viceberg_delete_sink.h b/be/src/exec/sink/viceberg_delete_sink.h index 55698ae0404b14..625134ba3a3d51 100644 --- a/be/src/exec/sink/viceberg_delete_sink.h +++ b/be/src/exec/sink/viceberg_delete_sink.h @@ -134,6 +134,7 @@ class VIcebergDeleteSink final : public AsyncResultWriter { Status _init_position_delete_output_exprs(); std::string _get_file_extension() const; void _cleanup_created_files(); + void _transfer_created_files_to_report_cleanup(); TDataSink _t_sink; RuntimeState* _state = nullptr; diff --git a/be/src/exec/sink/writer/async_result_writer.cpp b/be/src/exec/sink/writer/async_result_writer.cpp index 2fa2f6f92a7418..0999290370c1ac 100644 --- a/be/src/exec/sink/writer/async_result_writer.cpp +++ b/be/src/exec/sink/writer/async_result_writer.cpp @@ -43,6 +43,8 @@ Status AsyncResultWriter::sink(Block* block, bool eos) { add_block = _get_free_block(block, rows); } + // The pipeline reservation protects allocations performed after this block is dequeued. + auto reservation = thread_context()->thread_mem_tracker_mgr->take_reserved_memory(); std::lock_guard l(_m); // if io task failed, just return error status to // end the query @@ -54,9 +56,12 @@ Status AsyncResultWriter::sink(Block* block, bool eos) { if (_is_finished()) { _dependency->set_ready(); } - if (rows) { - _memory_used_counter->update(add_block->allocated_bytes()); - _data_queue.emplace_back(std::move(add_block)); + if (rows || eos) { + if (rows) { + _memory_used_counter->update(add_block->allocated_bytes()); + } + _data_queue.emplace_back(QueuedBlock { + .block = std::move(add_block), .reservation = std::move(reservation), .eos = eos}); if (!_data_queue_is_available() && !_is_finished()) { _dependency->block(); } @@ -70,17 +75,31 @@ Status AsyncResultWriter::sink(Block* block, bool eos) { return Status::OK(); } -std::unique_ptr AsyncResultWriter::_get_block_from_queue() { +AsyncResultWriter::QueuedBlock AsyncResultWriter::_get_block_from_queue() { std::lock_guard l(_m); DCHECK(!_data_queue.empty()); - auto block = std::move(_data_queue.front()); + auto queued = std::move(_data_queue.front()); _data_queue.pop_front(); + _queue_admission.begin_processing(); DCHECK(_dependency); if (_data_queue_is_available()) { _dependency->set_ready(); } - _memory_used_counter->update(-block->allocated_bytes()); - return block; + if (queued.block) { + _memory_used_counter->update(-queued.block->allocated_bytes()); + } + return queued; +} + +void AsyncResultWriter::_notify_block_processed() { + if (!_queue_admission.waits_for_processing()) { + return; + } + std::lock_guard l(_m); + _queue_admission.finish_processing(); + if (_data_queue_is_available()) { + _dependency->set_ready(); + } } Status AsyncResultWriter::start_writer(RuntimeState* state, RuntimeProfile* operator_profile) { @@ -130,6 +149,7 @@ void AsyncResultWriter::process_block(RuntimeState* state, RuntimeProfile* opera } DCHECK(_dependency); + bool reservation_held_for_finish = false; while (_writer_status.ok()) { ThreadCpuStopWatch cpu_time_stop_watch; cpu_time_stop_watch.start(); @@ -158,24 +178,48 @@ void AsyncResultWriter::process_block(RuntimeState* state, RuntimeProfile* opera //check if eos or writer error if ((_eos && _data_queue.empty()) || !_writer_status.ok()) { - _data_queue.clear(); break; } } //2) get the block from data queue and write to downstream - auto block = _get_block_from_queue(); - auto status = write(state, *block); + auto queued = _get_block_from_queue(); + thread_context()->thread_mem_tracker_mgr->adopt_reserved_memory( + std::move(queued.reservation)); + Status status = queued.block ? write(state, *queued.block) : Status::OK(); if (!status.ok()) [[unlikely]] { + thread_context()->thread_mem_tracker_mgr->shrink_reserved(); std::unique_lock l(_m); + _queue_admission.finish_processing(); _writer_status.update(status); - if (_is_finished()) { + if (_is_finished() || _data_queue_is_available()) { _dependency->set_ready(); } break; } - _return_free_block(std::move(block)); + if (queued.block) { + _return_free_block(std::move(queued.block)); + } + if (queued.eos) { + // Keep the final reservation through finish(), where buffered sorters are committed. + reservation_held_for_finish = true; + _notify_block_processed(); + break; + } + thread_context()->thread_mem_tracker_mgr->shrink_reserved(); + _notify_block_processed(); + } + + { + std::lock_guard l(_m); + drain_async_writer_queue(_data_queue, [this](const QueuedBlock& queued) { + if (queued.block) { + _memory_used_counter->update(-queued.block->allocated_bytes()); + } + }); + _queue_admission.finish_processing(); + _dependency->set_ready(); } bool need_finish = false; @@ -198,6 +242,9 @@ void AsyncResultWriter::process_block(RuntimeState* state, RuntimeProfile* opera Status st = finish(state); _writer_status.update(st); } + if (reservation_held_for_finish) { + thread_context()->thread_mem_tracker_mgr->shrink_reserved(); + } Status st = Status::OK(); { st = _writer_status.status(); } diff --git a/be/src/exec/sink/writer/async_result_writer.h b/be/src/exec/sink/writer/async_result_writer.h index 99d4f8eaa59eff..fe851a9171aae4 100644 --- a/be/src/exec/sink/writer/async_result_writer.h +++ b/be/src/exec/sink/writer/async_result_writer.h @@ -21,8 +21,10 @@ #include #include // IWYU pragma: keep +#include "exec/sink/writer/async_writer_queue_admission.h" #include "exec/sink/writer/result_writer.h" #include "exprs/vexpr_fwd.h" +#include "runtime/memory/thread_mem_tracker_mgr.h" #include "runtime/runtime_profile.h" namespace doris { @@ -36,6 +38,7 @@ class Dependency; class PipelineTask; class Block; + /* * In the pipeline execution engine, there are usually a large number of io operations on the sink side that * will block the limited execution threads of the pipeline execution engine, resulting in a sharp performance @@ -69,6 +72,10 @@ class AsyncResultWriter : public ResultWriter { void set_low_memory_mode(); + void wait_for_processing_before_next_sink() { + _queue_admission.wait_for_processing_before_next_sink(); + } + protected: Status _projection_block(Block& input_block, Block* output_block); const VExprContextSPtrs& _vec_output_expr_ctxs; @@ -77,21 +84,30 @@ class AsyncResultWriter : public ResultWriter { std::unique_ptr _get_free_block(Block*, size_t rows); private: + struct QueuedBlock { + std::unique_ptr block; + ReservedMemoryToken reservation; + bool eos = false; + }; + void process_block(RuntimeState* state, RuntimeProfile* operator_profile); - [[nodiscard]] bool _data_queue_is_available() const { return _data_queue.size() < QUEUE_SIZE; } + [[nodiscard]] bool _data_queue_is_available() const { + return _queue_admission.is_available(_data_queue.size()); + } [[nodiscard]] bool _is_finished() const { return !_writer_status.ok() || _eos; } void _set_ready_to_finish(); void _return_free_block(std::unique_ptr); - std::unique_ptr _get_block_from_queue(); + QueuedBlock _get_block_from_queue(); + void _notify_block_processed(); - static constexpr auto QUEUE_SIZE = 3; std::mutex _m; std::condition_variable _cv; - std::deque> _data_queue; + std::deque _data_queue; // Default value is ok AtomicStatus _writer_status; bool _eos = false; + AsyncWriterQueueAdmission _queue_admission; std::atomic_bool _low_memory_mode = false; std::shared_ptr _dependency; diff --git a/be/src/exec/sink/writer/async_writer_queue_admission.h b/be/src/exec/sink/writer/async_writer_queue_admission.h new file mode 100644 index 00000000000000..b5a73cb72aa881 --- /dev/null +++ b/be/src/exec/sink/writer/async_writer_queue_admission.h @@ -0,0 +1,53 @@ +// 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 + +namespace doris { + +inline constexpr size_t ASYNC_WRITER_QUEUE_SIZE = 3; + +class AsyncWriterQueueAdmission { +public: + void wait_for_processing_before_next_sink() { _wait_for_processing = true; } + void begin_processing() { _block_being_processed = _wait_for_processing; } + void finish_processing() { _block_being_processed = false; } + + [[nodiscard]] bool is_available(size_t queued_blocks) const { + return _wait_for_processing ? queued_blocks == 0 && !_block_being_processed + : queued_blocks < ASYNC_WRITER_QUEUE_SIZE; + } + + [[nodiscard]] bool waits_for_processing() const { return _wait_for_processing; } + +private: + bool _block_being_processed = false; + bool _wait_for_processing = false; +}; + +template +void drain_async_writer_queue(Queue& queue, BeforeRelease before_release) { + for (const auto& queued : queue) { + before_release(queued); + } + // Queued reservation tokens must be destroyed as soon as the writer reaches a terminal state. + queue.clear(); +} + +} // namespace doris diff --git a/be/src/exec/sink/writer/hive_multipart_compatibility.h b/be/src/exec/sink/writer/hive_multipart_compatibility.h new file mode 100644 index 00000000000000..c7546f314bfa8b --- /dev/null +++ b/be/src/exec/sink/writer/hive_multipart_compatibility.h @@ -0,0 +1,29 @@ +// 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 "io/fs/obj_storage_client.h" + +namespace doris { + +inline bool hive_multipart_protocol_supported(io::ObjStorageType provider, + bool supports_deferred_azure_multipart) { + return provider != io::ObjStorageType::AZURE || supports_deferred_azure_multipart; +} + +} // namespace doris diff --git a/be/src/exec/sink/writer/iceberg/iceberg_writer_compatibility.h b/be/src/exec/sink/writer/iceberg/iceberg_writer_compatibility.h new file mode 100644 index 00000000000000..ae06f4847a4c82 --- /dev/null +++ b/be/src/exec/sink/writer/iceberg/iceberg_writer_compatibility.h @@ -0,0 +1,35 @@ +// 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 "common/status.h" +#include "gen_cpp/PaloInternalService_types.h" + +namespace doris { + +inline Status validate_iceberg_external_file_report_ack(const TQueryOptions& query_options) { + if (!query_options.__isset.supports_external_file_report_ack || + !query_options.supports_external_file_report_ack) { + // A pre-ACK coordinator cannot safely take ownership of files created by this sink. + return Status::NotSupported( + "Iceberg writes require a coordinator that acknowledges external-file reports"); + } + return Status::OK(); +} + +} // namespace doris 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..396d81ad87e88a 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp @@ -17,6 +17,7 @@ #include "exec/sink/writer/iceberg/viceberg_sort_writer.h" +#include "exec/operator/iceberg_sorter_reserve_memory.h" #include "exec/spill/spill_file_manager.h" #include "exec/spill/spill_file_reader.h" #include "exec/spill/spill_file_writer.h" @@ -84,6 +85,27 @@ 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); + if (_sorter == nullptr) { + return {}; + } + auto reservation = _sorter->get_reserve_mem_size_components(state, eos); + if (eos && !_sorted_spill_files.empty()) { + size_t spill_file_count = _sorted_spill_files.size(); + if (_sorter->data_size() > 0) { + ++spill_file_count; + } + const size_t merge_workspace = + iceberg_spill_merge_workspace(spill_file_count, state->spill_buffer_size_bytes(), + state->spill_sort_merge_mem_limit_bytes()); + reservation.transient_workspace = + std::max(reservation.transient_workspace, merge_workspace); + } + return reservation; +} + Status VIcebergSortWriter::trigger_spill() { std::lock_guard lock(_sorter_mutex); if (_closed || _sorter == nullptr) { @@ -102,80 +124,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; + 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 = _sorter->prepare_for_read(false); - if (!internal_status.ok()) { - return internal_status; - } - 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 e3f8ed645edb26..0c8e65ee0ff541 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp @@ -27,6 +27,7 @@ #include "core/data_type/data_type_nullable.h" #include "core/data_type_serde/data_type_serde.h" #include "exec/sink/writer/iceberg/iceberg_partition_path.h" +#include "exec/sink/writer/iceberg/iceberg_writer_compatibility.h" #include "exec/sink/writer/iceberg/partition_transformers.h" #include "exec/sink/writer/iceberg/viceberg_partition_writer.h" #include "exec/sink/writer/iceberg/viceberg_sort_writer.h" @@ -45,12 +46,15 @@ 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) { _state = state; _operator_profile = profile; + RETURN_IF_ERROR(validate_iceberg_external_file_report_ack(state->query_options())); + // Get target file size from query options // If value is 0 or not set, use config::iceberg_sink_max_file_size _target_file_size_bytes = config::iceberg_sink_max_file_size; @@ -250,7 +254,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 @@ -267,6 +272,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()); @@ -276,6 +282,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); @@ -284,6 +291,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; } @@ -292,7 +300,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(); } @@ -310,6 +317,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()); @@ -319,6 +327,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); @@ -327,6 +336,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; } @@ -335,7 +345,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(); } @@ -384,10 +393,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(); @@ -396,8 +403,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( @@ -406,7 +413,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(); @@ -416,53 +422,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) { @@ -482,6 +488,7 @@ Status VIcebergTableWriter::close(Status status) { } } _partitions_to_writers.clear(); + _publish_active_writers(); } if (status.ok()) { SCOPED_TIMER(_operator_profile->total_time_counter()); @@ -497,7 +504,7 @@ Status VIcebergTableWriter::close(Status status) { if (!status.ok() || !result_status.ok()) { _cleanup_closed_files(); } else if (!_defer_file_cleanup_until_outer_close) { - _closed_files.clear(); + _transfer_closed_files_to_report_cleanup(); } return result_status; } @@ -508,11 +515,24 @@ void VIcebergTableWriter::finish_deferred_file_cleanup(Status outer_status) { if (!outer_status.ok()) { _cleanup_closed_files(); } else { - _closed_files.clear(); + _transfer_closed_files_to_report_cleanup(); } _defer_file_cleanup_until_outer_close = false; } +void VIcebergTableWriter::_transfer_closed_files_to_report_cleanup() { + DCHECK(_state != nullptr); + for (auto& closed_file : _closed_files) { + _state->add_rejected_external_file_report_cleanup( + [cleanup_fs = std::move(closed_file.first), + cleanup_path = std::move(closed_file.second)] { + WARN_IF_ERROR(cleanup_fs->delete_file(cleanup_path), + "failed to delete an Iceberg file after report failure"); + }); + } + _closed_files.clear(); +} + void VIcebergTableWriter::_cleanup_closed_files() { for (const auto& [fs, path] : _closed_files) { Status delete_status = fs->delete_file(path); 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 2cb83f73ed0691..ecda57961611f6 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h @@ -73,22 +73,16 @@ 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: FRIEND_TEST(VIcebergTableWriterTest, RejectMissingPartitionSource); + // The lifecycle fixture inspects snapshots to verify that cross-thread writer ownership stays stable. + friend class VIcebergTableWriterLifecycleTest; - // 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; + // 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, @@ -143,12 +137,14 @@ 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); void _cleanup_closed_files(); + void _transfer_closed_files_to_report_cleanup(); // Currently it is a copy, maybe it is better to use move semantics to eliminate it. TDataSink _t_sink; diff --git a/be/src/exec/sink/writer/vhive_partition_writer.cpp b/be/src/exec/sink/writer/vhive_partition_writer.cpp index 8331efac54bd47..658da3dbc8194f 100644 --- a/be/src/exec/sink/writer/vhive_partition_writer.cpp +++ b/be/src/exec/sink/writer/vhive_partition_writer.cpp @@ -21,10 +21,12 @@ #include "core/block/materialize_block.h" #include "core/column/column_map.h" +#include "exec/sink/writer/hive_multipart_compatibility.h" #include "format/transformer/vcsv_transformer.h" #include "format/transformer/vorc_transformer.h" #include "format/transformer/vparquet_transformer.h" #include "io/file_factory.h" +#include "io/fs/s3_file_system.h" #include "io/fs/s3_file_writer.h" #include "runtime/runtime_state.h" @@ -50,7 +52,10 @@ VHivePartitionWriter::VHivePartitionWriter(const TDataSink& t_sink, std::string _file_format_type(file_format_type), _hive_compress_type(hive_compress_type), _hive_serde_properties(hive_serde_properties), - _hadoop_conf(hadoop_conf) {} + _hadoop_conf(hadoop_conf), + _supports_deferred_azure_multipart( + t_sink.hive_table_sink.__isset.supports_deferred_azure_multipart && + t_sink.hive_table_sink.supports_deferred_azure_multipart) {} Status VHivePartitionWriter::open(RuntimeState* state, RuntimeProfile* operator_profile) { _state = state; @@ -64,6 +69,16 @@ Status VHivePartitionWriter::open(RuntimeState* state, RuntimeProfile* operator_ .path = fmt::format("{}/{}", _write_info.write_path, _get_target_file_name()), .fs_name {}}; _fs = DORIS_TRY(FileFactory::create_fs(fs_properties, file_description)); + if (auto* s3_fs = dynamic_cast(_fs.get()); + s3_fs != nullptr && + !hive_multipart_protocol_supported(s3_fs->client_holder()->s3_client_conf().provider, + _supports_deferred_azure_multipart)) { + // An old coordinator cannot publish namespaced Azure block IDs; lease expiry is not a + // compatibility fence, so reject before creating an upload that it could corrupt. + return Status::NotSupported( + "Azure Hive writes require a coordinator that supports deferred multipart " + "completion"); + } io::FileWriterOptions file_writer_options = {.used_by_s3_committer = true}; RETURN_IF_ERROR(_fs->create_file(file_description.path, &_file_writer, &file_writer_options)); @@ -147,6 +162,12 @@ Status VHivePartitionWriter::close(const Status& status) { } if (status_ok) { auto partition_update = _build_partition_update(); + if (partition_update.__isset.s3_mpu_pending_uploads) { + auto* s3_writer = dynamic_cast(_file_writer.get()); + DCHECK(s3_writer != nullptr); + // Until FE accepts the final report, BE remains the cleanup owner of staged uploads. + _state->add_rejected_external_file_report_cleanup(s3_writer->failed_report_cleanup()); + } _state->add_hive_partition_updates(partition_update); } return result_status; @@ -209,6 +230,10 @@ void VHivePartitionWriter::_add_s3_mpu_pending_upload_for_rollback() { if (!_build_s3_mpu_pending_upload(&s3_mpu_pending_upload)) { return; } + auto* s3_writer = dynamic_cast(_file_writer.get()); + DCHECK(s3_writer != nullptr); + // A failed write still relies on the final report to hand its staged upload to FE rollback. + _state->add_rejected_external_file_report_cleanup(s3_writer->failed_report_cleanup()); THivePartitionUpdate hive_partition_update; hive_partition_update.__set_name(_partition_name); diff --git a/be/src/exec/sink/writer/vhive_partition_writer.h b/be/src/exec/sink/writer/vhive_partition_writer.h index 0b124108623fa1..92e316a95c8e10 100644 --- a/be/src/exec/sink/writer/vhive_partition_writer.h +++ b/be/src/exec/sink/writer/vhive_partition_writer.h @@ -101,6 +101,7 @@ class VHivePartitionWriter { TFileCompressType::type _hive_compress_type; const THiveSerDeProperties* _hive_serde_properties; const std::map& _hadoop_conf; + bool _supports_deferred_azure_multipart = false; std::shared_ptr _fs = nullptr; 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..52dea823f18302 100644 --- a/be/src/io/fs/azure_obj_storage_client.cpp +++ b/be/src/io/fs/azure_obj_storage_client.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -28,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -35,6 +37,7 @@ #include #include #include +#include #include #include #include @@ -46,7 +49,6 @@ #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" using namespace Azure::Storage::Blobs; @@ -64,10 +66,37 @@ 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 encode_azure_block_id(std::string_view upload_id, int part_num) { + uint32_t upload_namespace = 0x811C9DC5U; + for (unsigned char byte : upload_id) { + upload_namespace = (upload_namespace ^ byte) * 0x01000193U; + } + uint32_t namespaced_part = upload_namespace + static_cast(part_num); + // Four decoded bytes remain compatible with legacy residual blocks. Writer isolation is + // enforced by the target blob lease because no 32-bit namespace can identify every upload. + std::array raw_id {}; + for (size_t i = 0; i < raw_id.size(); ++i) { + raw_id[i] = static_cast(namespaced_part >> (i * 8)); + } + Aws::Utils::ByteBuffer bytes(raw_id.data(), raw_id.size()); + return Aws::Utils::HashingUtils::Base64Encode(bytes); +} + +constexpr std::string_view MULTIPART_LEASE_PREFIX = "doris-azure-lease-v1:"; +constexpr std::chrono::seconds MULTIPART_LEASE_DURATION {60}; + +std::optional azure_multipart_lease_id(std::string_view upload_id) { + if (upload_id.starts_with(MULTIPART_LEASE_PREFIX) && + upload_id.size() > MULTIPART_LEASE_PREFIX.size()) { + return upload_id.substr(MULTIPART_LEASE_PREFIX.size()); + } + return std::nullopt; +} + +void renew_azure_multipart_lease(BlobClient& target_blob, std::string_view lease_id) { + BlobLeaseClient lease(target_blob, std::string(lease_id)); + // A renewal failure loses the upload-generation fence even if the same ID is acquirable later. + lease.Renew(); } // Rate limiting is applied by RateLimitedObjStorageClient, the decorator that @@ -79,6 +108,10 @@ constexpr char BlobNotFound[] = "BlobNotFound"; namespace doris::io { +std::string azure_multipart_block_id(std::string_view upload_id, int part_num) { + return encode_azure_block_id(upload_id, part_num); +} + // As Azure's doc said, the batch size is 256 // You can find out the num in https://learn.microsoft.com/en-us/rest/api/storageservices/blob-batch?tabs=microsoft-entra-id // > Each batch request supports a maximum of 256 subrequests. @@ -194,11 +227,29 @@ struct AzureBatchDeleter { std::vector> deferred_resps; }; -// Azure would do nothing ObjectStorageUploadResponse AzureObjStorageClient::create_multipart_upload( const ObjectStoragePathOptions& opts) { + auto target_blob = _client->GetBlobClient(opts.key); + auto target_client = target_blob.AsBlockBlobClient(); + std::string lease_id = BlobLeaseClient::CreateUniqueLeaseId(); + std::string upload_id = fmt::format("{}{}", MULTIPART_LEASE_PREFIX, lease_id); + auto resp = do_azure_client_call( + [&]() { + uint8_t empty = 0; + Azure::Core::IO::MemoryBodyStream empty_body(&empty, 0); + // The reservation makes an absent blob leaseable but remains uncommitted and + // invisible to normal listings until Put Block List publishes the real data. + target_client.StageBlock(azure_multipart_block_id(upload_id, 0), empty_body); + auto lease = + BlobLeaseClient(target_blob, lease_id).Acquire(MULTIPART_LEASE_DURATION); + upload_id = fmt::format("{}{}", MULTIPART_LEASE_PREFIX, lease.Value.LeaseId); + }, + opts, _tls_debug_context); return ObjectStorageUploadResponse { - .resp = ObjectStorageResponse::OK(), + .resp = resp, + .upload_id = resp.status.code == ErrorCode::OK + ? std::make_optional(std::move(upload_id)) + : std::nullopt, }; } @@ -216,35 +267,86 @@ ObjectStorageResponse AzureObjStorageClient::put_object(const ObjectStoragePathO ObjectStorageUploadResponse AzureObjStorageClient::upload_part(const ObjectStoragePathOptions& opts, std::string_view stream, int part_num) { - auto client = _client->GetBlockBlobClient(opts.key); + DCHECK(opts.upload_id.has_value()); + auto target_blob = _client->GetBlobClient(opts.key); + auto client = target_blob.AsBlockBlobClient(); + std::string block_id = azure_multipart_block_id(*opts.upload_id, 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); + auto lease_id = azure_multipart_lease_id(*opts.upload_id); + if (lease_id.has_value()) { + renew_azure_multipart_lease(target_blob, *lease_id); + StageBlockOptions stage_opts; + stage_opts.AccessConditions.LeaseId = std::string(*lease_id); + client.StageBlock(block_id, memory_body, stage_opts); + } else { + 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, }; } ObjectStorageResponse AzureObjStorageClient::complete_multipart_upload( const ObjectStoragePathOptions& opts, const std::vector& completed_parts) { - auto client = _client->GetBlockBlobClient(opts.key); + DCHECK(opts.upload_id.has_value()); + auto target_blob = _client->GetBlobClient(opts.key); + auto target_client = target_blob.AsBlockBlobClient(); 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); }); - return do_azure_client_call( + std::ranges::transform(completed_parts, std::back_inserter(string_block_ids), + [&opts](const ObjectCompleteMultiPart& i) { + return azure_multipart_block_id(*opts.upload_id, i.part_num); + }); + auto resp = do_azure_client_call( [&]() { SCOPED_BVAR_LATENCY(s3_bvar::s3_multi_part_upload_latency); - client.CommitBlockList(string_block_ids); + // Put Block List atomically replaces the committed blob; no scan-visible staging blob exists. + auto lease_id = azure_multipart_lease_id(*opts.upload_id); + if (lease_id.has_value()) { + renew_azure_multipart_lease(target_blob, *lease_id); + CommitBlockListOptions commit_opts; + commit_opts.AccessConditions.LeaseId = std::string(*lease_id); + target_client.CommitBlockList(string_block_ids, commit_opts); + } else { + target_client.CommitBlockList(string_block_ids); + } }, opts, _tls_debug_context); + if (resp.status.code == ErrorCode::OK) { + if (auto lease_id = azure_multipart_lease_id(*opts.upload_id); lease_id.has_value()) { + auto release_resp = do_azure_client_call( + [&]() { BlobLeaseClient(target_blob, std::string(*lease_id)).Release(); }, opts, + _tls_debug_context); + if (release_resp.status.code != ErrorCode::OK) { + LOG(WARNING) << "Azure multipart commit succeeded but its finite lease could not " + "be released; it will expire automatically"; + } + } + } + return resp; +} + +ObjectStorageResponse AzureObjStorageClient::abort_multipart_upload( + const ObjectStoragePathOptions& opts) { + DCHECK(opts.upload_id.has_value()); + if (auto lease_id = azure_multipart_lease_id(*opts.upload_id); lease_id.has_value()) { + auto target_blob = _client->GetBlobClient(opts.key); + return do_azure_client_call( + [&]() { BlobLeaseClient(target_blob, std::string(*lease_id)).Release(); }, opts, + _tls_debug_context); + } + // Azure cannot delete one upload's uncommitted blocks without changing the committed blob. + // Leaving them to service GC preserves the last successfully published value. + return ObjectStorageResponse::OK(); } ObjectStorageHeadResponse AzureObjStorageClient::head_object(const ObjectStoragePathOptions& opts) { diff --git a/be/src/io/fs/azure_obj_storage_client.h b/be/src/io/fs/azure_obj_storage_client.h index 7d1cecc502e44d..de4ea5459a7cf7 100644 --- a/be/src/io/fs/azure_obj_storage_client.h +++ b/be/src/io/fs/azure_obj_storage_client.h @@ -33,6 +33,7 @@ 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); +std::string azure_multipart_block_id(std::string_view upload_id, int part_num); class AzureObjStorageClient final : public ObjStorageClient { public: @@ -49,6 +50,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..63eb6e32c44c25 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,62 @@ 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(); +} + +std::function S3FileWriter::failed_report_cleanup() const { + auto client_holder = _obj_client; + auto path_opts = _obj_storage_path_opts; + return [client_holder = std::move(client_holder), path_opts = std::move(path_opts)] { + // The writer is already CLOSED, but ownership was not transferred; bypass abort()'s + // state guard while retaining the provider client and the exact upload identity. + const auto& client = client_holder->get(); + if (client == nullptr) { + LOG(WARNING) << "failed to abort a rejected external-file report: invalid object " + "storage client"; + return; + } + auto response = client->abort_multipart_upload(path_opts); + if (response.status.code != ErrorCode::OK) { + LOG(WARNING) << "failed to abort a rejected external-file report for " + << path_opts.path.native() << ": " << response.status.msg; + } + }; +} + +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 +151,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 +215,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 +229,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..83a2100c0ff943 100644 --- a/be/src/io/fs/s3_file_writer.h +++ b/be/src/io/fs/s3_file_writer.h @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -71,11 +72,14 @@ class S3FileWriter final : public FileWriter { } Status close(bool non_block = false) override; + Status abort() override; Status try_finish_close() override; + std::function failed_report_cleanup() const; + 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 +122,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/memory/thread_mem_tracker_mgr.cpp b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp index 3a0f0c7972fc6a..8e0f6cea6874a0 100644 --- a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp +++ b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp @@ -25,6 +25,46 @@ namespace doris { +void ReservedMemoryToken::release() { + if (_bytes == 0 && _untracked_bytes == 0) { + return; + } + // A queued item may be discarded after an async failure; its reservation still needs full rollback. + GlobalMemoryArbitrator::shrink_process_reserved(_bytes + _untracked_bytes); + _limiter_tracker->shrink_reserved(_bytes + _untracked_bytes); + _limiter_tracker->release(_bytes); + if (auto wg = _wg_wptr.lock()) { + wg->sub_wg_refresh_interval_memory_growth(_bytes); + } + _bytes = 0; + _untracked_bytes = 0; +} + +ReservedMemoryToken ThreadMemTrackerMgr::take_reserved_memory() { + CHECK(init()); + if (_reserved_mem == 0) { + return {}; + } + ReservedMemoryToken token(_limiter_tracker_sptr, _wg_wptr, _reserved_mem, _untracked_mem); + // Accounting remains reserved globally; only its thread-local ownership moves into the token. + _reserved_mem = 0; + _untracked_mem = 0; + return token; +} + +void ThreadMemTrackerMgr::adopt_reserved_memory(ReservedMemoryToken&& token) { + CHECK(init()); + if (token._bytes == 0 && token._untracked_bytes == 0) { + return; + } + flush_untracked_mem(); + CHECK(token._limiter_tracker == _limiter_tracker_sptr); + _reserved_mem += token._bytes; + _untracked_mem += token._untracked_bytes; + token._bytes = 0; + token._untracked_bytes = 0; +} + void ThreadMemTrackerMgr::attach_limiter_tracker( const std::shared_ptr& mem_tracker) { DCHECK(mem_tracker); diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.h b/be/src/runtime/memory/thread_mem_tracker_mgr.h index a24e32b205abe7..ec6461592468e9 100644 --- a/be/src/runtime/memory/thread_mem_tracker_mgr.h +++ b/be/src/runtime/memory/thread_mem_tracker_mgr.h @@ -25,6 +25,7 @@ #include #include #include +#include #include #include "common/be_mock_util.h" @@ -39,6 +40,42 @@ namespace doris { +class ReservedMemoryToken { +public: + ReservedMemoryToken() = default; + ReservedMemoryToken(const ReservedMemoryToken&) = delete; + ReservedMemoryToken& operator=(const ReservedMemoryToken&) = delete; + ReservedMemoryToken(ReservedMemoryToken&& other) noexcept { *this = std::move(other); } + ReservedMemoryToken& operator=(ReservedMemoryToken&& other) noexcept { + if (this != &other) { + release(); + _limiter_tracker = std::move(other._limiter_tracker); + _wg_wptr = std::move(other._wg_wptr); + _bytes = std::exchange(other._bytes, 0); + _untracked_bytes = std::exchange(other._untracked_bytes, 0); + } + return *this; + } + ~ReservedMemoryToken() { release(); } + [[nodiscard]] int64_t bytes() const { return _bytes; } + +private: + friend class ThreadMemTrackerMgr; + ReservedMemoryToken(std::shared_ptr limiter_tracker, + std::weak_ptr wg_wptr, int64_t bytes, + int64_t untracked_bytes) + : _limiter_tracker(std::move(limiter_tracker)), + _wg_wptr(std::move(wg_wptr)), + _bytes(bytes), + _untracked_bytes(untracked_bytes) {} + void release(); + + std::shared_ptr _limiter_tracker; + std::weak_ptr _wg_wptr; + int64_t _bytes = 0; + int64_t _untracked_bytes = 0; +}; + constexpr size_t SYNC_PROC_RESERVED_INTERVAL_BYTES = (1ULL << 20); // 1M static std::string MEMORY_ORPHAN_CHECK_MSG = "The ThreadContext of the current thread not attach a valid MemoryTracker. after the " @@ -99,6 +136,9 @@ class ThreadMemTrackerMgr { void shrink_reserved(); + ReservedMemoryToken take_reserved_memory(); + void adopt_reserved_memory(ReservedMemoryToken&& token); + MemTrackerLimiter* limiter_mem_tracker() { CHECK(init()); return _limiter_tracker; diff --git a/be/src/runtime/runtime_state.cpp b/be/src/runtime/runtime_state.cpp index 5dfd027d42d4ce..f2052910faf11b 100644 --- a/be/src/runtime/runtime_state.cpp +++ b/be/src/runtime/runtime_state.cpp @@ -52,12 +52,95 @@ #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)); + + // This is an early per-vector guard only; the assembled RPC is measured again before send. + constexpr size_t report_envelope_headroom = 1024 * 1024; + const size_t thrift_limit = coordinator_thrift_message_limit(); + const size_t commit_data_limit = + thrift_limit > report_envelope_headroom ? thrift_limit - report_envelope_headroom : 0; + std::lock_guard budget_lock(_external_file_report_state->mutex); + // Parallel task states share this budget because FE receives their vectors in one fragment report. + if (_external_file_report_state->iceberg_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); + _external_file_report_state->iceberg_serialized_bytes += serialized_size + sizeof(uint32_t); + _iceberg_commit_datas.emplace_back(std::move(iceberg_commit_data)); + return Status::OK(); +} + +size_t RuntimeState::coordinator_thrift_message_limit() const { + int32_t effective_thrift_limit = std::max(config::thrift_max_message_size, 0); + if (_query_options.__isset.coordinator_thrift_max_message_size && + _query_options.coordinator_thrift_max_message_size > 0) { + // An older FE omits this field; otherwise the receiver's smaller limit is authoritative. + effective_thrift_limit = std::min(effective_thrift_limit, + _query_options.coordinator_thrift_max_message_size); + } + return static_cast(effective_thrift_limit); +} + +void RuntimeState::append_external_file_commit_data(TReportExecStatusParams* params, + bool final_report) const { + if (!final_report) { + // Ownership-bearing commit vectors must only appear in the final report that transfers them. + return; + } + if (auto updates = hive_partition_updates(); !updates.empty()) { + params->__isset.hive_partition_updates = true; + params->hive_partition_updates.insert(params->hive_partition_updates.end(), updates.begin(), + updates.end()); + } + append_iceberg_commit_datas(¶ms->iceberg_commit_datas); + if (!params->iceberg_commit_datas.empty()) { + params->__isset.iceberg_commit_datas = true; + } + if (auto commit_datas = mc_commit_datas(); !commit_datas.empty()) { + params->__isset.mc_commit_datas = true; + params->mc_commit_datas.insert(params->mc_commit_datas.end(), commit_datas.begin(), + commit_datas.end()); + } +} + +void RuntimeState::add_rejected_external_file_report_cleanup(std::function cleanup) { + std::lock_guard lock(_external_file_report_state->mutex); + _external_file_report_state->rejected_report_cleanups.emplace_back(std::move(cleanup)); +} + +void RuntimeState::finalize_external_file_report_cleanup(ExternalFileReportOutcome outcome) { + std::vector> cleanups; + { + std::lock_guard lock(_external_file_report_state->mutex); + if (outcome == ExternalFileReportOutcome::ACKNOWLEDGED) { + _external_file_report_state->rejected_report_cleanups.clear(); + return; + } + if (outcome == ExternalFileReportOutcome::AMBIGUOUS) { + // A consumed request with a lost ACK may already be publishing these files. + return; + } + cleanups.swap(_external_file_report_state->rejected_report_cleanups); + } + for (auto& cleanup : cleanups) { + cleanup(); + } +} + 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..8ce0595e91f05b 100644 --- a/be/src/runtime/runtime_state.h +++ b/be/src/runtime/runtime_state.h @@ -74,6 +74,17 @@ class RuntimeFilterConsumer; class RuntimeFilterProducer; class TaskExecutionContext; +class ExternalFileReportState { + friend class RuntimeState; + +private: + std::mutex mutex; + size_t iceberg_serialized_bytes = 0; + std::vector> rejected_report_cleanups; +}; + +enum class ExternalFileReportOutcome { ACKNOWLEDGED, REJECTED, AMBIGUOUS }; + // 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 +534,27 @@ 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); + + size_t coordinator_thrift_message_limit() const; + + void append_external_file_commit_data(TReportExecStatusParams* params, bool final_report) const; + + void add_rejected_external_file_report_cleanup(std::function cleanup); + + void finalize_external_file_report_cleanup(ExternalFileReportOutcome outcome); + + void set_external_file_report_state(std::shared_ptr report_state) { + _external_file_report_state = std::move(report_state); + } + + const std::shared_ptr& external_file_report_state() const { + return _external_file_report_state; } std::vector mc_commit_datas() const { @@ -976,6 +1000,8 @@ class RuntimeState { mutable std::mutex _iceberg_commit_datas_mutex; std::vector _iceberg_commit_datas; + std::shared_ptr _external_file_report_state = + 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..92aa896daebcf2 --- /dev/null +++ b/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp @@ -0,0 +1,121 @@ +// 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 + +#include "core/block/block.h" +#include "core/column/column_string.h" +#include "core/data_type/data_type_string.h" +#include "exec/operator/iceberg_sorter_reserve_memory.h" +#include "exec/sink/writer/async_writer_queue_admission.h" +#include "exec/sink/writer/hive_multipart_compatibility.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)); +} + +TEST(SpillIcebergTableSinkOperatorTest, ReservesIncomingBlockBeforeAnyPartitionWriterExists) { + std::vector no_published_sorters; + + EXPECT_EQ(6 * 1024 * 1024, iceberg_reserve_size(no_published_sorters, 6 * 1024 * 1024)); +} + +TEST(SpillIcebergTableSinkOperatorTest, ColdWriterReserveUsesFirstBlockLargerThanOperatorFloor) { + constexpr size_t operator_floor = 32 * 1024 * 1024; + auto strings = ColumnString::create(); + std::string payload(40 * 1024 * 1024, 'x'); + strings->insert_data(payload.data(), payload.size()); + Block block; + block.insert({std::move(strings), std::make_shared(), "payload"}); + + ASSERT_GT(block.allocated_bytes(), operator_floor); + EXPECT_GT(iceberg_cold_writer_reserve_size(block, operator_floor), 2 * block.allocated_bytes()); +} + +TEST(SpillIcebergTableSinkOperatorTest, ReservesAllMergeInputsAndOutputAtEos) { + constexpr size_t MB = 1024 * 1024; + + EXPECT_EQ(72 * MB, iceberg_spill_merge_workspace(12, 8 * MB, 64 * MB)); + EXPECT_EQ(32 * MB, iceberg_spill_merge_workspace(3, 8 * MB, 64 * MB)); +} + +TEST(SpillIcebergTableSinkOperatorTest, WaitsUntilDequeuedBlockUpdatesSorterState) { + AsyncWriterQueueAdmission stateful_admission; + stateful_admission.wait_for_processing_before_next_sink(); + + EXPECT_TRUE(stateful_admission.is_available(0)); + EXPECT_FALSE(stateful_admission.is_available(1)); + stateful_admission.begin_processing(); + // Dequeueing does not admit block 2 until block 1 changes the state sampled by admission. + EXPECT_FALSE(stateful_admission.is_available(0)); + stateful_admission.finish_processing(); + EXPECT_TRUE(stateful_admission.is_available(0)); + + // Writers without state-dependent admission retain the existing three-block queue behavior. + AsyncWriterQueueAdmission buffered_admission; + buffered_admission.begin_processing(); + EXPECT_TRUE(buffered_admission.is_available(2)); + EXPECT_FALSE(buffered_admission.is_available(3)); +} + +TEST(SpillIcebergTableSinkOperatorTest, TerminalWriterDrainsQueuedReservations) { + int live_reservations = 0; + struct Reservation { + explicit Reservation(int* live) : live(live) { ++*live; } + ~Reservation() { --*live; } + int* live; + }; + struct Queued { + size_t bytes; + std::unique_ptr reservation; + }; + std::deque queue; + queue.push_back({7, std::make_unique(&live_reservations)}); + queue.push_back({11, std::make_unique(&live_reservations)}); + size_t released_bytes = 0; + + drain_async_writer_queue(queue, [&](const Queued& queued) { released_bytes += queued.bytes; }); + + EXPECT_TRUE(queue.empty()); + EXPECT_EQ(0, live_reservations); + EXPECT_EQ(18, released_bytes); +} + +TEST(SpillIcebergTableSinkOperatorTest, AzureDeferredMultipartRequiresCoordinatorCapability) { + EXPECT_TRUE(hive_multipart_protocol_supported(io::ObjStorageType::AWS, false)); + EXPECT_FALSE(hive_multipart_protocol_supported(io::ObjStorageType::AZURE, false)); + EXPECT_TRUE(hive_multipart_protocol_supported(io::ObjStorageType::AZURE, true)); +} + +} // namespace doris diff --git a/be/test/exec/sink/viceberg_delete_sink_test.cpp b/be/test/exec/sink/viceberg_delete_sink_test.cpp index 7faa77ed702c68..e2897125e497f3 100644 --- a/be/test/exec/sink/viceberg_delete_sink_test.cpp +++ b/be/test/exec/sink/viceberg_delete_sink_test.cpp @@ -94,6 +94,18 @@ TEST_F(VIcebergDeleteSinkTest, TestInitProperties) { ASSERT_TRUE(status.ok()); } +TEST_F(VIcebergDeleteSinkTest, RejectsCoordinatorWithoutExternalFileReportAck) { + VExprContextSPtrs output_exprs; + auto sink = std::make_shared(_t_data_sink, output_exprs, nullptr, nullptr); + RuntimeState state; + RuntimeProfile profile("test"); + + Status status = sink->open(&state, &profile); + + EXPECT_TRUE(status.is()); + EXPECT_NE(std::string::npos, status.to_string().find("acknowledges external-file reports")); +} + TEST_F(VIcebergDeleteSinkTest, TestGetRowIdColumnIndex) { VExprContextSPtrs output_exprs; auto sink = std::make_shared(_t_data_sink, output_exprs, nullptr, nullptr); 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..d486672c0c83db --- /dev/null +++ b/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp @@ -0,0 +1,156 @@ +// 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" +#include "runtime/runtime_state.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 VIcebergTableWriterLifecycleTest : 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(VIcebergTableWriterLifecycleTest, RejectsCoordinatorWithoutExternalFileReportAck) { + VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr); + RuntimeState state; + RuntimeProfile profile("test"); + + Status status = writer.open(&state, &profile); + + EXPECT_TRUE(status.is()); + EXPECT_NE(std::string::npos, status.to_string().find("acknowledges external-file reports")); +} + +TEST_F(VIcebergTableWriterLifecycleTest, 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(VIcebergTableWriterLifecycleTest, 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(VIcebergTableWriterLifecycleTest, 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/exec/sink/writer/vhive_partition_writer_report_lifecycle_test.cpp b/be/test/exec/sink/writer/vhive_partition_writer_report_lifecycle_test.cpp new file mode 100644 index 00000000000000..21234083984ca4 --- /dev/null +++ b/be/test/exec/sink/writer/vhive_partition_writer_report_lifecycle_test.cpp @@ -0,0 +1,208 @@ +// 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 +#include +#include + +#include "exec/pipeline/pipeline_fragment_context.h" +#include "exec/sink/writer/vhive_partition_writer.h" +#include "format/transformer/vfile_format_transformer.h" +#include "io/fs/s3_file_system.h" +#include "io/fs/s3_file_writer.h" +#include "runtime/exec_env.h" +#include "testutil/mock/mock_runtime_state.h" + +namespace doris { +namespace { + +class RecordingObjStorageClient final : public io::ObjStorageClient { +public: + io::ObjectStorageUploadResponse create_multipart_upload( + const io::ObjectStoragePathOptions&) override { + return {.resp = io::ObjectStorageResponse::OK(), .upload_id = "upload-id"}; + } + + io::ObjectStorageResponse put_object(const io::ObjectStoragePathOptions&, + std::string_view) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageUploadResponse upload_part(const io::ObjectStoragePathOptions&, + std::string_view, int) override { + return {.resp = io::ObjectStorageResponse::OK(), .etag = "etag"}; + } + + io::ObjectStorageResponse complete_multipart_upload( + const io::ObjectStoragePathOptions&, + const std::vector&) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageResponse abort_multipart_upload( + const io::ObjectStoragePathOptions& opts) override { + ++abort_count; + aborted_upload_id = opts.upload_id.value_or(""); + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageHeadResponse head_object(const io::ObjectStoragePathOptions&) override { + return {.resp = io::ObjectStorageResponse::OK(), .file_size = 0}; + } + + io::ObjectStorageResponse get_object(const io::ObjectStoragePathOptions&, void*, size_t, size_t, + size_t*) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageResponse list_objects(const io::ObjectStoragePathOptions&, + std::vector*) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageResponse delete_objects(const io::ObjectStoragePathOptions&, + std::vector) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageResponse delete_object(const io::ObjectStoragePathOptions&) override { + return io::ObjectStorageResponse::OK(); + } + + io::ObjectStorageResponse delete_objects_recursively( + const io::ObjectStoragePathOptions&) override { + return io::ObjectStorageResponse::OK(); + } + + std::string generate_presigned_url(const io::ObjectStoragePathOptions&, int64_t, + const S3ClientConf&) override { + return {}; + } + + int abort_count = 0; + std::string aborted_upload_id; +}; + +class FixedLengthTransformer final : public VFileFormatTransformer { +public: + explicit FixedLengthTransformer(const VExprContextSPtrs& output_exprs) + : VFileFormatTransformer(nullptr, output_exprs, false) {} + + Status open() override { return Status::OK(); } + Status write(const Block&) override { return Status::OK(); } + Status close() override { return Status::OK(); } + int64_t written_len() override { return 64; } +}; + +std::unique_ptr create_closed_hive_writer( + RuntimeState* state, const VExprContextSPtrs& output_exprs, + const std::shared_ptr& client) { + THiveTableSink hive_sink; + TDataSink sink; + sink.__set_type(TDataSinkType::HIVE_TABLE_SINK); + sink.__set_hive_table_sink(hive_sink); + VHivePartitionWriter::WriteInfo write_info {.write_path = "s3://bucket/staging", + .original_write_path = "s3://bucket/table", + .target_path = "s3://bucket/table", + .file_type = TFileType::FILE_S3, + .broker_addresses = {}}; + static const std::map hadoop_conf; + auto writer = std::make_unique( + sink, "", TUpdateMode::APPEND, output_exprs, std::vector {}, + std::move(write_info), "part", 0, TFileFormatType::FORMAT_PARQUET, + TFileCompressType::PLAIN, nullptr, hadoop_conf); + + auto holder = std::make_shared(S3ClientConf {}); + holder->_client = client; + io::FileWriterOptions options {.used_by_s3_committer = true}; + auto file_writer = + std::make_unique(holder, "bucket", "table/part.parquet", &options); + file_writer->_obj_storage_path_opts.upload_id = "upload-id"; + file_writer->_state = io::FileWriter::State::CLOSED; + writer->_file_writer = std::move(file_writer); + writer->_file_format_transformer = std::make_unique(output_exprs); + writer->_state = state; + EXPECT_TRUE(writer->close(Status::OK()).ok()); + return writer; +} + +std::shared_ptr create_fragment_context(TUniqueId query_id) { + auto query_ctx = MockQueryContext::create(query_id); + return std::make_shared(query_id, TPipelineFragmentParams(), query_ctx, + ExecEnv::GetInstance(), + [](RuntimeState*, Status*) {}); +} + +ReportStatusRequest report_request(RuntimeState* state, bool done) { + TNetworkAddress address; + address.hostname = "external"; + return {.status = Status::OK(), + .runtime_states = {}, + .done = done, + .coord_addr = address, + .query_id = TUniqueId(), + .fragment_id = 0, + .fragment_instance_id = TUniqueId(), + .backend_num = 0, + .runtime_state = state, + .load_error_url = "", + .first_error_msg = "", + .cancel_fn = [](const Status&) {}}; +} + +} // namespace + +TEST(VHivePartitionWriterReportLifecycleTest, + PeriodicReportDefersMetadataAndRejectedFinalReportAbortsProviderUpload) { + MockRuntimeState state; + VExprContextSPtrs output_exprs; + auto client = std::make_shared(); + auto writer = create_closed_hive_writer(&state, output_exprs, client); + auto context = create_fragment_context(TUniqueId()); + + TReportExecStatusParams periodic_params; + auto periodic_request = report_request(&state, false); + context->_append_external_file_commit_data(periodic_request, &periodic_params); + + EXPECT_FALSE(periodic_params.__isset.hive_partition_updates); + EXPECT_EQ(0, client->abort_count); + + TReportExecStatusParams final_params; + auto final_request = report_request(&state, true); + context->_append_external_file_commit_data(final_request, &final_params); + ASSERT_TRUE(final_params.__isset.hive_partition_updates); + ASSERT_EQ(1, final_params.hive_partition_updates.size()); + ASSERT_TRUE(final_params.hive_partition_updates[0].__isset.s3_mpu_pending_uploads); + ASSERT_EQ(1, final_params.hive_partition_updates[0].s3_mpu_pending_uploads.size()); + const auto& pending_upload = final_params.hive_partition_updates[0].s3_mpu_pending_uploads[0]; + EXPECT_EQ("bucket", pending_upload.bucket); + EXPECT_EQ("table/part.parquet", pending_upload.key); + EXPECT_EQ("upload-id", pending_upload.upload_id); + + // A definite coordinator rejection must consume the same cleanup owner that was retained + // while the periodic report deliberately withheld the pending-upload record. + context->_coordinator_callback(final_request); + + EXPECT_EQ(1, client->abort_count); + EXPECT_EQ("upload-id", client->aborted_upload_id); +} + +} // namespace doris diff --git a/be/test/io/client/s3_file_system_test.cpp b/be/test/io/client/s3_file_system_test.cpp index 1cc11876bde4c8..ce88c111789cf0 100644 --- a/be/test/io/client/s3_file_system_test.cpp +++ b/be/test/io/client/s3_file_system_test.cpp @@ -927,9 +927,8 @@ TEST_F(S3FileSystemTest, RateLimiterGetDownloadTest) { // Test: S3 rate limiter for PUT operations - multipart upload TEST_F(S3FileSystemTest, RateLimiterPutMultipartTest) { - // Skip if using Azure provider - Azure's create_multipart_upload is a no-op and doesn't - // consume rate limiter quota, while S3's CreateMultipartUpload does. This causes different - // failure timing that makes the test assertions invalid for Azure. + // This test asserts the S3 provider's exact multipart request/failure sequence; Azure uses + // lease coordination and therefore has different failure timing. if (config_->get_provider() == "AZURE") { GTEST_SKIP() << "This test relies on S3-specific multipart upload quota consumption " "behavior, not applicable for Azure"; @@ -1436,9 +1435,8 @@ TEST_F(S3FileSystemTest, RateLimiterGetDeleteDirectoryListTest) { // Test: S3 rate limiter for PUT operations - multipart upload with UploadPart failure TEST_F(S3FileSystemTest, RateLimiterPutMultipartUploadPartFailureTest) { - // Skip if using Azure provider - Azure's create_multipart_upload is a no-op and doesn't - // consume rate limiter quota, while S3's CreateMultipartUpload does. This causes different - // failure timing that makes the test assertions invalid for Azure. + // This test asserts the S3 provider's exact multipart request/failure sequence; Azure uses + // lease coordination and therefore has different failure timing. if (config_->get_provider() == "AZURE") { GTEST_SKIP() << "This test relies on S3-specific multipart upload quota consumption " "behavior, not applicable for Azure"; @@ -1865,8 +1863,7 @@ TEST_F(S3FileSystemTest, RateLimiterPutDeleteDirectoryDeleteObjectsTest) { // Test: S3 CreateMultipartUpload failure - simulates error when initiating multipart upload TEST_F(S3FileSystemTest, CreateMultipartUploadFailureTest) { - // Skip if using Azure provider - SyncPoint mechanism is S3-specific - // Also, Azure's create_multipart_upload is a no-op that always succeeds + // Skip if using Azure provider because the SyncPoint mechanism is S3-specific. if (config_->get_provider() == "AZURE") { GTEST_SKIP() << "This test uses S3-specific SyncPoint mechanism and multipart semantics, " "not applicable for Azure"; 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..dc7a2314dddb01 100644 --- a/be/test/io/fs/azure_obj_storage_client_test.cpp +++ b/be/test/io/fs/azure_obj_storage_client_test.cpp @@ -19,11 +19,15 @@ #include +#include + #include "io/fs/file_system.h" #include "io/fs/obj_storage_client.h" #include "util/s3_util.h" #ifdef USE_AZURE +#include + #include #include #include @@ -34,6 +38,18 @@ namespace doris { #ifdef USE_AZURE +TEST(AzureObjStorageClientMultipartHelperTest, fixed_length_namespace_requires_target_lease) { + EXPECT_EQ("p3w3DA==", io::azure_multipart_block_id("upload-a", 1)); + EXPECT_EQ("Sc7grw==", io::azure_multipart_block_id( + "doris-azure-lease-v1:09492e3d-e231-4ed9-bf84-b6fc772cda54", 1)); + EXPECT_EQ("Sc7grw==", io::azure_multipart_block_id( + "doris-azure-lease-v1:06996d15-1c2e-4ddd-8853-43816ea84a07", 1)); + EXPECT_EQ(io::azure_multipart_block_id("upload-a", 1).size(), + io::azure_multipart_block_id("upload-a", 999).size()); + EXPECT_EQ(4, Aws::Utils::HashingUtils::Base64Decode(io::azure_multipart_block_id("upload-a", 1)) + .GetLength()); +} + using namespace Azure::Storage::Blobs; TEST(AzureObjStorageClientTlsHelperTest, detects_tls_ca_error) { @@ -156,6 +172,88 @@ TEST_F(AzureObjStorageClientTest, delete_objects_recursively) { EXPECT_EQ(response.status.code, ErrorCode::OK); EXPECT_EQ(files.size(), 0); } + +TEST_F(AzureObjStorageClientTest, abort_multipart_upload_leaves_no_visible_blob) { + io::ObjectStoragePathOptions opts {.key = "AzureObjStorageClientTest/abort_multipart_upload"}; + 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.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 {.key = "AzureObjStorageClientTest/abort_preserves_put_blob"}; + auto put_response = AzureObjStorageClientTest::obj_storage_client->put_object(opts, "original"); + ASSERT_EQ(put_response.status.code, ErrorCode::OK); + 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.upload_id = create_response.upload_id; + + 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); +} + +TEST_F(AzureObjStorageClientTest, concurrent_multipart_uploads_do_not_share_staged_blocks) { + io::ObjectStoragePathOptions first {.key = "AzureObjStorageClientTest/concurrent_multipart"}; + io::ObjectStoragePathOptions second = first; + auto first_create = obj_storage_client->create_multipart_upload(first); + auto second_create = obj_storage_client->create_multipart_upload(second); + ASSERT_EQ(first_create.resp.status.code, ErrorCode::OK); + ASSERT_NE(second_create.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(first_create.upload_id.has_value()); + first.upload_id = first_create.upload_id; + + ASSERT_EQ(obj_storage_client->upload_part(first, "first", 1).resp.status.code, ErrorCode::OK); + ASSERT_EQ(obj_storage_client->complete_multipart_upload(first, {{.part_num = 1}}).status.code, + ErrorCode::OK); + + second_create = obj_storage_client->create_multipart_upload(second); + ASSERT_EQ(second_create.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(second_create.upload_id.has_value()); + second.upload_id = second_create.upload_id; + ASSERT_EQ(obj_storage_client->upload_part(second, "second", 1).resp.status.code, ErrorCode::OK); + ASSERT_EQ(obj_storage_client->complete_multipart_upload(second, {{.part_num = 1}}).status.code, + ErrorCode::OK); + + std::array contents {}; + size_t size_return = 0; + ASSERT_EQ(obj_storage_client + ->get_object(second, contents.data(), 0, contents.size(), &size_return) + .status.code, + ErrorCode::OK); + EXPECT_EQ(std::string_view(contents.data(), size_return), "second"); + EXPECT_EQ(obj_storage_client->delete_object(second).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..7b7686e187de56 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; @@ -455,7 +478,7 @@ TEST(RateLimitedObjStorageClientTest, recursive_delete_charges_one_put_qps) { EXPECT_EQ(4, fake->delete_objects_recursively_provider_calls); } -TEST(RateLimitedObjStorageClientTest, azure_noop_multipart_create_charges_one_put_qps) { +TEST(RateLimitedObjStorageClientTest, multipart_create_charges_one_logical_put_qps) { RateLimiterConfigGuard guard; config::enable_s3_rate_limiter = true; auto& manager = S3RateLimiterManager::instance(); @@ -463,20 +486,20 @@ TEST(RateLimitedObjStorageClientTest, azure_noop_multipart_create_charges_one_pu manager.bytes_limiter(S3RateLimitType::PUT)->reset(0, 0, 0); auto fake = std::make_shared(); - // Azure implements create_multipart_upload as a provider-side no-op. - fake->create_multipart_upload_provider_calls_per_logical_call = 0; + // Provider coordination may need multiple requests, but admission remains per logical API call. + fake->create_multipart_upload_provider_calls_per_logical_call = 2; RateLimitedObjStorageClient client(fake); ObjectStoragePathOptions opts {.bucket = "b", .key = "k"}; EXPECT_EQ(0, client.create_multipart_upload(opts).resp.status.code); EXPECT_EQ(1, fake->create_multipart_upload_calls); - EXPECT_EQ(0, fake->create_multipart_upload_provider_calls); + EXPECT_EQ(2, fake->create_multipart_upload_provider_calls); auto resp = client.create_multipart_upload(opts); EXPECT_EQ(ErrorCode::EXCEEDED_LIMIT, resp.resp.status.code); EXPECT_EQ(0, resp.resp.http_code); EXPECT_EQ(1, fake->create_multipart_upload_calls); - EXPECT_EQ(0, fake->create_multipart_upload_provider_calls); + EXPECT_EQ(2, fake->create_multipart_upload_provider_calls); } TEST(RateLimitedObjStorageClientTest, presigned_url_bypasses_rate_limiters) { diff --git a/be/test/io/fs/s3_file_writer_test.cpp b/be/test/io/fs/s3_file_writer_test.cpp index 3937d6e38561fe..c0c42239a1858a 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(); @@ -1294,8 +1319,9 @@ class SimpleMockObjStorageClient : public io::ObjStorageClient { * @return A tuple containing the mock S3 client and the S3FileWriter. */ std::tuple, std::shared_ptr> -create_s3_client(const std::string& path) { +create_s3_client(const std::string& path, bool used_by_s3_committer = false) { doris::io::FileWriterOptions opts; + opts.used_by_s3_committer = used_by_s3_committer; io::FileWriterPtr file_writer; auto st = s3_fs->create_file(path, &file_writer, &opts); EXPECT_TRUE(st.ok()) << st; @@ -1307,6 +1333,32 @@ 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()); +} + +TEST_F(S3FileWriterTest, failedReportCleanupAbortsDeferredProviderUploadAfterClose) { + auto [client, writer] = create_s3_client("deferred_report_rejected", true); + std::string data(config::s3_write_buffer_size, 'a'); + ASSERT_TRUE(writer->append(Slice(data)).ok()); + ASSERT_TRUE(writer->close().ok()); + ASSERT_EQ(FileWriter::State::CLOSED, writer->state()); + auto cleanup = writer->failed_report_cleanup(); + + cleanup(); + + EXPECT_EQ(1, client->abort_multipart_count); +} + /** * 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/memory/thread_mem_tracker_mgr_test.cpp b/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp index cbefa422ad2a77..d88d2b5ec4bc19 100644 --- a/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp +++ b/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp @@ -330,6 +330,39 @@ TEST_F(ThreadMemTrackerMgrTest, ReserveMemory) { EXPECT_EQ(doris::GlobalMemoryArbitrator::process_reserved_memory(), 0); } +TEST_F(ThreadMemTrackerMgrTest, TransfersReservationBetweenAsyncTasks) { + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "UT-TransferReservation"); + auto resource_context = ResourceContext::create_shared(); + resource_context->memory_context()->set_mem_tracker(tracker); + ThreadContext producer; + ThreadContext consumer; + producer.attach_task(resource_context); + consumer.attach_task(resource_context); + constexpr int64_t reservation = 4 * 1024 * 1024; + + ASSERT_TRUE(producer.thread_mem_tracker_mgr->try_reserve(reservation).ok()); + auto token = producer.thread_mem_tracker_mgr->take_reserved_memory(); + EXPECT_EQ(producer.thread_mem_tracker_mgr->reserved_mem(), 0); + EXPECT_EQ(token.bytes(), reservation); + + consumer.thread_mem_tracker_mgr->adopt_reserved_memory(std::move(token)); + EXPECT_EQ(consumer.thread_mem_tracker_mgr->reserved_mem(), reservation); + consumer.thread_mem_tracker_mgr->consume(reservation); + EXPECT_EQ(consumer.thread_mem_tracker_mgr->reserved_mem(), 0); + + ASSERT_TRUE(producer.thread_mem_tracker_mgr->try_reserve(reservation).ok()); + { + auto abandoned = producer.thread_mem_tracker_mgr->take_reserved_memory(); + EXPECT_EQ(abandoned.bytes(), reservation); + } + EXPECT_EQ(GlobalMemoryArbitrator::process_reserved_memory(), 0); + + producer.detach_task(); + consumer.detach_task(); + EXPECT_EQ(GlobalMemoryArbitrator::process_reserved_memory(), 0); +} + TEST_F(ThreadMemTrackerMgrTest, NestedReserveMemory) { std::unique_ptr thread_context = std::make_unique(); std::shared_ptr t = MemTrackerLimiter::create_shared( diff --git a/be/test/runtime/runtime_state_block_budget_test.cpp b/be/test/runtime/runtime_state_block_budget_test.cpp index 22ebc5ebf8a0ee..16a078bd64c8bf 100644 --- a/be/test/runtime/runtime_state_block_budget_test.cpp +++ b/be/test/runtime/runtime_state_block_budget_test.cpp @@ -18,12 +18,120 @@ #include #include "common/config.h" +#include "exec/pipeline/report_exec_status_size.h" #include "runtime/runtime_state.h" #include "testutil/mock/mock_runtime_state.h" #include "util/block_budget.h" 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_external_file_report_state(budget); + second.set_external_file_report_state(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()); +} + +TEST(RuntimeStateIcebergCommitDataTest, UsesTheSmallerCoordinatorThriftLimit) { + RuntimeState state; + const int32_t saved_limit = config::thrift_max_message_size; + config::thrift_max_message_size = 4 * 1024 * 1024; + state._query_options.__set_coordinator_thrift_max_message_size(1024 * 1024 + 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()); +} + +TEST(RuntimeStateIcebergCommitDataTest, ValidatesTheCompleteReportEnvelope) { + TReportExecStatusParams params; + params.__set_error_log({std::string(2 * 1024 * 1024, 'x')}); + + EXPECT_FALSE(validate_report_exec_status_size(params, 1024 * 1024).ok()); + EXPECT_TRUE(validate_report_exec_status_size(params, 3 * 1024 * 1024).ok()); +} + +TEST(RuntimeStateIcebergCommitDataTest, PeriodicReportOmitsExternalCommitData) { + RuntimeState state; + THivePartitionUpdate hive_update; + state.add_hive_partition_updates(hive_update); + TIcebergCommitData iceberg_data; + iceberg_data.__set_file_path("data.parquet"); + ASSERT_TRUE(state.add_iceberg_commit_datas(iceberg_data).ok()); + TMCCommitData mc_data; + state.add_mc_commit_datas(mc_data); + TReportExecStatusParams periodic_params; + + state.append_external_file_commit_data(&periodic_params, false); + + EXPECT_FALSE(periodic_params.__isset.hive_partition_updates); + EXPECT_FALSE(periodic_params.__isset.iceberg_commit_datas); + EXPECT_FALSE(periodic_params.__isset.mc_commit_datas); + + TReportExecStatusParams final_params; + state.append_external_file_commit_data(&final_params, true); + EXPECT_TRUE(final_params.__isset.hive_partition_updates); + EXPECT_TRUE(final_params.__isset.iceberg_commit_datas); + EXPECT_TRUE(final_params.__isset.mc_commit_datas); +} + +TEST(RuntimeStateIcebergCommitDataTest, RetainsFileCleanupUntilReportAcknowledgement) { + RuntimeState coordinator_state; + RuntimeState task_state; + auto report_state = std::make_shared(); + coordinator_state.set_external_file_report_state(report_state); + task_state.set_external_file_report_state(report_state); + int cleanup_count = 0; + task_state.add_rejected_external_file_report_cleanup([&] { ++cleanup_count; }); + + coordinator_state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::REJECTED); + coordinator_state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::REJECTED); + + EXPECT_EQ(1, cleanup_count); + + task_state.add_rejected_external_file_report_cleanup([&] { ++cleanup_count; }); + coordinator_state.finalize_external_file_report_cleanup( + ExternalFileReportOutcome::ACKNOWLEDGED); + EXPECT_EQ(1, cleanup_count); + + task_state.add_rejected_external_file_report_cleanup([&] { ++cleanup_count; }); + coordinator_state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::AMBIGUOUS); + EXPECT_EQ(1, cleanup_count); + coordinator_state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::REJECTED); + EXPECT_EQ(2, cleanup_count); +} + // --------------------------------------------------------------------------- // RuntimeState::batch_size() // --------------------------------------------------------------------------- diff --git a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorTransaction.java b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorTransaction.java index 978144329f8bde..f5958649219a4e 100644 --- a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorTransaction.java +++ b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveConnectorTransaction.java @@ -221,10 +221,22 @@ public void beginWrite(ConnectorSession session, String db, String tableName, Hi @Override public void commit() { - // The classification (finishInsertTable) ran from the executor in the legacy class; the unified SPI - // exposes only commit(), so it runs here (before the committer) to populate the action maps. If it - // throws, the committer was never created and the engine's subsequent rollback() cleans up. - finishInsertTable(nameMapping); + try { + // Object-store files remain unpublished until FE consumes one completion record per file. + validateObjectStoreCommitRecords(); + // The classification (finishInsertTable) ran from the executor in the legacy class; the unified + // SPI exposes only commit(), so it runs here to populate the action maps. + finishInsertTable(nameMapping); + } catch (Throwable t) { + // The transaction manager removes this connector before commit(), so this is the last owner + // capable of aborting deferred uploads when pre-commit validation or classification fails. + try { + rollback(); + } catch (Throwable cleanupFailure) { + t.addSuppressed(new Exception("Failed to clean up after pre-commit failure", cleanupFailure)); + } + throw t; + } hmsCommitter = new HmsCommitter(); try { for (Map.Entry> entry : tableActions.entrySet()) { @@ -282,6 +294,25 @@ public void commit() { } } + private void validateObjectStoreCommitRecords() { + if (fileType != TFileType.FILE_S3) { + return; + } + for (THivePartitionUpdate update : hivePartitionUpdates) { + int fileCount = update.getFileNames() == null ? 0 : update.getFileNames().size(); + List uploads = update.getS3MpuPendingUploads(); + int uploadCount = uploads == null ? 0 : uploads.size(); + boolean completeRecords = uploads != null && uploads.stream() + .allMatch(HiveConnectorTransaction::isCompleteObjectStoreUpload); + if (fileCount != uploadCount || (fileCount > 0 && !completeRecords)) { + throw new DorisConnectorException(String.format( + "Object-store write reported %d file(s) but %d valid multipart completion record(s); " + + "all backends must support deferred multipart completion before metadata commit", + fileCount, completeRecords ? uploadCount : 0)); + } + } + } + @Override public void rollback() { if (hmsCommitter == null) { @@ -454,15 +485,32 @@ void finishInsertTable(NameMapping nameMapping) { private void collectUncompletedMpuPendingUploads(List hivePartitionUpdates) { for (THivePartitionUpdate pu : hivePartitionUpdates) { - if (pu.getS3MpuPendingUploads() != null) { - for (TS3MPUPendingUpload s3MpuPendingUpload : pu.getS3MpuPendingUploads()) { - uncompletedMpuPendingUploads.add( - new UncompletedMpuPendingUpload(s3MpuPendingUpload, pu.getLocation().getWritePath())); + List uploads = pu.getS3MpuPendingUploads(); + if (uploads == null) { + continue; + } + String writePath = pu.getLocation() == null ? null : pu.getLocation().getWritePath(); + if (writePath == null || writePath.isEmpty()) { + // A malformed record must not prevent other valid uploads from being aborted. + LOG.warn("Skipping MPU cleanup record without a write path"); + continue; + } + for (TS3MPUPendingUpload upload : uploads) { + if (!isCompleteObjectStoreUpload(upload)) { + LOG.warn("Skipping incomplete MPU cleanup record for write path {}", writePath); + continue; } + uncompletedMpuPendingUploads.add(new UncompletedMpuPendingUpload(upload, writePath)); } } } + private static boolean isCompleteObjectStoreUpload(TS3MPUPendingUpload upload) { + return upload != null && upload.getUploadId() != null && !upload.getUploadId().isEmpty() + && upload.getBucket() != null && !upload.getBucket().isEmpty() + && upload.getKey() != null && !upload.getKey().isEmpty(); + } + private void convertToInsertExistingPartitionAction( NameMapping nameMapping, List> partitions) { diff --git a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWritePlanProvider.java b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWritePlanProvider.java index 703d3e47479b50..86423552fbe748 100644 --- a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWritePlanProvider.java +++ b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWritePlanProvider.java @@ -239,6 +239,8 @@ private THiveTableSink buildSink(ConnectorSession session, HiveTableHandle table // Hadoop config (BE-canonical static creds; hive has no vended overlay). tSink.setHadoopConfig(buildHadoopConfig()); + // New coordinators publish Azure's exact staged block IDs after BE writers finish. + tSink.setSupportsDeferredAzureMultipart(true); tSink.setOverwrite(handle.isOverwrite()); return tSink; } diff --git a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWriteUtils.java b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWriteUtils.java index 8e4c46dee4c92e..80c5199d7f491a 100644 --- a/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWriteUtils.java +++ b/fe/fe-connector/fe-connector-hive/src/main/java/org/apache/doris/connector/hive/HiveWriteUtils.java @@ -55,7 +55,11 @@ static List mergePartitions(List hiv THivePartitionUpdate old = merged.get(pu.getName()); old.setFileSize(old.getFileSize() + pu.getFileSize()); old.setRowCount(old.getRowCount() + pu.getRowCount()); - if (old.getS3MpuPendingUploads() != null && pu.getS3MpuPendingUploads() != null) { + if (pu.getS3MpuPendingUploads() != null && !pu.getS3MpuPendingUploads().isEmpty()) { + // A missing legacy list is empty state, not ownership of later completion records. + if (old.getS3MpuPendingUploads() == null) { + old.setS3MpuPendingUploads(new ArrayList<>()); + } old.getS3MpuPendingUploads().addAll(pu.getS3MpuPendingUploads()); } old.getFileNames().addAll(pu.getFileNames()); diff --git a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorTransactionTest.java b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorTransactionTest.java index 8e8436e2ec4ca8..7c5f018aa5436e 100644 --- a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorTransactionTest.java +++ b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveConnectorTransactionTest.java @@ -386,6 +386,69 @@ public void testCommitCompletesMultipartUploads() throws TException { "an unpartitioned INSERT_EXISTING must also update the table statistics; calls=" + client.calls); } + @Test + public void testCommitRejectsBaseBeObjectStoreUpdateWithoutPendingUpload() throws TException { + RecordingHmsClient client = new RecordingHmsClient(); + client.table = table(false, Collections.emptyMap()); + HiveConnectorTransaction txn = newTxn(client); + txn.beginWrite(null, DB, TBL, ctx(false)); + // The base Azure BE reports the file but omits this field because initiation returned no upload ID. + txn.addCommitData(serialize(pu("", TUpdateMode.APPEND, "s3://bucket/db/t", + Collections.singletonList("base-be-file"), 100, 4))); + + DorisConnectorException ex = Assertions.assertThrows(DorisConnectorException.class, txn::commit); + + Assertions.assertTrue(ex.getMessage().contains("multipart completion"), ex.getMessage()); + Assertions.assertFalse(client.calls.stream().anyMatch(c -> c.startsWith("updateTableStatistics")), + "metadata must remain unchanged when the object is still uncommitted"); + } + + @Test + public void testCommitValidationFailureAbortsEveryValidPendingUpload() throws TException { + RecordingHmsClient client = new RecordingHmsClient(); + client.table = table(false, Collections.emptyMap()); + RecordingObjStorage objStorage = new RecordingObjStorage(); + HiveConnectorTransaction txn = newTxnWithFs(client, new RecordingObjFileSystem(objStorage)); + txn.beginWrite(null, DB, TBL, ctx(false)); + txn.addCommitData(serialize(puWithMpu("", TUpdateMode.APPEND, "s3://bucket/db/t", + "bucket", "db/t/valid", "upload-valid", Collections.singletonMap(1, "etag-1")))); + THivePartitionUpdate malformed = puWithMpu("", TUpdateMode.APPEND, "s3://bucket/db/t", + "bucket", "db/t/malformed", "upload-malformed", Collections.singletonMap(1, "etag-1")); + malformed.unsetLocation(); + txn.addCommitData(serialize(malformed)); + txn.addCommitData(serialize(pu("", TUpdateMode.APPEND, "s3://bucket/db/t", + Collections.singletonList("missing-completion"), 100, 4))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + + Assertions.assertEquals(Collections.singletonList("abort:s3://bucket/db/t/valid:upload-valid"), + objStorage.calls, + "a pre-committer validation failure must abort every valid provider upload it rejects"); + Assertions.assertFalse(client.calls.stream().anyMatch(c -> c.startsWith("updateTableStatistics")), + "validation failure must not publish HMS metadata"); + } + + @Test + public void testCommitClassificationFailureAbortsPendingUpload() throws TException { + RecordingHmsClient client = new RecordingHmsClient(); + client.table = table(false, Collections.emptyMap()); + RecordingObjStorage objStorage = new RecordingObjStorage(); + HiveConnectorTransaction txn = newTxnWithFs(client, new RecordingObjFileSystem(objStorage)); + txn.beginWrite(null, DB, TBL, ctx(false)); + txn.addCommitData(serialize(puWithMpu("", TUpdateMode.NEW, "s3://bucket/db/t", + "bucket", "db/t/unclassified", "upload-unclassified", + Collections.singletonMap(1, "etag-1")))); + + Assertions.assertThrows(RuntimeException.class, txn::commit); + + Assertions.assertEquals( + Collections.singletonList("abort:s3://bucket/db/t/unclassified:upload-unclassified"), + objStorage.calls, + "a classification failure before HmsCommitter creation must abort its provider upload"); + Assertions.assertFalse(client.calls.stream().anyMatch(c -> c.startsWith("updateTableStatistics")), + "classification failure must not publish HMS metadata"); + } + @Test public void testRollbackAbortsPendingMultipartUploads() throws TException { // rollback() is NOT a no-op for hive (D9): data files are staged before commit, so a rollback must @@ -430,15 +493,17 @@ public void testCommitAddsNewPartitionOnce() throws TException { // GAP-7: the 20-at-a-time batching moved INTO ThriftHmsClient.addPartitions, so the committer must // call addPartitions ONCE with the whole list (not re-batch it). GAP-4: the new partition's storage // descriptor (values/location/columns) is rebuilt from the table at commit time. A genuinely-new - // partition takes the ADD path; on FILE_S3 the write path == target path, so no rename/MPU runs and - // the object-store FileSystem is never resolved (hence newTxn, not newTxnWithFs). + // partition takes the ADD path; on FILE_S3 the write path == target path, so FE completes the deferred + // multipart upload before adding HMS metadata. RecordingHmsClient client = new RecordingHmsClient(); client.table = table(true, Collections.emptyMap()); client.partitionExistsResult = false; - HiveConnectorTransaction txn = newTxn(client); + RecordingObjStorage objStorage = new RecordingObjStorage(); + HiveConnectorTransaction txn = newTxnWithFs(client, new RecordingObjFileSystem(objStorage)); txn.beginWrite(null, DB, TBL, ctx(false)); - txn.addCommitData(serialize(pu("dt=2024-01-01", TUpdateMode.NEW, "s3://bucket/db/t/dt=2024-01-01", - Collections.singletonList("f1"), 100, 4))); + txn.addCommitData(serialize(puWithMpu("dt=2024-01-01", TUpdateMode.NEW, + "s3://bucket/db/t/dt=2024-01-01", "bucket", "db/t/dt=2024-01-01/f1", + "upload-1", Collections.singletonMap(1, "etag-1")))); txn.commit(); txn.close(); diff --git a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWritePlanProviderTest.java b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWritePlanProviderTest.java index fcb889c6b2df10..b900e220733317 100644 --- a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWritePlanProviderTest.java +++ b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWritePlanProviderTest.java @@ -172,6 +172,17 @@ public void planWriteSetsBucketInfo() { Assertions.assertEquals(8, sink.getBucketInfo().getBucketCount()); } + @Test + public void planWriteAdvertisesDeferredAzureMultipartProtocol() { + RecordingHmsClient client = new RecordingHmsClient(); + client.table = tableBuilder().build(); + + THiveTableSink sink = planSink(client, new RecordingConnectorContext(), handle()); + + Assertions.assertTrue(sink.isSetSupportsDeferredAzureMultipart()); + Assertions.assertTrue(sink.isSupportsDeferredAzureMultipart()); + } + // ───────────────────────────── file format ───────────────────────────── @Test diff --git a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWriteUtilsTest.java b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWriteUtilsTest.java index e4633e6e032fbe..cad9fd6cc9a86a 100644 --- a/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWriteUtilsTest.java +++ b/fe/fe-connector/fe-connector-hive/src/test/java/org/apache/doris/connector/hive/HiveWriteUtilsTest.java @@ -25,6 +25,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -100,6 +101,20 @@ public void mergePartitionsToleratesNullPendingUploads() { Assertions.assertEquals(3L, merged.get(0).getFileSize()); } + @Test + public void mergePartitionsPreservesNewPendingUploadAfterLegacyUpdate() { + THivePartitionUpdate legacy = update("p=1", 1L, 1L, "legacy-file"); + THivePartitionUpdate current = update("p=1", 2L, 2L, "current-file"); + current.setS3MpuPendingUploads(new ArrayList<>(Collections.singletonList( + new TS3MPUPendingUpload().setUploadId("upload-1")))); + + THivePartitionUpdate merged = HiveWriteUtils.mergePartitions(Arrays.asList(legacy, current)).get(0); + + Assertions.assertNotNull(merged.getS3MpuPendingUploads()); + Assertions.assertEquals(1, merged.getS3MpuPendingUploads().size(), + "a legacy first update must not erase a later completion token"); + } + @Test public void isSubDirectoryHappyPath() { Assertions.assertTrue(HiveWriteUtils.isSubDirectory("/warehouse/table", "/warehouse/table/p=1")); 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 3b32742842c27d..4406c6ad9bf82b 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 4c3f687c5f5bf4..12600c78470280 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 @@ -154,8 +154,8 @@ public class IcebergConnectorTransaction implements ConnectorTransaction, Rewrit private Map staticPartitionValues = Collections.emptyMap(); private String branchName; private IcebergWriteSchemaContext writeSchemaContext; - // 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; @@ -394,12 +394,13 @@ private void applyBeginGuards(IcebergWriteContext ctx, String tableName) { // scan used, S_read), threaded onto the write handle and carried on the ctx. The commit-time // removeDeletes (option D) re-derives from baseSnapshotId, and BE unions the scan-time (S_read) // old deletes into the new DV — anchoring both at S_read keeps supply and remove on one snapshot - // (no resurrection under a concurrent commit in the read->begin-write window). A -1 readSnapshotId - // (no pin: a caller without the threaded handle) falls back to the begin-time current snapshot. + // (no resurrection under a concurrent commit in the read->begin-write window). An explicitly pinned + // -1 is the empty-table generation and must remain an OCC fence; only an unpinned caller may fall + // back to the begin-time current snapshot. long pinnedReadSnapshot = ctx.getReadSnapshotId(); // Keep both ternary arms boxed (Long): getSnapshotIdIfPresent returns null for an empty table // (no snapshot), and a primitive arm would force-unbox that null into an NPE. - this.baseSnapshotId = pinnedReadSnapshot >= 0 + this.baseSnapshotId = ctx.isReadSnapshotPinned() ? Long.valueOf(pinnedReadSnapshot) : getSnapshotIdIfPresent(table); if (table instanceof HasTableOperations) { int formatVersion = ((HasTableOperations) table).operations().current().formatVersion(); @@ -410,7 +411,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); @@ -420,10 +420,51 @@ 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"); } + // Keep the empty generation pinned across Transactions.newTransaction(), whose refresh may + // otherwise adopt a first snapshot committed between the guard and transaction creation. + return readSnapshotId; + } + 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 @@ -685,6 +726,8 @@ private void commitReplaceTxn(List pendingResults) { if (branchName != null) { overwriteFiles = overwriteFiles.toBranch(branchName); } + // 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); @@ -704,6 +747,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."); @@ -729,6 +777,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, @@ -738,6 +787,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 536babad0e517c..a75beb97ad12a5 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 024a036df11158..318761a52ccbde 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 @@ -437,6 +437,9 @@ public List planScan(ConnectorSession session, ConnectorScan public long streamingSplitEstimate(ConnectorSession session, ConnectorTableHandle handle, Optional filter, boolean countPushdown) { IcebergTableHandle iceHandle = (IcebergTableHandle) handle; + if (isExplicitEmptySnapshot(iceHandle)) { + return -1; + } if (iceHandle.isSystemTable() || !sessionBool(session, ENABLE_EXTERNAL_TABLE_BATCH_MODE, true)) { return -1; } @@ -613,6 +616,11 @@ private List planScanInternal( Optional filter, boolean countPushdown) { IcebergTableHandle iceHandle = (IcebergTableHandle) handle; + if (isExplicitEmptySnapshot(iceHandle)) { + // The absence of a snapshot is part of the statement's MVCC state; consulting newScan() here + // would silently replace it with a snapshot committed after statement start. + return Collections.emptyList(); + } if (iceHandle.isSystemTable()) { // System tables take a metadata-table path, never the data-file path below (no count pushdown, no // data-file ranges) — mirrors legacy IcebergScanNode branching on isSystemTable. $position_deletes @@ -904,7 +912,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 { @@ -1126,6 +1134,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 8ce93420fdc6c1..89653e97743374 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 @@ -43,26 +43,35 @@ final class IcebergWriteContext { private final Optional branchName; private final long readSnapshotId; private final IcebergWriteSchemaContext writeSchemaContext; + 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, null); } IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, Map staticPartitionValues, Optional branchName, long readSnapshotId) { - 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.writeSchemaContext = null; + this(writeOperation, overwrite, staticPartitionValues, branchName, readSnapshotId, true, null); + } + + IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, + Map staticPartitionValues, Optional branchName, long readSnapshotId, + boolean readSnapshotPinned) { + this(writeOperation, overwrite, staticPartitionValues, branchName, + readSnapshotId, readSnapshotPinned, null); } IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, Map staticPartitionValues, Optional branchName, long readSnapshotId, IcebergWriteSchemaContext writeSchemaContext) { + this(writeOperation, overwrite, staticPartitionValues, branchName, + readSnapshotId, false, writeSchemaContext); + } + + IcebergWriteContext(WriteOperation writeOperation, boolean overwrite, + Map staticPartitionValues, Optional branchName, long readSnapshotId, + boolean readSnapshotPinned, IcebergWriteSchemaContext writeSchemaContext) { this.writeOperation = writeOperation; this.overwrite = overwrite; this.staticPartitionValues = staticPartitionValues == null @@ -70,6 +79,7 @@ final class IcebergWriteContext { this.branchName = branchName == null ? Optional.empty() : branchName; this.readSnapshotId = readSnapshotId; this.writeSchemaContext = writeSchemaContext; + this.readSnapshotPinned = readSnapshotPinned; } WriteOperation getWriteOperation() { @@ -95,7 +105,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. @@ -107,4 +118,8 @@ long getReadSnapshotId() { IcebergWriteSchemaContext getWriteSchemaContext() { return writeSchemaContext; } + + 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 6594c2589d85c0..fdbce14de8f2fb 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 @@ -390,13 +390,15 @@ private IcebergWriteContext buildWriteContext( // 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, schemaContext); + handle.getBranchName(), readSnapshotId, readSnapshotPinned, schemaContext); } 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 37d949c1d7abcf..7b785c9ac31674 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 to scan for orphan files", + null, ArgumentParsers.nonEmptyString(LOCATION)); + namedArguments.registerOptionalArgument(DRY_RUN, "Only count orphan files", true, + ArgumentParsers.booleanValue(DRY_RUN)); + namedArguments.registerOptionalArgument(ALLOW_UNSAFE_LOCATION, + "Allow an explicitly supplied location whose table ownership cannot be proved", + false, ArgumentParsers.booleanValue(ALLOW_UNSAFE_LOCATION)); + } + + @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"); + } + long olderThan = namedArguments.getLong(OLDER_THAN); + // Reject an unsafe cutoff before opening any metadata or manifest file. + if (olderThan > System.currentTimeMillis() - MIN_RETENTION_MS) { + throw new DorisConnectorException("older_than must retain at least 24 hours of files"); + } + List scanScopes = resolveScanScopes(table); + + try { + ReachableIndex reachable = collectReachableFiles(table); + long orphanCount = 0; + long deletedCount = 0; + boolean dryRun = namedArguments.getBoolean(DRY_RUN); + for (ScanScope scope : scanScopes) { + // Object stores use raw prefix matching, so the separator excludes sibling prefixes. + String listingPrefix = scope.root.endsWith("/") ? scope.root : scope.root + "/"; + for (FileInfo file : ((SupportsPrefixOperations) table.io()).listPrefix(listingPrefix)) { + if (scope.owns(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 resolveScanScopes(Table table) { + String tableRoot = normalizeLocation(table.location()); + String requested = namedArguments.getString(LOCATION); + if (requested != null) { + String normalized = normalizeLocation(requested); + if (isWithin(normalized, tableRoot)) { + return Lists.newArrayList(ScanScope.exclusive(normalized)); + } + if (!namedArguments.getBoolean(ALLOW_UNSAFE_LOCATION)) { + throw new DorisConnectorException( + "Cannot prove that location is owned by this table; set allow_unsafe_location=true " + + "only after verifying the prefix is exclusive to the table"); + } + // This explicit escape hatch also covers historical roots after a table-location migration. + return Lists.newArrayList(ScanScope.exclusive(normalized)); + } + if (nonEmpty(table.properties().get(TableProperties.WRITE_LOCATION_PROVIDER_IMPL)) != null) { + throw new DorisConnectorException( + "remove_orphan_files cannot infer ownership for a custom write.location-provider.impl; " + + "provide location with allow_unsafe_location=true after verifying exclusivity"); + } + String metadataRoot = nonEmpty(table.properties().get(TableProperties.WRITE_METADATA_LOCATION)); + if (metadataRoot != null && !isWithin(normalizeLocation(metadataRoot), tableRoot)) { + throw new DorisConnectorException( + "Cannot prove that the configured external metadata location is table-exclusive; " + + "provide location with allow_unsafe_location=true after verifying exclusivity"); + } + List scopes = new ArrayList<>(); + scopes.add(ScanScope.exclusive(tableRoot)); + if (Boolean.parseBoolean(table.properties().get(TableProperties.OBJECT_STORE_ENABLED))) { + // Match Iceberg's ObjectStoreLocationProvider precedence exactly. + String objectRoot = nonEmpty(table.properties().get(TableProperties.WRITE_DATA_LOCATION)); + if (objectRoot == null) { + objectRoot = nonEmpty(table.properties().get(TableProperties.OBJECT_STORE_PATH)); + } + if (objectRoot == null) { + objectRoot = nonEmpty(table.properties().get(TableProperties.WRITE_FOLDER_STORAGE_LOCATION)); + } + if (objectRoot != null) { + String normalizedObjectRoot = normalizeLocation(objectRoot); + if (!isWithin(normalizedObjectRoot, tableRoot)) { + // Iceberg's hashed path retains only a suffix of the table location; that suffix is not + // a globally unique ownership key when multiple catalogs share an object-store root. + throw new DorisConnectorException( + "Cannot prove that the configured object-store root is table-exclusive; " + + "provide location with allow_unsafe_location=true after verifying exclusivity"); + } + } + } else { + String externalDataRoot = nonEmpty(table.properties().get(TableProperties.WRITE_DATA_LOCATION)); + if (externalDataRoot == null) { + externalDataRoot = nonEmpty( + table.properties().get(TableProperties.WRITE_FOLDER_STORAGE_LOCATION)); + } + if (externalDataRoot != null && !isWithin(normalizeLocation(externalDataRoot), tableRoot)) { + throw new DorisConnectorException( + "Cannot prove that the configured external data location is table-exclusive; " + + "provide location with allow_unsafe_location=true after verifying exclusivity"); + } + } + return scopes; + } + + private static String nonEmpty(String location) { + return location == null || location.isEmpty() ? null : location; + } + + private boolean isWithin(String location, String root) { + return isWithinLocation(location, root); + } + + private static boolean isWithinLocation(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 ReachableIndex collectReachableFiles(Table table) throws IOException { + ReachableIndex reachable = new ReachableIndex(MAX_REACHABLE_FILES); + reachable.addAll(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); + FileIdentity retainedIdentity = reachable.byPath.get(candidateIdentity.path); + if (candidateIdentity.equals(retainedIdentity)) { + return true; + } + if (retainedIdentity != null) { + // 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)); + } + + 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); + } + } + + static void verifyReachableIndexLimit(Set locations, int maxEntries) { + ReachableIndex index = new ReachableIndex(maxEntries); + index.addAll(locations); + } + + private static final class ReachableIndex { + private final Map byPath = new LinkedHashMap<>(); + private final int maxEntries; + + private ReachableIndex(int maxEntries) { + this.maxEntries = maxEntries; + } + + private void addAll(Iterable locations) { + locations.forEach(this::add); + } + + private void add(String location) { + FileIdentity identity = FileIdentity.of(location); + FileIdentity existing = byPath.putIfAbsent(identity.path, identity); + if (existing != null && !existing.equals(identity)) { + throw new DorisConnectorException( + "Cannot determine whether reachable file locations are equivalent"); + } + if (existing == null && byPath.size() > maxEntries) { + throw new DorisConnectorException( + "Reachable file index exceeds the safe in-memory limit of " + maxEntries); + } + } + } + + private static final class ScanScope { + private final String root; + + private ScanScope(String root) { + this.root = root; + } + + private static ScanScope exclusive(String root) { + return new ScanScope(root); + } + + private boolean owns(String candidate) { + return isWithinLocation(candidate, root); + } + } + + private static 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 a1afce8ed8111d..b1c949685edc81 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 5f4c2d1c0713b8..7ded2a2e77e42e 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 @@ -32,6 +32,7 @@ import org.apache.doris.thrift.TIcebergColumnStats; import org.apache.doris.thrift.TIcebergCommitData; +import org.apache.iceberg.BaseTable; import org.apache.iceberg.DataFile; import org.apache.iceberg.DataFiles; import org.apache.iceberg.DeleteFile; @@ -39,10 +40,12 @@ import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileMetadata; import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.HasTableOperations; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.Snapshot; import org.apache.iceberg.Table; +import org.apache.iceberg.Transaction; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.expressions.Expression; @@ -131,6 +134,11 @@ private static IcebergWriteContext overwriteToBranch(String branch) { WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.of(branch)); } + private static IcebergWriteContext overwriteCtxPinned(long readSnapshotId) { + return new IcebergWriteContext(WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.empty(), + readSnapshotId); + } + private static IcebergWriteContext overwriteStaticCtx(Table table, Map staticValues) { IcebergWriteSchemaContext schemaContext = IcebergWriteSchemaContext.create(table, table.name(), Optional.empty(), false, false); @@ -510,6 +518,22 @@ public void beginMergeHonorsPinnedReadSnapshotOverCurrent() { "MERGE must anchor baseSnapshotId at the pinned read snapshot, not the current snapshot"); } + @Test + public void beginDeletePreservesExplicitEmptyReadSnapshot() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table table = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), props("format-version", "2")); + table.newAppend().appendFile( + dataFile(table.spec(), "s3://b/db1/t1/concurrent.parquet", 1L)).commit(); + + IcebergConnectorTransaction txn = txnFor( + opsReturning(catalog.loadTable(id)), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", deleteCtxPinned(-1L)); + + Assertions.assertEquals(Long.valueOf(-1L), txn.getBaseSnapshotId(), + "an explicit empty read is an OCC fence, not an absent pin"); + } + @Test public void beginInsertDoesNotCaptureBaseSnapshotId() { InMemoryCatalog catalog = freshCatalog(); @@ -585,6 +609,138 @@ 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)); + Assertions.assertEquals(-1L, txn.getBaseSnapshotId(), + "the empty-read generation must remain the transaction OCC fence"); + + table.newAppend().appendFile(dataFile(table.spec(), + "s3://b/db1/t1/after-begin.parquet", 1L)).commit(); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit); + } + + @Test + public void overwriteRejectsFirstSnapshotCommittedDuringTransactionRefresh() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table loaded = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), + props("write.format.default", "parquet")); + Table racing = new BaseTable(((HasTableOperations) loaded).operations(), loaded.name()) { + private boolean injected; + + @Override + public Transaction newTransaction() { + if (!injected) { + injected = true; + Table concurrent = catalog.loadTable(id); + concurrent.newAppend().appendFile(dataFile(concurrent.spec(), + "s3://b/db1/t1/during-refresh.parquet", 1L)).commit(); + } + return super.newTransaction(); + } + }; + IcebergConnectorTransaction txn = txnFor( + opsReturning(racing), new RecordingConnectorContext()); + + txn.beginWrite(SESSION, "db1", "t1", overwriteCtxPinned(-1L)); + txn.addCommitData(commitBytes(dataFileItem( + "s3://b/db1/t1/replacement.parquet", 1L, 1024L, Collections.emptyList()))); + + 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(); @@ -635,6 +791,24 @@ public void overwriteEmptyUnpartitionedBranchClearsOnlyBranchFiles() { } } + @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(); @@ -655,6 +829,53 @@ 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(table, 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(); @@ -758,6 +979,23 @@ public void deleteDetectsConcurrentDataFileConflict() { "a concurrent data-file append since the base snapshot must be detected as a conflict"); } + @Test + public void deleteFromExplicitEmptySnapshotDetectsFirstConcurrentCommit() { + InMemoryCatalog catalog = freshCatalog(); + TableIdentifier id = TableIdentifier.of("db1", "t1"); + Table table = catalog.createTable(id, SCHEMA, PartitionSpec.unpartitioned(), props("format-version", "2")); + IcebergConnectorTransaction txn = txnFor(opsReturning(table), new RecordingConnectorContext()); + txn.beginWrite(SESSION, "db1", "t1", deleteCtxPinned(-1L)); + + catalog.loadTable(id).newAppend().appendFile( + dataFile(table.spec(), "s3://b/db1/t1/concurrent.parquet", 7L)).commit(); + txn.addCommitData(commitBytes(positionDeleteItem( + "s3://b/db1/t1/del.parquet", 1L, "s3://b/db1/t1/concurrent.parquet"))); + + Assertions.assertThrows(DorisConnectorException.class, txn::commit, + "validateFromSnapshot(-1) must reject the first snapshot committed after an empty read"); + } + @Test public void deletePassesValidationSuiteWhenNoConcurrentChange() { 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 2ab67da81fd841..3d996d4926e93b 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/IcebergScanPlanProviderTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java index 426f89c7d427b6..db175d9f0a1541 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergScanPlanProviderTest.java @@ -318,6 +318,23 @@ public void planScanEnumeratesOneRangePerDataFile() { Assertions.assertEquals(2048L, ranges.get(1).getLength()); } + @Test + public void explicitEmptySnapshotDoesNotDriftToFirstConcurrentCommit() { + Table table = createTable("t1", SCHEMA, PartitionSpec.unpartitioned()); + IcebergTableHandle emptyPin = new IcebergTableHandle("db1", "t1") + .withSnapshot(-1L, null, -1L); + table.newAppend().appendFile( + dataFile(table.spec(), "s3://b/db/t1/concurrent.parquet", 1024, null, null)).commit(); + IcebergScanPlanProvider provider = new IcebergScanPlanProvider( + Collections.emptyMap(), opsReturning(table)); + + List ranges = provider.planScan(null, + ConnectorScanRequest.builder(emptyPin, Collections.emptyList()).build()); + + Assertions.assertTrue(ranges.isEmpty(), + "an explicit empty MVCC pin must not be replaced by the table's first snapshot"); + } + @Test public void planScanRewriteFileScopeKeepsOnlyRawScopedFiles() { Table table = createTable("t1", SCHEMA, PartitionSpec.unpartitioned()); 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 d1978fe7741fd7..76d4f34b57e2d1 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.spi.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..bce05901cead38 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergRemoveOrphanFilesActionTest.java @@ -0,0 +1,427 @@ +// 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.spi.DorisConnectorException; +import org.apache.doris.connector.spi.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()); + Set manifestPaths = new HashSet<>(); + table.snapshots().forEach(snapshot -> snapshot.allManifests(table.io()) + .forEach(manifest -> manifestPaths.add(manifest.path()))); + RecordingFileIO recordingFileIO = new RecordingFileIO(table.io(), manifestPaths); + Table recordingTable = new BaseTable( + new StaticTableOperations(((HasTableOperations) table).operations().current(), recordingFileIO), + table.name()); + Path uncommitted = createOldFile(temp.resolve("table/data/uncommitted.parquet")); + IcebergRemoveOrphanFilesAction action = action(System.currentTimeMillis(), false); + action.validate(); + + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(recordingTable, ActionTestTables.session("UTC"))); + Assertions.assertEquals(0, recordingFileIO.manifestOpenCount()); + 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 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 rejectsUnprovenExternalDataRootByDefault(@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(); + + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(table, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(orphan)); + } + + @Test + public void guardedExplicitDataRootDeletesButUnguardedArbitraryRootIsRejected(@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(), true); + 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 guardedExplicitLocationCoversFormerTableRootAfterMigration(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("current-table"), Collections.emptyMap()); + Path sharedFormerRoot = temp.resolve("former-root-now-shared"); + Path formerDataRoot = sharedFormerRoot.resolve("old-table-data"); + Path formerMetadataRoot = sharedFormerRoot.resolve("old-table-metadata"); + Path dataOrphan = createOldFile(formerDataRoot.resolve("orphan.parquet")); + Path metadataOrphan = createOldFile(formerMetadataRoot.resolve("orphan.metadata.json")); + Path neighborFile = createOldFile(sharedFormerRoot.resolve("neighbor-table/live.parquet")); + + IcebergRemoveOrphanFilesAction unguarded = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + sharedFormerRoot.toUri().toString()); + unguarded.validate(); + Assertions.assertThrows(DorisConnectorException.class, + () -> unguarded.execute(table, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(dataOrphan)); + Assertions.assertTrue(Files.exists(metadataOrphan)); + Assertions.assertTrue(Files.exists(neighborFile)); + + IcebergRemoveOrphanFilesAction dataAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + formerDataRoot.toUri().toString(), true); + dataAction.validate(); + dataAction.execute(table, ActionTestTables.session("UTC")); + IcebergRemoveOrphanFilesAction metadataAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + formerMetadataRoot.toUri().toString(), true); + metadataAction.validate(); + metadataAction.execute(table, ActionTestTables.session("UTC")); + + Assertions.assertFalse(Files.exists(dataOrphan)); + Assertions.assertFalse(Files.exists(metadataOrphan)); + Assertions.assertTrue(Files.exists(neighborFile)); + } + + @Test + public void objectStoreOwnershipExcludesNeighborTableAndFolderRootFailsClosed(@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.WRITE_DATA_LOCATION, objectRoot.toUri().toString()); + objectProperties.put(TableProperties.OBJECT_STORE_PATH, + temp.resolve("lower-precedence-object-path").toUri().toString()); + Files.createDirectories(objectRoot); + Table objectTable = createTable(temp.resolve("object-metadata"), objectProperties); + String ownLocation = objectTable.locationProvider().newDataLocation("own.parquet"); + Path ownOrphan = createOldFile(Path.of(java.net.URI.create(ownLocation))); + Table neighborTable = createTable(temp.resolve("neighbor"), objectProperties); + String neighborLocation = neighborTable.locationProvider().newDataLocation("live.parquet"); + Path neighborFile = createOldFile(Path.of(java.net.URI.create(neighborLocation))); + 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(); + Assertions.assertThrows(DorisConnectorException.class, + () -> objectAction.execute(objectTable, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(ownOrphan)); + Assertions.assertTrue(Files.exists(neighborFile)); + + IcebergRemoveOrphanFilesAction guardedObjectAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + ownOrphan.getParent().toUri().toString(), true); + guardedObjectAction.validate(); + guardedObjectAction.execute(objectTable, ActionTestTables.session("UTC")); + Assertions.assertFalse(Files.exists(ownOrphan)); + Assertions.assertTrue(Files.exists(neighborFile)); + + Map prefixCollisionProperties = new HashMap<>(); + prefixCollisionProperties.put(TableProperties.OBJECT_STORE_ENABLED, "true"); + prefixCollisionProperties.put(TableProperties.WRITE_DATA_LOCATION, + temp.resolve("prefix-table-shared").toUri().toString()); + Table prefixCollisionTable = createTable(temp.resolve("prefix-table"), prefixCollisionProperties); + Assertions.assertThrows(DorisConnectorException.class, + () -> objectAction.execute(prefixCollisionTable, ActionTestTables.session("UTC"))); + + IcebergRemoveOrphanFilesAction folderAction = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + folderAction.validate(); + Assertions.assertThrows(DorisConnectorException.class, + () -> folderAction.execute(folderTable, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(folderOrphan)); + } + + @Test + public void sharedObjectRootWithSameTableSuffixFailsClosed(@TempDir Path temp) throws Exception { + Path sharedRoot = temp.resolve("shared-object-root"); + Map properties = new HashMap<>(); + properties.put(TableProperties.OBJECT_STORE_ENABLED, "true"); + properties.put(TableProperties.WRITE_DATA_LOCATION, sharedRoot.toUri().toString()); + Table first = createTable(temp.resolve("catalog-a/db/t"), properties); + Table second = createTable(temp.resolve("catalog-b/db/t"), properties); + Path firstFile = createOldFile(Path.of(java.net.URI.create( + first.locationProvider().newDataLocation("first.parquet")))); + Path secondFile = createOldFile(Path.of(java.net.URI.create( + second.locationProvider().newDataLocation("second.parquet")))); + + IcebergRemoveOrphanFilesAction action = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + action.validate(); + + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(first, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(firstFile)); + Assertions.assertTrue(Files.exists(secondFile)); + } + + @Test + public void reachableIndexHasExplicitSafetyCap() { + Set files = Set.of("s3://bucket/table/a", "s3://bucket/table/b", "s3://bucket/table/c"); + Assertions.assertThrows(DorisConnectorException.class, + () -> IcebergRemoveOrphanFilesAction.verifyReachableIndexLimit(files, 2)); + } + + @Test + public void customLocationProviderRequiresGuardedExplicitLocation(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("table"), Collections.singletonMap( + TableProperties.WRITE_LOCATION_PROVIDER_IMPL, "example.CustomProvider")); + Path providerRoot = temp.resolve("provider-data"); + Path orphan = createOldFile(providerRoot.resolve("orphan.parquet")); + IcebergRemoveOrphanFilesAction action = action( + System.currentTimeMillis() - MIN_RETENTION_MS, true); + action.validate(); + Assertions.assertThrows(DorisConnectorException.class, + () -> action.execute(table, ActionTestTables.session("UTC"))); + Assertions.assertTrue(Files.exists(orphan)); + + IcebergRemoveOrphanFilesAction guarded = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false, + providerRoot.toUri().toString(), true); + guarded.validate(); + guarded.execute(table, ActionTestTables.session("UTC")); + Assertions.assertFalse(Files.exists(orphan)); + } + + private static IcebergRemoveOrphanFilesAction action(long olderThan, boolean dryRun) { + return action(olderThan, dryRun, null); + } + + private static IcebergRemoveOrphanFilesAction action(long olderThan, boolean dryRun, String location) { + return action(olderThan, dryRun, location, false); + } + + private static IcebergRemoveOrphanFilesAction action( + long olderThan, boolean dryRun, String location, boolean allowUnsafeLocation) { + Map properties = new HashMap<>(); + properties.put(IcebergRemoveOrphanFilesAction.OLDER_THAN, String.valueOf(olderThan)); + properties.put(IcebergRemoveOrphanFilesAction.DRY_RUN, String.valueOf(dryRun)); + properties.put(IcebergRemoveOrphanFilesAction.ALLOW_UNSAFE_LOCATION, + String.valueOf(allowUnsafeLocation)); + 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/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java index 5af8be9c147bbf..40f6e263ebc446 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java @@ -54,7 +54,6 @@ import org.apache.doris.datasource.doris.RemoteDorisExternalTable; import org.apache.doris.datasource.doris.RemoteOlapTable; import org.apache.doris.datasource.doris.source.RemoteDorisScanNode; -import org.apache.doris.datasource.mvcc.MvccUtil; import org.apache.doris.datasource.plugin.PluginDrivenExternalCatalog; import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; import org.apache.doris.datasource.plugin.PluginDrivenMetadata; @@ -646,13 +645,11 @@ private PluginDrivenTableSink buildPluginRowLevelDmlSink( "Connector '" + catalog.getName() + "' (type: " + catalog.getType() + ") does not support row-level DML operations"); } - providerTableHandle = PluginDrivenScanNode.applyMvccSnapshotPin( - metadata, connSession, providerTableHandle, MvccUtil.getSnapshotFromContext(targetTable)); - // writeSortInfo == null: a row-level DML has no engine-resolved write sort (MERGE's sort lives in the // connector's TIcebergMergeSink.sort_fields, DELETE is unsorted). return new PluginDrivenTableSink(targetTable, writePlanProvider, connSession, - providerTableHandle, connectorColumns, null, writeOperation, requireMergeCardinalityCheck); + providerTableHandle, connectorColumns, null, writeOperation, + requireMergeCardinalityCheck, metadata); } @Override @@ -702,16 +699,6 @@ public PlanFragment visitPhysicalConnectorTableSink( + ") does not support INSERT operations"); } - // Thread the statement's MVCC snapshot pin onto the WRITE handle, reusing the exact scan-side pin - // logic so a DML's write anchors at the SAME snapshot its scan read (the pin is keyed by - // catalog/db/table in StatementContext, so the write target resolves the scan's pin). WHY: an MVCC - // connector's RowDelta DELETE/MERGE re-derives the deletes to remove from the write's base snapshot, - // while BE unions the scan-time deletes into the new DV — pinning both at the read snapshot keeps - // them on one snapshot ([SHOULD-2] / Fix B). A no-op for non-MVCC tables (jdbc/maxcompute) and any - // connector whose handle is not snapshot-pinned, so it is byte-identical for every current write path. - providerTableHandle = PluginDrivenScanNode.applyMvccSnapshotPin( - metadata, connSession, providerTableHandle, MvccUtil.getSnapshotFromContext(targetTable)); - // The connector declares its write-sort columns (e.g. an iceberg WRITE ORDERED BY) as positions // into the sink's full-schema output; the engine resolves them to bound slots and builds the // TSortInfo here (the connector's planWrite has no bound exprs). Empty for connectors with no @@ -728,7 +715,7 @@ public PlanFragment visitPhysicalConnectorTableSink( ? WriteOperation.REWRITE : WriteOperation.INSERT; PluginDrivenTableSink providerSink = new PluginDrivenTableSink(targetTable, writePlanProvider, connSession, providerTableHandle, connectorColumns, writeSortInfo, - writeOperation); + writeOperation, false, metadata); rootFragment.setSink(providerSink); return rootFragment; } 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 c9e5f60a71a2c3..c24931a1a55ba9 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/main/java/org/apache/doris/planner/PluginDrivenTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/planner/PluginDrivenTableSink.java index e9e539ea018e2f..691e0942995e60 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/PluginDrivenTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/PluginDrivenTableSink.java @@ -17,15 +17,19 @@ package org.apache.doris.planner; +import org.apache.doris.analysis.TableScanParams; import org.apache.doris.common.AnalysisException; import org.apache.doris.connector.spi.ConnectorColumn; +import org.apache.doris.connector.spi.ConnectorMetadata; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.write.ConnectorSinkPlan; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; +import org.apache.doris.datasource.mvcc.MvccUtil; import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; +import org.apache.doris.datasource.scan.PluginDrivenScanNode; import org.apache.doris.nereids.trees.plans.commands.insert.InsertCommandContext; import org.apache.doris.nereids.trees.plans.commands.insert.PluginDrivenInsertCommandContext; import org.apache.doris.thrift.TDataSink; @@ -57,6 +61,7 @@ public class PluginDrivenTableSink extends BaseExternalTableDataSink { private final ConnectorWritePlanProvider writePlanProvider; private final ConnectorSession connectorSession; private final ConnectorTableHandle tableHandle; + private final ConnectorMetadata connectorMetadata; private final List connectorColumns; // The engine-built BE sort instruction for a connector that declares write-sort columns (iceberg // WRITE ORDERED BY); null when the target needs no write sort. The connector cannot build it (the @@ -109,7 +114,7 @@ public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, ConnectorTableHandle tableHandle, List connectorColumns, TSortInfo writeSortInfo, WriteOperation writeOperation) { this(targetTable, writePlanProvider, connectorSession, tableHandle, connectorColumns, - writeSortInfo, writeOperation, false); + writeSortInfo, writeOperation, false, null); } /** @@ -121,11 +126,21 @@ public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, ConnectorWritePlanProvider writePlanProvider, ConnectorSession connectorSession, ConnectorTableHandle tableHandle, List connectorColumns, TSortInfo writeSortInfo, WriteOperation writeOperation, boolean requireMergeCardinalityCheck) { + this(targetTable, writePlanProvider, connectorSession, tableHandle, connectorColumns, + writeSortInfo, writeOperation, requireMergeCardinalityCheck, null); + } + + public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, + ConnectorWritePlanProvider writePlanProvider, ConnectorSession connectorSession, + ConnectorTableHandle tableHandle, List connectorColumns, + TSortInfo writeSortInfo, WriteOperation writeOperation, boolean requireMergeCardinalityCheck, + ConnectorMetadata connectorMetadata) { super(); this.targetTable = targetTable; this.writePlanProvider = writePlanProvider; this.connectorSession = connectorSession; this.tableHandle = tableHandle; + this.connectorMetadata = connectorMetadata; this.connectorColumns = connectorColumns; this.writeSortInfo = writeSortInfo; this.writeOperation = writeOperation == null ? WriteOperation.INSERT : writeOperation; @@ -187,8 +202,18 @@ public void bindDataSink(Optional insertCtx) writeContext = ctx.getStaticPartitionSpec(); branchName = ctx.getBranchName(); } + ConnectorTableHandle boundTableHandle = tableHandle; + if (connectorMetadata != null && targetTable != null) { + Optional scanParams = branchName.map(branch -> + new TableScanParams(TableScanParams.BRANCH, Collections.emptyMap(), + Collections.singletonList(branch))); + // The write target must use the pin for its exact branch, not another reference of the same table. + boundTableHandle = PluginDrivenScanNode.applyMvccSnapshotPin( + connectorMetadata, connectorSession, tableHandle, + MvccUtil.getSnapshotFromContext(targetTable, Optional.empty(), scanParams)); + } ConnectorWriteHandle handle = new PluginDrivenWriteHandle( - tableHandle, connectorColumns, overwrite, writeContext, writeSortInfo, branchName, + boundTableHandle, connectorColumns, overwrite, writeContext, writeSortInfo, branchName, writeOperation, requireMergeCardinalityCheck); ConnectorSinkPlan sinkPlan = writePlanProvider.planWrite(connectorSession, handle); this.tDataSink = sinkPlan.getDataSink(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java index 54e607dd27afa0..991f839a511bf4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java @@ -84,14 +84,18 @@ public void tryFinishSchedule() { } @Override - public final void updateFragmentExecStatus(TReportExecStatusParams params) { + public final boolean updateFragmentExecStatus(TReportExecStatusParams params) { if (params.status.status_code == TStatusCode.FINISHED) { params.status = new TStatus(TStatusCode.OK); } SingleFragmentPipelineTask fragmentTask = backendFragmentTasks.get().get( new BackendFragmentId(params.getBackendId(), params.getFragmentId())); if (fragmentTask == null) { - return; + if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() + || params.isSetMcCommitDatas()) { + throw new IllegalStateException("Missing fragment handler for external-file report"); + } + return false; } TUniqueId queryId = coordinatorContext.queryId; @@ -117,6 +121,8 @@ public final void updateFragmentExecStatus(TReportExecStatusParams params) { } } doProcessReportExecStatus(params, fragmentTask); + return !params.isSetHivePartitionUpdates() && !params.isSetIcebergCommitDatas() + && !params.isSetMcCommitDatas() || fragmentTask.isDone(); } private Map buildBackendFragmentTasks( diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 82979d901db526..35780a456b26f3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -2558,7 +2558,7 @@ private void updateScanRangeNumByScanRange(TScanRangeParams param) { } // update job progress from BE - public void updateFragmentExecStatus(TReportExecStatusParams params) { + public boolean updateFragmentExecStatus(TReportExecStatusParams params) { if (params.isSetLoadedRows() && jobId != -1) { if (params.isSetFragmentInstanceReports()) { for (TFragmentInstanceReport report : params.getFragmentInstanceReports()) { @@ -2578,82 +2578,102 @@ public void updateFragmentExecStatus(TReportExecStatusParams params) { } PipelineExecContext ctx = pipelineExecContexts.get(Pair.of(params.getFragmentId(), params.getBackendId())); - if (ctx == null || !ctx.updatePipelineStatus(params)) { + boolean hasExternalCommitData = params.isSetHivePartitionUpdates() + || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas(); + if (ctx == null) { + if (hasExternalCommitData) { + throw new IllegalStateException("Missing fragment handler for external-file report"); + } + return false; + } + if (!ctx.updatePipelineStatus(params)) { + if (hasExternalCommitData && !ctx.done) { + throw new IllegalStateException("External-file report was not a completed fragment report"); + } LOG.debug("Fragment {} is not done, ignore report status: {}", params.getFragmentId(), params.toString()); - return; + return ctx.done; } - Status status = new Status(params.status); - // for now, abort the query if we see any error except if the error is cancelled - // and returned_all_results_ is true. - // (UpdateStatus() initiates cancellation, if it hasn't already been initiated) - if (!status.ok()) { - if (returnedAllResults && status.isCancelled()) { - LOG.warn("Query {} has returned all results, fragment_id={} instance_id={}, be={}" - + " is reporting failed status {}", - DebugUtil.printId(queryId), params.getFragmentId(), - DebugUtil.printId(params.getFragmentInstanceId()), - params.getBackendId(), - status.toString()); - } else { - LOG.warn("one instance report fail, query_id={} fragment_id={} instance_id={}, be={}," - + " error message: {}", - DebugUtil.printId(queryId), params.getFragmentId(), - DebugUtil.printId(params.getFragmentInstanceId()), - params.getBackendId(), status.toString()); - updateStatus(status); + boolean accepted = false; + try { + Status status = new Status(params.status); + // for now, abort the query if we see any error except if the error is cancelled + // and returned_all_results_ is true. + // (UpdateStatus() initiates cancellation, if it hasn't already been initiated) + if (!status.ok()) { + if (returnedAllResults && status.isCancelled()) { + LOG.warn("Query {} has returned all results, fragment_id={} instance_id={}, be={}" + + " is reporting failed status {}", + DebugUtil.printId(queryId), params.getFragmentId(), + DebugUtil.printId(params.getFragmentInstanceId()), + params.getBackendId(), + status.toString()); + } else { + LOG.warn("one instance report fail, query_id={} fragment_id={} instance_id={}, be={}," + + " error message: {}", + DebugUtil.printId(queryId), params.getFragmentId(), + DebugUtil.printId(params.getFragmentInstanceId()), + params.getBackendId(), status.toString()); + updateStatus(status); + } } - } - if (params.isSetDeltaUrls() && deltaUrls != null) { - updateDeltas(params.getDeltaUrls()); - } - if (params.isSetLoadCounters() && loadCounters != null) { - updateLoadCounters(params.getLoadCounters()); - } - if (params.isSetTrackingUrl()) { - LOG.info("query_id={} tracking_url: {}", DebugUtil.printId(queryId), params.getTrackingUrl()); - trackingUrl = params.getTrackingUrl(); - } - if (params.isSetFirstErrorMsg()) { - LOG.info("query_id={} first_error_msg: {}", DebugUtil.printId(queryId), params.getFirstErrorMsg()); - firstErrorMsg = params.getFirstErrorMsg(); - } - if (params.isSetTxnId()) { - txnId = params.getTxnId(); - } - if (params.isSetLabel()) { - label = params.getLabel(); - } - if (params.isSetExportFiles()) { - updateExportFiles(params.getExportFiles()); - } - if (params.isSetCommitInfos()) { - updateCommitInfos(params.getCommitInfos()); - } - if (params.isSetErrorTabletInfos()) { - updateErrorTabletInfos(params.getErrorTabletInfos()); - } - if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas()) { - Transaction txn = Env.getCurrentEnv().getGlobalExternalTransactionInfoMgr().getTxnById(txnId); - if (params.isSetHivePartitionUpdates()) { - CommitDataSerializer.feed(txn, params.getHivePartitionUpdates()); + if (params.isSetDeltaUrls() && deltaUrls != null) { + updateDeltas(params.getDeltaUrls()); } - if (params.isSetIcebergCommitDatas()) { - CommitDataSerializer.feed(txn, params.getIcebergCommitDatas()); + if (params.isSetLoadCounters() && loadCounters != null) { + updateLoadCounters(params.getLoadCounters()); } - if (params.isSetMcCommitDatas()) { - CommitDataSerializer.feed(txn, params.getMcCommitDatas()); + if (params.isSetTrackingUrl()) { + LOG.info("query_id={} tracking_url: {}", DebugUtil.printId(queryId), params.getTrackingUrl()); + trackingUrl = params.getTrackingUrl(); } + if (params.isSetFirstErrorMsg()) { + LOG.info("query_id={} first_error_msg: {}", DebugUtil.printId(queryId), params.getFirstErrorMsg()); + firstErrorMsg = params.getFirstErrorMsg(); + } + if (params.isSetTxnId()) { + txnId = params.getTxnId(); + } + if (params.isSetLabel()) { + label = params.getLabel(); + } + if (params.isSetExportFiles()) { + updateExportFiles(params.getExportFiles()); + } + if (params.isSetCommitInfos()) { + updateCommitInfos(params.getCommitInfos()); + } + if (params.isSetErrorTabletInfos()) { + updateErrorTabletInfos(params.getErrorTabletInfos()); + } + if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() + || params.isSetMcCommitDatas()) { + Transaction txn = Env.getCurrentEnv().getGlobalExternalTransactionInfoMgr().getTxnById(txnId); + if (params.isSetHivePartitionUpdates()) { + CommitDataSerializer.feed(txn, params.getHivePartitionUpdates()); + } + if (params.isSetIcebergCommitDatas()) { + CommitDataSerializer.feed(txn, params.getIcebergCommitDatas()); + } + if (params.isSetMcCommitDatas()) { + CommitDataSerializer.feed(txn, params.getMcCommitDatas()); + } + } + + accepted = true; + } finally { + ctx.finishPipelineStatus(accepted); } - if (ctx.done) { + if (accepted) { if (LOG.isDebugEnabled()) { LOG.debug("Query {} fragment {} is marked done", DebugUtil.printId(queryId), ctx.fragmentId); } fragmentsDoneLatch.markedCountDown(params.getFragmentId(), params.getBackendId()); } + return accepted; } /* @@ -3061,6 +3081,7 @@ public static class PipelineExecContext { PlanFragmentId fragmentId; boolean initiated; boolean done; + boolean processingDoneReport; TNetworkAddress brpcAddress; TNetworkAddress address; @@ -3117,10 +3138,30 @@ public synchronized boolean updatePipelineStatus(TReportExecStatusParams params) // duplicate packet return false; } - this.done = true; + while (processingDoneReport) { + try { + wait(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted while waiting for a duplicate report", e); + } + if (this.done) { + return false; + } + } + // Serialize ownership processing so no duplicate can be acknowledged before acceptance finishes. + processingDoneReport = true; return true; } + public synchronized void finishPipelineStatus(boolean accepted) { + if (accepted) { + this.done = true; + } + processingDoneReport = false; + notifyAll(); + } + public boolean isBackendStateHealthy() { if (backend.getLastMissingHeartbeatTime() > lastMissingHeartbeatTime && !backend.isAlive()) { LOG.warn("backend {} is down while joining the coordinator. job id: {}", diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java index 6365f26c4f9c2a..bc2b991032adef 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/JobProcessor.java @@ -26,7 +26,7 @@ public interface JobProcessor { void cancel(Status cancelReason); - void updateFragmentExecStatus(TReportExecStatusParams params); + boolean updateFragmentExecStatus(TReportExecStatusParams params); void tryFinishSchedule(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java index bc4c74335ca0b7..35bc335e30468f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/NereidsCoordinator.java @@ -271,8 +271,8 @@ public boolean isDone() { } @Override - public void updateFragmentExecStatus(TReportExecStatusParams params) { - coordinatorContext.getJobProcessor().updateFragmentExecStatus(params); + public boolean updateFragmentExecStatus(TReportExecStatusParams params) { + return coordinatorContext.getJobProcessor().updateFragmentExecStatus(params); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java index 2c4202c3078757..eafbb23dfc704e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java @@ -37,6 +37,8 @@ import org.apache.doris.thrift.TUniqueId; import com.google.common.base.Strings; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; import com.google.common.collect.Maps; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -47,6 +49,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; public final class QeProcessorImpl implements QeProcessor { @@ -57,6 +60,10 @@ public final class QeProcessorImpl implements QeProcessor { private Map queryToInstancesNum; private Map userToInstancesCount; private ExecutorService writeProfileExecutor; + private final Cache acceptedExternalFileReports = CacheBuilder.newBuilder() + .maximumSize(1_000_000) + .expireAfterWrite(30, TimeUnit.MINUTES) + .build(); private final QueryFinishCallbackRegistry queryFinishCallbackRegistry = new QueryFinishCallbackRegistry(); public static final QeProcessor INSTANCE; @@ -284,22 +291,67 @@ public TReportExecStatusResult reportExecStatus(TReportExecStatusParams params, } } + boolean hasExternalCommitData = hasExternalCommitData(params); + String reportKey = hasExternalCommitData ? externalFileReportKey(params) : null; + if (hasExternalCommitData && reportKey == null) { + return rejectedExternalFileReport(result, "External-file report is missing its identity fields"); + } + if (hasExternalCommitData && acceptedExternalFileReports.getIfPresent(reportKey) != null) { + // Keep acceptance available after coordinator removal so a lost response is retry-safe. + result.setStatus(new TStatus(TStatusCode.OK)); + result.setExternalFileCommitDataAccepted(true); + return result; + } + final QueryInfo info = coordinatorMap.get(params.query_id); result.setStatus(new TStatus(TStatusCode.OK)); if (info == null) { // Currently, the execution of query is splited from the exec status process. // So, it is very likely that when exec status arrived on FE asynchronously, coordinator // has been removed from coordinatorMap. - return result; + return hasExternalCommitData + ? rejectedExternalFileReport(result, "Coordinator no longer owns this external-file report") + : result; } try { - info.getCoord().updateFragmentExecStatus(params); + boolean accepted = info.getCoord().updateFragmentExecStatus(params); + if (hasExternalCommitData && !accepted) { + return rejectedExternalFileReport(result, "FE has not accepted the external-file report"); + } } catch (Exception e) { LOG.warn("Exception during handle report, response: {}, query: {}, instance: {}", result.toString(), DebugUtil.printId(params.query_id), DebugUtil.printId(params.fragment_instance_id), e); - return result; + return hasExternalCommitData + ? rejectedExternalFileReport(result, "FE did not accept the external-file report") + : result; } result.setStatus(new TStatus(TStatusCode.OK)); + if (hasExternalCommitData) { + // Publish the retry token before replying; a transport loss cannot revoke FE ownership. + acceptedExternalFileReports.put(reportKey, Boolean.TRUE); + result.setExternalFileCommitDataAccepted(true); + } + return result; + } + + private static boolean hasExternalCommitData(TReportExecStatusParams params) { + return params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas(); + } + + private static String externalFileReportKey(TReportExecStatusParams params) { + if (!params.isSetQueryId() || !params.isSetFragmentId() || !params.isSetBackendId()) { + return null; + } + return params.getQueryId().getHi() + ":" + params.getQueryId().getLo() + ":" + + params.getFragmentId() + ":" + params.getBackendId(); + } + + private static TReportExecStatusResult rejectedExternalFileReport( + TReportExecStatusResult result, String message) { + TStatus status = new TStatus(TStatusCode.INTERNAL_ERROR); + status.addToErrorMsgs(message); + result.setStatus(status); + result.setExternalFileCommitDataAccepted(false); return result; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index 1eba8466a56912..b943c3bb127d92 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -5694,6 +5694,9 @@ public boolean isRequireSequenceInInsert() { */ public TQueryOptions toThrift() { TQueryOptions tResult = new TQueryOptions(); + // Fragment reports are decoded by FE, whose limit can be lower than a rolling-upgrade BE's. + tResult.setCoordinatorThriftMaxMessageSize(Config.thrift_max_message_size); + tResult.setSupportsExternalFileReportAck(true); tResult.setMemLimit(maxExecMemByte); tResult.setMaxScanMemRatio(maxScanMemRatio); tResult.setEnableAdaptiveScan(enableAdaptiveScan); diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java index 069cd5e8a8924b..d4878ff99a3d6c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java @@ -186,12 +186,37 @@ protected void doProcessReportExecStatus(TReportExecStatusParams params, SingleF } } - if (!fragmentTask.processReportExecStatus(params)) { + if (!fragmentTask.processReportExecStatus(params, () -> acceptFinalReport(params))) { + if ((params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() + || params.isSetMcCommitDatas()) && !fragmentTask.isDone()) { + throw new IllegalStateException("External-file report was not a completed fragment report"); + } LOG.debug("Fragment {} is not done, ignore report status: {}", params.getFragmentId(), params.toString()); return; } + if (fragmentTask.isDone()) { + if (LOG.isDebugEnabled()) { + LOG.debug("Query {} fragment {} is marked done", + DebugUtil.printId(coordinatorContext.queryId), params.getFragmentId()); + } + MarkedCountDownLatch latch = this.latch.get(); + latch.markedCountDown(params.getFragmentId(), params.getBackendId()); + + int topFragmentId = coordinatorContext.topDistributedPlan + .getFragmentJob().getFragment().getFragmentId().asInt(); + if (topFragmentId == params.getFragmentId()) { + MarkedCountDownLatch topFragmentLatch = this.topFragmentLatch.get(); + topFragmentLatch.markedCountDown(params.getFragmentId(), params.getBackendId()); + if (topFragmentLatch.getCount() == 0) { + tryFinishSchedule(); + } + } + } + } + + private void acceptFinalReport(TReportExecStatusParams params) { LoadContext loadContext = coordinatorContext.asLoadProcessor().loadContext; if (params.isSetDeltaUrls()) { loadContext.updateDeltaUrls(params.getDeltaUrls()); @@ -233,25 +258,6 @@ protected void doProcessReportExecStatus(TReportExecStatusParams params, SingleF CommitDataSerializer.feed(txn, params.getMcCommitDatas()); } } - - if (fragmentTask.isDone()) { - if (LOG.isDebugEnabled()) { - LOG.debug("Query {} fragment {} is marked done", - DebugUtil.printId(coordinatorContext.queryId), params.getFragmentId()); - } - MarkedCountDownLatch latch = this.latch.get(); - latch.markedCountDown(params.getFragmentId(), params.getBackendId()); - - int topFragmentId = coordinatorContext.topDistributedPlan - .getFragmentJob().getFragment().getFragmentId().asInt(); - if (topFragmentId == params.getFragmentId()) { - MarkedCountDownLatch topFragmentLatch = this.topFragmentLatch.get(); - topFragmentLatch.markedCountDown(params.getFragmentId(), params.getBackendId()); - if (topFragmentLatch.getCount() == 0) { - tryFinishSchedule(); - } - } - } } // Check backend health for every unfinished load fragment task. diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTask.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTask.java index c6110d6a35be01..2b5b685bdd0f27 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTask.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTask.java @@ -58,12 +58,19 @@ public SingleFragmentPipelineTask(Backend backend, int fragmentId, Set> fragments) { try { TSerializer serializer = new TSerializer(new TBinaryProtocol.Factory()); - for (TBase fragment : fragments) { - txn.addCommitData(serializer.serialize(fragment)); + List serializedFragments = fragments.stream().map(fragment -> { + try { + return serializer.serialize(fragment); + } catch (TException e) { + throw new CommitDataSerializationException(e); + } + }).collect(Collectors.toList()); + // Serialize the complete vector before mutating the transaction so malformed input is retry-safe. + for (byte[] serializedFragment : serializedFragments) { + txn.addCommitData(serializedFragment); } } catch (TException e) { - throw new RuntimeException("failed to serialize connector commit data", e); + throw new RuntimeException("failed to initialize connector commit-data serialization", e); + } catch (CommitDataSerializationException e) { + throw new RuntimeException("failed to serialize connector commit data", e.getCause()); + } + } + + private static final class CommitDataSerializationException extends RuntimeException { + private CommitDataSerializationException(TException cause) { + super(cause); } } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorIcebergRowLevelDmlTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorIcebergRowLevelDmlTest.java index b18d613d13cb99..a5b41dd630a95e 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorIcebergRowLevelDmlTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorIcebergRowLevelDmlTest.java @@ -33,8 +33,10 @@ import org.apache.doris.connector.spi.ConnectorStatementScope; import org.apache.doris.connector.spi.ConnectorType; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; +import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot; +import org.apache.doris.connector.spi.write.ConnectorSinkPlan; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; import org.apache.doris.datasource.mvcc.MvccUtil; import org.apache.doris.datasource.mvcc.PluginDrivenMvccSnapshot; @@ -50,6 +52,7 @@ import org.apache.doris.planner.DataSink; import org.apache.doris.planner.PlanFragment; import org.apache.doris.planner.PluginDrivenTableSink; +import org.apache.doris.thrift.TDataSink; import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Assertions; @@ -248,11 +251,10 @@ public void mergePluginArmRunsMaterializedNameLoopSoBeResolvesOperationColumn() } @Test - public void rowLevelDmlThreadsMvccReadSnapshotPinOntoTheWriteHandle() { + public void rowLevelDmlThreadsMvccReadSnapshotPinOntoTheWriteHandle() throws Exception { // Fix B: the write handle must carry the statement's pinned MVCC read snapshot, so a DELETE/MERGE - // re-derives its deletes from the SAME snapshot its scan read. The pin decision itself is unit-tested in - // PluginDrivenScanNodeMvccPinTest; this pins that the row-level-DML helper actually wires it onto the - // write handle (a mutation dropping the applyMvccSnapshotPin call would leave the raw, unpinned handle). + // re-derives its deletes from the SAME snapshot its scan read. Binding is intentionally late because + // that is the first point where the exact target branch is available. Plugin plugin = pluginTable(); ConnectorMvccSnapshot connectorSnapshot = Mockito.mock(ConnectorMvccSnapshot.class); PluginDrivenMvccSnapshot pinned = new PluginDrivenMvccSnapshot( @@ -271,13 +273,18 @@ public void rowLevelDmlThreadsMvccReadSnapshotPinOntoTheWriteHandle() { PlanTranslatorContext context = new PlanTranslatorContext(); PhysicalPlanTranslator translator = new PhysicalPlanTranslator(context, null); + translator.visitPhysicalExternalRowLevelDeleteSink(sink, context); + PluginDrivenTableSink pluginSink = capturePluginSink(childFragment); try (MockedStatic mvcc = Mockito.mockStatic(MvccUtil.class)) { - mvcc.when(() -> MvccUtil.getSnapshotFromContext(plugin.table)).thenReturn(Optional.of(pinned)); - translator.visitPhysicalExternalRowLevelDeleteSink(sink, context); + mvcc.when(() -> MvccUtil.getSnapshotFromContext( + plugin.table, Optional.empty(), Optional.empty())).thenReturn(Optional.of(pinned)); + pluginSink.bindDataSink(Optional.empty()); } - PluginDrivenTableSink pluginSink = capturePluginSink(childFragment); - Assertions.assertSame(pinnedHandle, Deencapsulation.getField(pluginSink, "tableHandle"), + ConnectorWritePlanProvider provider = Deencapsulation.getField(pluginSink, "writePlanProvider"); + ArgumentCaptor handle = ArgumentCaptor.forClass(ConnectorWriteHandle.class); + Mockito.verify(provider).planWrite(Mockito.any(), handle.capture()); + Assertions.assertSame(pinnedHandle, handle.getValue().getTableHandle(), "the row-level DML write handle must carry the snapshot-pinned table handle (Fix B), not the raw" + " latest-read handle"); } @@ -362,6 +369,8 @@ private static Plugin pluginTable() { // provider and admits on ITS supportedOperations containing DELETE/MERGE. Mockito.when(provider.supportedOperations()) .thenReturn(EnumSet.of(WriteOperation.DELETE, WriteOperation.MERGE)); + Mockito.when(provider.planWrite(Mockito.any(), Mockito.any())) + .thenReturn(new ConnectorSinkPlan(new TDataSink())); Mockito.when(connector.getMetadata(Mockito.any())).thenReturn(metadata); Mockito.when(metadata.getTableHandle(Mockito.any(), Mockito.any(), Mockito.any())) .thenReturn(Optional.of(handle)); 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 86569395c902c9..75fd22c20ce220 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-core/src/test/java/org/apache/doris/planner/PluginDrivenTableSinkBindingTest.java b/fe/fe-core/src/test/java/org/apache/doris/planner/PluginDrivenTableSinkBindingTest.java index 5d8f4f28052a1b..4f80f34b6fe393 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/planner/PluginDrivenTableSinkBindingTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/planner/PluginDrivenTableSinkBindingTest.java @@ -17,18 +17,26 @@ package org.apache.doris.planner; +import org.apache.doris.analysis.TableScanParams; import org.apache.doris.common.AnalysisException; import org.apache.doris.connector.ConnectorSessionBuilder; +import org.apache.doris.connector.spi.ConnectorMetadata; import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; +import org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot; import org.apache.doris.connector.spi.write.ConnectorSinkPlan; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; +import org.apache.doris.datasource.mvcc.MvccUtil; +import org.apache.doris.datasource.mvcc.PluginDrivenMvccSnapshot; +import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; import org.apache.doris.nereids.trees.plans.commands.insert.PluginDrivenInsertCommandContext; import org.apache.doris.thrift.TDataSink; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import java.util.Collections; import java.util.HashMap; @@ -88,6 +96,40 @@ public void absentContextDefaultsToNonOverwriteEmptySpec() throws AnalysisExcept "a plain INSERT must pass an empty static partition spec"); } + @Test + public void branchTargetUsesItsExactVersionAwareSnapshotPin() throws AnalysisException { + RecordingWritePlanProvider provider = new RecordingWritePlanProvider(); + ConnectorSession session = ConnectorSessionBuilder.create().withCatalogName("iceberg").build(); + ConnectorTableHandle baseHandle = Mockito.mock(ConnectorTableHandle.class); + ConnectorTableHandle pinnedHandle = Mockito.mock(ConnectorTableHandle.class); + ConnectorMetadata metadata = Mockito.mock(ConnectorMetadata.class); + PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); + ConnectorMvccSnapshot connectorSnapshot = Mockito.mock(ConnectorMvccSnapshot.class); + PluginDrivenMvccSnapshot snapshot = new PluginDrivenMvccSnapshot( + connectorSnapshot, Collections.emptyMap(), Collections.emptyMap()); + Mockito.when(metadata.applySnapshot(session, baseHandle, connectorSnapshot)).thenReturn(pinnedHandle); + PluginDrivenTableSink sink = new PluginDrivenTableSink(table, provider, session, baseHandle, + Collections.emptyList(), null, null, false, metadata); + PluginDrivenInsertCommandContext ctx = new PluginDrivenInsertCommandContext(); + ctx.setBranchName(Optional.of("audit")); + + try (MockedStatic mvcc = Mockito.mockStatic(MvccUtil.class)) { + mvcc.when(() -> MvccUtil.getSnapshotFromContext( + Mockito.eq(table), Mockito.eq(Optional.empty()), Mockito.any())) + .thenAnswer(invocation -> { + Optional selector = invocation.getArgument(2); + Assertions.assertEquals(TableScanParams.BRANCH, + selector.orElseThrow().getParamType()); + Assertions.assertEquals(Collections.singletonList("audit"), + selector.orElseThrow().getListParams()); + return Optional.of(snapshot); + }); + sink.bindDataSink(Optional.of(ctx)); + } + + Assertions.assertSame(pinnedHandle, provider.capturedHandle.getTableHandle()); + } + private static PluginDrivenTableSink newPlanProviderSink(ConnectorWritePlanProvider provider) { ConnectorSession session = ConnectorSessionBuilder.create().withCatalogName("mc_cat").build(); ConnectorTableHandle tableHandle = new ConnectorTableHandle() { }; diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java new file mode 100644 index 00000000000000..fd8f34d90e73ab --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java @@ -0,0 +1,136 @@ +// 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.qe; + +import org.apache.doris.common.profile.ExecutionProfile; +import org.apache.doris.planner.PlanFragmentId; +import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TIcebergCommitData; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TQueryOptions; +import org.apache.doris.thrift.TReportExecStatusParams; +import org.apache.doris.thrift.TReportExecStatusResult; +import org.apache.doris.thrift.TStatus; +import org.apache.doris.thrift.TStatusCode; +import org.apache.doris.thrift.TUniqueId; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.util.Collections; + +class QeProcessorImplReportAckTest { + private TUniqueId registeredQueryId; + + @AfterEach + void cleanup() { + if (registeredQueryId != null) { + QeProcessorImpl.INSTANCE.unregisterQuery(registeredQueryId); + } + } + + @Test + void rejectsExternalReportWithoutCoordinator() { + TReportExecStatusResult result = report(params(new TUniqueId(12345, 1))); + + Assertions.assertEquals(TStatusCode.INTERNAL_ERROR, result.getStatus().getStatusCode()); + Assertions.assertFalse(result.isExternalFileCommitDataAccepted()); + } + + @Test + void rejectsExternalReportWhenHandlerThrows() throws Exception { + TUniqueId queryId = new TUniqueId(12345, 2); + Coordinator coordinator = register(queryId); + Mockito.doThrow(new RuntimeException("injected acceptance failure")) + .when(coordinator).updateFragmentExecStatus(Mockito.any()); + + TReportExecStatusResult result = report(params(queryId)); + + Assertions.assertEquals(TStatusCode.INTERNAL_ERROR, result.getStatus().getStatusCode()); + Assertions.assertFalse(result.isExternalFileCommitDataAccepted()); + } + + @Test + void rejectsExternalReportWhenHandlerDoesNotAcceptIt() throws Exception { + TUniqueId queryId = new TUniqueId(12345, 4); + register(queryId); + + TReportExecStatusResult result = report(params(queryId)); + + Assertions.assertEquals(TStatusCode.INTERNAL_ERROR, result.getStatus().getStatusCode()); + Assertions.assertFalse(result.isExternalFileCommitDataAccepted()); + } + + @Test + void retriesAcceptedExternalReportAfterCoordinatorRemoval() throws Exception { + TUniqueId queryId = new TUniqueId(12345, 3); + Coordinator coordinator = register(queryId); + Mockito.when(coordinator.updateFragmentExecStatus(Mockito.any())).thenReturn(true); + TReportExecStatusParams params = params(queryId); + + TReportExecStatusResult first = report(params); + QeProcessorImpl.INSTANCE.unregisterQuery(queryId); + registeredQueryId = null; + TReportExecStatusResult retry = report(params); + + Assertions.assertTrue(first.isExternalFileCommitDataAccepted()); + Assertions.assertTrue(retry.isExternalFileCommitDataAccepted()); + Assertions.assertEquals(TStatusCode.OK, retry.getStatus().getStatusCode()); + Mockito.verify(coordinator, Mockito.times(1)).updateFragmentExecStatus(params); + } + + @Test + void legacyCoordinatorRetriesFailedAcceptanceBeforeMarkingDone() { + Backend backend = Mockito.mock(Backend.class); + Mockito.when(backend.getHost()).thenReturn("127.0.0.1"); + ExecutionProfile profile = Mockito.mock(ExecutionProfile.class); + Coordinator.PipelineExecContext context = new Coordinator.PipelineExecContext( + new PlanFragmentId(7), null, backend, profile, -1); + TReportExecStatusParams report = new TReportExecStatusParams().setDone(true); + + Assertions.assertTrue(context.updatePipelineStatus(report)); + context.finishPipelineStatus(false); + Assertions.assertTrue(context.updatePipelineStatus(report)); + context.finishPipelineStatus(true); + Assertions.assertFalse(context.updatePipelineStatus(report)); + } + + private Coordinator register(TUniqueId queryId) throws Exception { + Coordinator coordinator = Mockito.mock(Coordinator.class); + Mockito.when(coordinator.getQueryOptions()).thenReturn(new TQueryOptions()); + QeProcessorImpl.INSTANCE.registerQuery(queryId, new QeProcessorImpl.QueryInfo(coordinator)); + registeredQueryId = queryId; + return coordinator; + } + + private static TReportExecStatusParams params(TUniqueId queryId) { + return new TReportExecStatusParams() + .setQueryId(queryId) + .setFragmentId(7) + .setBackendId(9) + .setDone(true) + .setStatus(new TStatus(TStatusCode.OK)) + .setIcebergCommitDatas(Collections.emptyList()); + } + + private static TReportExecStatusResult report(TReportExecStatusParams params) { + return QeProcessorImpl.INSTANCE.reportExecStatus(params, new TNetworkAddress("127.0.0.1", 9050)); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java index 75bc69ba7c5531..5d10692f33738b 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java @@ -401,4 +401,14 @@ public void testFileCacheQueryLimitBytesToThrift() throws Exception { Assertions.assertTrue(queryOptions.isSetFileCacheQueryLimitBytes()); Assertions.assertEquals(262144L, queryOptions.getFileCacheQueryLimitBytes()); } + + @Test + public void testCoordinatorThriftLimitPropagatesToBackends() { + TQueryOptions queryOptions = new SessionVariable().toThrift(); + + Assertions.assertTrue(queryOptions.isSetCoordinatorThriftMaxMessageSize()); + Assertions.assertEquals(Config.thrift_max_message_size, + queryOptions.getCoordinatorThriftMaxMessageSize()); + Assertions.assertTrue(queryOptions.isSupportsExternalFileReportAck()); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTaskTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTaskTest.java index 31becae01dc9db..281da9ab3dfe1d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTaskTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/runtime/SingleFragmentPipelineTaskTest.java @@ -19,6 +19,7 @@ import org.apache.doris.common.Status; import org.apache.doris.system.Backend; +import org.apache.doris.thrift.TReportExecStatusParams; import org.apache.doris.thrift.TStatusCode; import org.apache.doris.thrift.TUniqueId; @@ -28,6 +29,20 @@ import java.util.Collections; class SingleFragmentPipelineTaskTest { + @Test + void failedAcceptanceLeavesFinalReportRetryable() { + SingleFragmentPipelineTask task = createTask(createBackend(100L)); + TReportExecStatusParams report = new TReportExecStatusParams().setDone(true); + + Assertions.assertThrows(RuntimeException.class, + () -> task.processReportExecStatus(report, () -> { + throw new RuntimeException("injected failure"); + })); + Assertions.assertFalse(task.isDone()); + Assertions.assertTrue(task.processReportExecStatus(report, () -> { })); + Assertions.assertTrue(task.isDone()); + } + @Test void backendWithUnchangedProcessEpochIsHealthy() { Backend backend = createBackend(100L); 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..a269ceca30c5f0 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 @@ -25,6 +25,8 @@ import com.azure.core.http.rest.PagedIterable; import com.azure.core.http.rest.PagedResponse; +import com.azure.core.util.BinaryData; +import com.azure.core.util.Context; import com.azure.identity.ClientSecretCredentialBuilder; import com.azure.storage.blob.BlobClient; import com.azure.storage.blob.BlobContainerClient; @@ -33,10 +35,14 @@ import com.azure.storage.blob.BlobServiceClientBuilder; import com.azure.storage.blob.models.BlobItem; import com.azure.storage.blob.models.BlobProperties; +import com.azure.storage.blob.models.BlobRequestConditions; import com.azure.storage.blob.models.BlobStorageException; import com.azure.storage.blob.models.ListBlobsOptions; +import com.azure.storage.blob.options.BlockBlobCommitBlockListOptions; import com.azure.storage.blob.sas.BlobSasPermission; import com.azure.storage.blob.sas.BlobServiceSasSignatureValues; +import com.azure.storage.blob.specialized.BlobLeaseClient; +import com.azure.storage.blob.specialized.BlobLeaseClientBuilder; import com.azure.storage.blob.specialized.BlockBlobClient; import com.azure.storage.common.StorageSharedKeyCredential; import com.azure.storage.common.sas.SasProtocol; @@ -48,6 +54,7 @@ import java.io.InputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; import java.time.OffsetDateTime; import java.util.ArrayList; import java.util.Base64; @@ -56,6 +63,7 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.UUID; /** * Azure Blob Storage implementation of {@link ObjStorage}. @@ -69,6 +77,8 @@ public class AzureObjStorage implements ObjStorage { private static final Logger LOG = LogManager.getLogger(AzureObjStorage.class); private static final int HTTP_NOT_FOUND = 404; + private static final int MULTIPART_LEASE_SECONDS = 60; + private static final String MULTIPART_LEASE_PREFIX = "doris-azure-lease-v1:"; /** Validity period for presigned (SAS) URLs, in seconds. */ private static final int SESSION_EXPIRE_SECONDS = 3600; @@ -218,9 +228,23 @@ public void copyObject(String srcPath, String dstPath) throws IOException { @Override public String initiateMultipartUpload(String remotePath) throws IOException { - // Azure block blobs don't have an explicit "initiate" API. - // Return the path itself as the upload ID; block IDs are derived from part numbers. - return remotePath; + try { + AzureUri uri = AzureUri.parse(remotePath); + BlobClient blobClient = getClient().getBlobContainerClient(uri.container()) + .getBlobClient(uri.key()); + BlockBlobClient blockBlobClient = blobClient.getBlockBlobClient(); + String leaseId = UUID.randomUUID().toString(); + String uploadId = MULTIPART_LEASE_PREFIX + leaseId; + // A zero-byte uncommitted block materializes an absent target without exposing it to + // normal listings, so Azure can fence every later block operation with a blob lease. + blockBlobClient.stageBlock(multipartBlockId(uploadId, 0), BinaryData.fromBytes(new byte[0])); + String acquiredLeaseId = createLeaseClient(blobClient, leaseId) + .acquireLease(MULTIPART_LEASE_SECONDS); + return MULTIPART_LEASE_PREFIX + acquiredLeaseId; + } catch (BlobStorageException e) { + throw new IOException("initiateMultipartUpload failed for " + remotePath + + ": " + e.getMessage(), e); + } } @Override @@ -228,10 +252,18 @@ public UploadPartResult uploadPart(String remotePath, String uploadId, int partN RequestBody body) throws IOException { try { AzureUri uri = AzureUri.parse(remotePath); - BlockBlobClient blockBlobClient = getClient().getBlobContainerClient(uri.container()) - .getBlobClient(uri.key()).getBlockBlobClient(); - String blockId = toBlockId(partNum); - blockBlobClient.stageBlock(blockId, body.content(), body.contentLength()); + BlobClient blobClient = getClient().getBlobContainerClient(uri.container()) + .getBlobClient(uri.key()); + BlockBlobClient blockBlobClient = blobClient.getBlockBlobClient(); + String blockId = multipartBlockId(uploadId, partNum); + String leaseId = multipartLeaseId(uploadId); + if (leaseId == null) { + blockBlobClient.stageBlock(blockId, body.content(), body.contentLength()); + } else { + renewMultipartLease(blobClient, leaseId); + blockBlobClient.stageBlockWithResponse(blockId, body.content(), body.contentLength(), + null, leaseId, null, Context.NONE); + } return new UploadPartResult(partNum, blockId); } catch (BlobStorageException e) { throw new IOException("uploadPart failed for " + remotePath + " part " + partNum @@ -244,15 +276,38 @@ public void completeMultipartUpload(String remotePath, String uploadId, List parts) throws IOException { try { AzureUri uri = AzureUri.parse(remotePath); - BlockBlobClient blockBlobClient = getClient().getBlobContainerClient(uri.container()) - .getBlobClient(uri.key()).getBlockBlobClient(); + BlobContainerClient containerClient = getClient().getBlobContainerClient(uri.container()); List blockIds = new ArrayList<>(); List sorted = new ArrayList<>(parts); sorted.sort((a, b) -> Integer.compare(a.partNumber(), b.partNumber())); for (UploadPartResult part : sorted) { - blockIds.add(toBlockId(part.partNumber())); + // Guessing a legacy ID cannot recover an old BE payload and can select another writer's block. + if (part.etag() == null || part.etag().isEmpty()) { + throw new IOException("Azure multipart completion requires the exact staged block ID " + + "for part " + part.partNumber()); + } + blockIds.add(part.etag()); + } + // Put Block List is the atomic publication point and does not expose a staging blob to scans. + BlobClient blobClient = containerClient.getBlobClient(uri.key()); + BlockBlobClient blockBlobClient = blobClient.getBlockBlobClient(); + String leaseId = multipartLeaseId(uploadId); + if (leaseId == null) { + blockBlobClient.commitBlockList(blockIds); + } else { + BlobLeaseClient leaseClient = renewMultipartLease(blobClient, leaseId); + BlobRequestConditions conditions = new BlobRequestConditions().setLeaseId(leaseId); + blockBlobClient.commitBlockListWithResponse( + new BlockBlobCommitBlockListOptions(blockIds).setRequestConditions(conditions), + null, Context.NONE); + try { + leaseClient.releaseLease(); + } catch (BlobStorageException e) { + // Publication is already durable; the finite lease will expire without + // turning a successful commit into a retry that could overwrite new data. + LOG.warn("Failed to release Azure multipart lease after commit for {}", remotePath, e); + } } - blockBlobClient.commitBlockList(blockIds); } catch (BlobStorageException e) { throw new IOException("completeMultipartUpload failed for " + remotePath + ": " + e.getMessage(), e); @@ -261,43 +316,41 @@ public void completeMultipartUpload(String remotePath, String uploadId, @Override public void abortMultipartUpload(String remotePath, String uploadId) throws IOException { - // Azure has no native "abort multipart upload" API; the closest equivalent is to - // commit an empty block list (which atomically discards any uncommitted blocks - // for that blob) and then delete the resulting empty blob so no trace remains. - // - // SAFETY: commitBlockList(empty) overwrites whatever is at the target blob, so we - // MUST refuse to run when a committed blob already exists at this path — otherwise - // an abort call could destroy real user data. In that case the staged blocks are - // left to expire on their own (Azure GCs them after the service-side timeout). + String leaseId = multipartLeaseId(uploadId); + if (leaseId == null) { + // Azure cannot selectively remove legacy uncommitted blocks without rewriting the blob. + return; + } try { AzureUri uri = AzureUri.parse(remotePath); - BlockBlobClient blockBlobClient = getClient().getBlobContainerClient(uri.container()) - .getBlobClient(uri.key()).getBlockBlobClient(); - boolean committedBlobExists; - try { - blockBlobClient.getProperties(); - committedBlobExists = true; - } catch (BlobStorageException e) { - if (e.getStatusCode() != HTTP_NOT_FOUND) { - throw e; - } - committedBlobExists = false; - } - if (committedBlobExists) { - LOG.warn("abortMultipartUpload skipped for {}: a committed blob already exists; " - + "uncommitted blocks will expire automatically.", remotePath); - return; - } - blockBlobClient.commitBlockList(Collections.emptyList()); - blockBlobClient.delete(); + BlobClient blobClient = getClient().getBlobContainerClient(uri.container()) + .getBlobClient(uri.key()); + createLeaseClient(blobClient, leaseId).releaseLease(); } catch (BlobStorageException e) { - // Best-effort: log and swallow rather than mask the original failure that - // triggered the abort path. Uncommitted blocks will be GC'd by the service. - LOG.warn("abortMultipartUpload best-effort cleanup failed for {}: {}", - remotePath, e.getMessage()); + throw new IOException("abortMultipartUpload failed for " + remotePath + + ": " + e.getMessage(), e); } } + protected BlobLeaseClient createLeaseClient(BlobClient blobClient, String leaseId) { + return new BlobLeaseClientBuilder().blobClient(blobClient).leaseId(leaseId).buildClient(); + } + + private BlobLeaseClient renewMultipartLease(BlobClient blobClient, String leaseId) { + BlobLeaseClient leaseClient = createLeaseClient(blobClient, leaseId); + // A renewal failure loses the upload-generation fence even if the same ID is acquirable later. + leaseClient.renewLease(); + return leaseClient; + } + + private static String multipartLeaseId(String uploadId) { + if (uploadId != null && uploadId.startsWith(MULTIPART_LEASE_PREFIX) + && uploadId.length() > MULTIPART_LEASE_PREFIX.length()) { + return uploadId.substring(MULTIPART_LEASE_PREFIX.length()); + } + return null; + } + /** * Opens an InputStream to download the blob at the given path. * Used by {@link AzureFileSystem} for read operations. @@ -524,4 +577,15 @@ private static String toBlockId(int partNum) { byte[] bytes = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(partNum).array(); return Base64.getEncoder().encodeToString(bytes); } + + static String multipartBlockId(String uploadId, int partNum) { + int uploadNamespace = 0x811C9DC5; + for (byte value : uploadId.getBytes(StandardCharsets.UTF_8)) { + uploadNamespace = (uploadNamespace ^ (value & 0xFF)) * 0x01000193; + } + int namespacedPart = uploadNamespace + partNum; + // Match the legacy four-byte length so a retry can coexist with pre-upgrade residual blocks. + byte[] rawId = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(namespacedPart).array(); + return Base64.getEncoder().encodeToString(rawId); + } } 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..b7aa4bcf9f19f0 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,21 +17,30 @@ package org.apache.doris.filesystem.azure; +import org.apache.doris.filesystem.UploadPartResult; import org.apache.doris.filesystem.spi.RemoteObjects; +import org.apache.doris.filesystem.spi.RequestBody; import com.azure.storage.blob.BlobClient; import com.azure.storage.blob.BlobContainerClient; import com.azure.storage.blob.BlobServiceClient; import com.azure.storage.blob.models.BlobProperties; +import com.azure.storage.blob.models.BlobRequestConditions; import com.azure.storage.blob.models.BlobStorageException; +import com.azure.storage.blob.options.BlockBlobCommitBlockListOptions; +import com.azure.storage.blob.specialized.BlobLeaseClient; +import com.azure.storage.blob.specialized.BlockBlobClient; import com.azure.storage.common.StorageSharedKeyCredential; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import java.io.ByteArrayInputStream; import java.io.IOException; import java.time.OffsetDateTime; +import java.util.ArrayList; import java.util.Arrays; +import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -379,47 +388,274 @@ void deleteObjectsByKeys_attachesPerKeyExceptionsAsSuppressed() throws Exception // ------------------------------------------------------------------ @Test - void abortMultipartUpload_safeNoopWhenCommittedBlobExists() throws Exception { - com.azure.storage.blob.models.BlobProperties props = - Mockito.mock(com.azure.storage.blob.models.BlobProperties.class); - Mockito.when(props.getBlobSize()).thenReturn(1024L); + void multipartBlockId_keepsLegacyLengthButCannotIdentifyWriter() { + Assertions.assertEquals("p3w3DA==", AzureObjStorage.multipartBlockId("upload-a", 1)); + Assertions.assertEquals("Sc7grw==", AzureObjStorage.multipartBlockId( + "doris-azure-lease-v1:09492e3d-e231-4ed9-bf84-b6fc772cda54", 1)); + Assertions.assertEquals("Sc7grw==", AzureObjStorage.multipartBlockId( + "doris-azure-lease-v1:06996d15-1c2e-4ddd-8853-43816ea84a07", 1)); + Assertions.assertEquals( + AzureObjStorage.multipartBlockId("upload-a", 1).length(), + AzureObjStorage.multipartBlockId("upload-a", 999).length()); + Assertions.assertEquals(4, + Base64.getDecoder().decode(AzureObjStorage.multipartBlockId("upload-a", 1)).length); + } - com.azure.storage.blob.specialized.BlockBlobClient blockClient = - Mockito.mock(com.azure.storage.blob.specialized.BlockBlobClient.class); - Mockito.when(blockClient.getProperties()).thenReturn(props); + @Test + void initiateMultipartUpload_reservesTargetAndAcquiresLease() throws Exception { + BlockBlobClient blockClient = Mockito.mock(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); + BlobLeaseClient leaseClient = Mockito.mock(BlobLeaseClient.class); + Mockito.when(leaseClient.acquireLease(60)).thenReturn("lease-id"); + TestableAzureObjStorage storage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, leaseClient); + + String uploadId = storage.initiateMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob"); + + Assertions.assertEquals("doris-azure-lease-v1:lease-id", uploadId); + Mockito.verify(blockClient).stageBlock( + Mockito.anyString(), Mockito.any(com.azure.core.util.BinaryData.class)); + Mockito.verify(leaseClient).acquireLease(60); + } + @Test + void uploadPart_renewsLeaseAndFencesStagedBlock() throws Exception { + BlockBlobClient blockClient = Mockito.mock(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); + BlobLeaseClient leaseClient = Mockito.mock(BlobLeaseClient.class); + TestableAzureObjStorage storage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, leaseClient); + + storage.uploadPart("wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "doris-azure-lease-v1:lease-id", 1, + RequestBody.of(new ByteArrayInputStream(new byte[]{1}), 1)); + + Mockito.verify(leaseClient).renewLease(); + Mockito.verify(blockClient).stageBlockWithResponse(Mockito.anyString(), + Mockito.any(java.io.InputStream.class), Mockito.eq(1L), Mockito.isNull(), + Mockito.eq("lease-id"), Mockito.isNull(), Mockito.any()); + } + @Test + void completeMultipartUpload_fencesCommitAndReleasesLease() throws Exception { + BlockBlobClient blockClient = Mockito.mock(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); + BlobLeaseClient leaseClient = Mockito.mock(BlobLeaseClient.class); + TestableAzureObjStorage storage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, leaseClient); + + storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "doris-azure-lease-v1:lease-id", + Collections.singletonList(new UploadPartResult(1, "AQAAAA=="))); + + Mockito.verify(leaseClient).renewLease(); + org.mockito.ArgumentCaptor options = + org.mockito.ArgumentCaptor.forClass(BlockBlobCommitBlockListOptions.class); + Mockito.verify(blockClient).commitBlockListWithResponse( + options.capture(), Mockito.isNull(), Mockito.any()); + BlobRequestConditions conditions = options.getValue().getRequestConditions(); + Assertions.assertEquals("lease-id", conditions.getLeaseId()); + Mockito.verify(leaseClient).releaseLease(); + } + @Test + void completeMultipartUpload_lostLeaseFailsBeforePublication() throws Exception { + BlockBlobClient blockClient = Mockito.mock(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); + BlobLeaseClient leaseClient = Mockito.mock(BlobLeaseClient.class); + BlobStorageException lostLease = Mockito.mock(BlobStorageException.class); + Mockito.when(leaseClient.renewLease()).thenThrow(lostLease); + TestableAzureObjStorage storage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, leaseClient); - TestableAzureObjStorage storage = new TestableAzureObjStorage(buildBasicProps(), serviceClient); + Assertions.assertThrows(IOException.class, () -> storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "doris-azure-lease-v1:lease-id", + Collections.singletonList(new UploadPartResult(1, "AQAAAA==")))); + + Mockito.verify(blockClient, Mockito.never()).commitBlockList(Mockito.anyList()); + Mockito.verify(blockClient, Mockito.never()).commitBlockListWithResponse( + Mockito.any(BlockBlobCommitBlockListOptions.class), Mockito.isNull(), Mockito.any()); + } + + @Test + void completeMultipartUpload_expiredLeaseFailsClosedAfterCollidingWriterStagesAndReleases() throws Exception { + BlockBlobClient blockClient = Mockito.mock(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); + BlobLeaseClient staleLeaseClient = Mockito.mock(BlobLeaseClient.class); + BlobLeaseClient competingLeaseClient = Mockito.mock(BlobLeaseClient.class); + BlobStorageException expiredLease = Mockito.mock(BlobStorageException.class); + Mockito.when(expiredLease.getStatusCode()).thenReturn(409); + Mockito.when(staleLeaseClient.renewLease()).thenThrow(expiredLease); + TestableAzureObjStorage staleStorage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, staleLeaseClient); + TestableAzureObjStorage competingStorage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, competingLeaseClient); + String staleUpload = "doris-azure-lease-v1:09492e3d-e231-4ed9-bf84-b6fc772cda54"; + String competingUpload = "doris-azure-lease-v1:06996d15-1c2e-4ddd-8853-43816ea84a07"; + String collidingBlockId = AzureObjStorage.multipartBlockId(staleUpload, 1); + + competingStorage.uploadPart( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + competingUpload, 1, RequestBody.of(new ByteArrayInputStream(new byte[]{2}), 1)); + competingStorage.abortMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", competingUpload); + + Assertions.assertThrows(IOException.class, () -> staleStorage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + staleUpload, Collections.singletonList(new UploadPartResult(1, collidingBlockId)))); + + Mockito.verify(blockClient).stageBlockWithResponse(Mockito.eq(collidingBlockId), + Mockito.any(java.io.InputStream.class), Mockito.eq(1L), Mockito.isNull(), + Mockito.eq("06996d15-1c2e-4ddd-8853-43816ea84a07"), Mockito.isNull(), Mockito.any()); + Mockito.verify(competingLeaseClient).releaseLease(); + Mockito.verify(staleLeaseClient, Mockito.never()).acquireLease(Mockito.anyInt()); + Mockito.verify(blockClient, Mockito.never()).commitBlockListWithResponse( + Mockito.any(BlockBlobCommitBlockListOptions.class), Mockito.isNull(), Mockito.any()); + } + + @Test + void abortMultipartUpload_releasesLeasedSessionWithoutRewritingTarget() throws Exception { + BlobClient blobClient = Mockito.mock(BlobClient.class); + 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); + BlobLeaseClient leaseClient = Mockito.mock(BlobLeaseClient.class); + TestableAzureObjStorage storage = + new TestableAzureObjStorage(buildBasicProps(), serviceClient, leaseClient); storage.abortMultipartUpload( - "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", "uploadId"); + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "doris-azure-lease-v1:lease-id"); - // The committed blob must NOT be touched (no commitBlockList, no delete). - Mockito.verify(blockClient, Mockito.never()).commitBlockList(Mockito.anyList()); - Mockito.verify(blockClient, Mockito.never()).delete(); + Mockito.verify(leaseClient).releaseLease(); + Mockito.verify(blobClient, Mockito.never()).delete(); } @Test - void abortMultipartUpload_commitsEmptyAndDeletesWhenNoCommittedBlob() throws Exception { + void uploadPart_acceptsLegacyResidualBlockLength() throws Exception { com.azure.storage.blob.specialized.BlockBlobClient blockClient = Mockito.mock(com.azure.storage.blob.specialized.BlockBlobClient.class); - BlobStorageException notFoundEx = Mockito.mock(BlobStorageException.class); - Mockito.when(notFoundEx.getStatusCode()).thenReturn(404); - Mockito.when(blockClient.getProperties()).thenThrow(notFoundEx); + List stagedBlockIds = new ArrayList<>(Collections.singletonList("AQAAAA==")); + Mockito.doAnswer(invocation -> { + String blockId = invocation.getArgument(0); + int requiredDecodedLength = Base64.getDecoder().decode(stagedBlockIds.get(0)).length; + if (Base64.getDecoder().decode(blockId).length != requiredDecodedLength) { + throw new IllegalStateException("Azure would reject a different block ID length"); + } + stagedBlockIds.add(blockId); + return null; + }).when(blockClient).stageBlock( + Mockito.anyString(), Mockito.any(java.io.InputStream.class), Mockito.anyLong()); + 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); + UploadPartResult result = storage.uploadPart( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "new-upload-id", 1, RequestBody.of(new ByteArrayInputStream(new byte[]{1}), 1)); + + Assertions.assertEquals(Arrays.asList("AQAAAA==", result.etag()), stagedBlockIds); + } + + @Test + void completeMultipartUpload_usesExactBlockIdsReportedByBe() throws Exception { + com.azure.storage.blob.specialized.BlockBlobClient blockClient = + Mockito.mock(com.azure.storage.blob.specialized.BlockBlobClient.class); + BlobClient targetBlob = Mockito.mock(BlobClient.class); + Mockito.when(targetBlob.getBlockBlobClient()).thenReturn(blockClient); + BlobContainerClient containerClient = Mockito.mock(BlobContainerClient.class); + Mockito.when(containerClient.getBlobClient("stage/blob")).thenReturn(targetBlob); + 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)); + Mockito.verify(targetBlob, Mockito.never()).beginCopy(Mockito.anyString(), Mockito.isNull()); + } + + @Test + void completeMultipartUpload_rejectsOlderBeWithoutBlockIds() 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); + Assertions.assertThrows(IOException.class, () -> storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "legacy-upload-id", Collections.singletonList(new UploadPartResult(1, "")))); + + Mockito.verify(blockClient, Mockito.never()).commitBlockList(Mockito.anyList()); + } + + @Test + void completeMultipartUpload_rejectsMixedExactAndMissingBlockIds() 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); + + Assertions.assertThrows(IOException.class, () -> storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "mixed-upload-id", + Arrays.asList(new UploadPartResult(1, "exact-id"), new UploadPartResult(2, "")))); + + Mockito.verify(blockClient, Mockito.never()).commitBlockList(Mockito.anyList()); + } + + @Test + void abortMultipartUpload_doesNotMutatePublishedTarget() throws Exception { + BlobContainerClient containerClient = Mockito.mock(BlobContainerClient.class); BlobServiceClient serviceClient = Mockito.mock(BlobServiceClient.class); Mockito.when(serviceClient.getBlobContainerClient("mycontainer")).thenReturn(containerClient); @@ -429,8 +665,8 @@ void abortMultipartUpload_commitsEmptyAndDeletesWhenNoCommittedBlob() throws Exc storage.abortMultipartUpload( "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", "uploadId"); - Mockito.verify(blockClient).commitBlockList(Collections.emptyList()); - Mockito.verify(blockClient).delete(); + Mockito.verify(containerClient, Mockito.never()).getBlobClient("stage/blob"); + Mockito.verifyNoInteractions(containerClient); } // ------------------------------------------------------------------ @@ -450,13 +686,20 @@ private Map buildBasicProps() { */ private static class TestableAzureObjStorage extends AzureObjStorage { private final BlobServiceClient mockClient; + private final BlobLeaseClient mockLeaseClient; String stubbedSasUrl = "https://stubbed-sas-url"; String lastGenerateSasContainer; String lastGenerateSasBlobKey; TestableAzureObjStorage(Map props, BlobServiceClient mockClient) { + this(props, mockClient, null); + } + + TestableAzureObjStorage(Map props, BlobServiceClient mockClient, + BlobLeaseClient mockLeaseClient) { super(props); this.mockClient = mockClient; + this.mockLeaseClient = mockLeaseClient; } @Override @@ -464,6 +707,11 @@ protected BlobServiceClient buildClient() { return mockClient; } + @Override + protected BlobLeaseClient createLeaseClient(BlobClient blobClient, String leaseId) { + return mockLeaseClient; + } + @Override protected String generateSasUrl(String endpoint, String container, String blobKey, StorageSharedKeyCredential credential, OffsetDateTime expiresOn) { diff --git a/gensrc/thrift/DataSinks.thrift b/gensrc/thrift/DataSinks.thrift index 9efe29a51eb421..d2cb4d534bb49b 100644 --- a/gensrc/thrift/DataSinks.thrift +++ b/gensrc/thrift/DataSinks.thrift @@ -384,6 +384,7 @@ struct THiveTableSink { 10: optional bool overwrite 11: optional THiveSerDeProperties serde_properties 12: optional list broker_addresses; + 13: optional bool supports_deferred_azure_multipart } enum TUpdateMode { diff --git a/gensrc/thrift/FrontendService.thrift b/gensrc/thrift/FrontendService.thrift index 39bdbf022fad34..d2e622142ef320 100644 --- a/gensrc/thrift/FrontendService.thrift +++ b/gensrc/thrift/FrontendService.thrift @@ -183,6 +183,8 @@ struct TListPrivilegesResult{ struct TReportExecStatusResult { // required in V1 1: optional Status.TStatus status + // Set only after FE accepts the external-file commit vectors for this report. + 2: optional bool external_file_commit_data_accepted } // Service Protocol Details diff --git a/gensrc/thrift/PaloInternalService.thrift b/gensrc/thrift/PaloInternalService.thrift index 549118e26a1ea7..37b53f7d75d3ca 100644 --- a/gensrc/thrift/PaloInternalService.thrift +++ b/gensrc/thrift/PaloInternalService.thrift @@ -512,6 +512,10 @@ struct TQueryOptions { 226: optional bool enable_prune_nested_column = false; 227: optional bool new_version_bitmap_op_count = false; 228: optional bool enable_local_exchange_before_streaming_agg = false; + // FE is the receiver of fragment reports, so BE must also honor its message limit. + 229: optional i32 coordinator_thrift_max_message_size; + // FE can explicitly and idempotently acknowledge external-file commit reports. + 230: optional bool supports_external_file_report_ack = false; // For cloud, to control if the content would be written into file cache // In write path, to control if the content would be written into file cache. // In read path, read from file cache or remote storage when execute query.