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..8b09a0af4dfbfc --- /dev/null +++ b/be/src/exec/operator/iceberg_sorter_reserve_memory.h @@ -0,0 +1,132 @@ +// 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 retained_growth_trigger_bytes = 0; + size_t transient_workspace = 0; +}; + +inline size_t iceberg_saturating_add(size_t lhs, size_t rhs) { + return std::min(std::numeric_limits::max() - lhs, rhs) + lhs; +} + +inline size_t bounded_iceberg_reserve_size( + const std::vector& per_partition_reservations, + size_t incoming_rows = std::numeric_limits::max(), + size_t incoming_bytes = std::numeric_limits::max()) { + size_t transient_workspace = 0; + for (const auto& reservation : per_partition_reservations) { + transient_workspace = std::max(transient_workspace, reservation.transient_workspace); + } + + std::vector growth_candidates; + growth_candidates.reserve(per_partition_reservations.size()); + for (const auto& reservation : per_partition_reservations) { + if (reservation.retained_growth > 0) { + growth_candidates.push_back(&reservation); + } + } + + std::sort(growth_candidates.begin(), growth_candidates.end(), + [](const auto* lhs, const auto* rhs) { + return lhs->retained_growth > rhs->retained_growth; + }); + size_t row_bound = 0; + for (size_t i = 0; i < std::min(incoming_rows, growth_candidates.size()); ++i) { + row_bound = iceberg_saturating_add(row_bound, growth_candidates[i]->retained_growth); + } + + size_t byte_bound = 0; + std::vector positive_trigger_candidates; + positive_trigger_candidates.reserve(growth_candidates.size()); + for (const auto* reservation : growth_candidates) { + if (reservation->retained_growth_trigger_bytes == 0) { + byte_bound = iceberg_saturating_add(byte_bound, reservation->retained_growth); + } else { + positive_trigger_candidates.push_back(reservation); + } + } + std::sort(positive_trigger_candidates.begin(), positive_trigger_candidates.end(), + [](const auto* lhs, const auto* rhs) { + return static_cast(lhs->retained_growth) * + rhs->retained_growth_trigger_bytes > + static_cast(rhs->retained_growth) * + lhs->retained_growth_trigger_bytes; + }); + size_t remaining_bytes = incoming_bytes; + for (const auto* reservation : positive_trigger_candidates) { + if (reservation->retained_growth_trigger_bytes <= remaining_bytes) { + byte_bound = iceberg_saturating_add(byte_bound, reservation->retained_growth); + remaining_bytes -= reservation->retained_growth_trigger_bytes; + continue; + } + const auto numerator = + static_cast(reservation->retained_growth) * remaining_bytes + + reservation->retained_growth_trigger_bytes - 1; + const auto fractional_growth = + std::min(numerator / reservation->retained_growth_trigger_bytes, + std::numeric_limits::max()); + byte_bound = iceberg_saturating_add(byte_bound, static_cast(fractional_growth)); + break; + } + + // A block's rows and bytes are divided across partition sorters. The two fractional-relaxation + // bounds avoid charging the complete input block to every active partition while remaining safe. + const size_t retained_growth = std::min(row_bound, byte_bound); + return iceberg_saturating_add(retained_growth, transient_workspace); +} + +inline size_t iceberg_reserve_size( + const std::vector& per_partition_reservations, + size_t incoming_block_reserve, size_t incoming_rows = std::numeric_limits::max(), + size_t incoming_bytes = std::numeric_limits::max()) { + size_t sorter_reserve = + bounded_iceberg_reserve_size(per_partition_reservations, incoming_rows, incoming_bytes); + // The incoming block creates cold partition writers before they can appear in the published snapshot. + return iceberg_saturating_add(sorter_reserve, incoming_block_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..c7e4615568d127 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 dispatch_copies = block_bytes > std::numeric_limits::max() / 4 + ? std::numeric_limits::max() + : block_bytes * 4; + size_t reserve = iceberg_saturating_add(writer_workspace_bytes, dispatch_copies); + // A transform, selected blocks, and retained sorter copies can coexist; one extra block-sized + // allowance covers allocator fragmentation when a block is split into many tiny partitions. + return iceberg_saturating_add(reserve, row_index_bytes); +} + 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,49 @@ 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; + const size_t incoming_rows = block == nullptr ? 0 : block->rows(); + const size_t incoming_bytes = block == nullptr ? 0 : block->allocated_bytes(); + 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, incoming_rows, incoming_bytes); + per_partition_reservations.push_back( + {.retained_growth = reservation.retained_growth, + .retained_growth_trigger_bytes = reservation.retained_growth_trigger_bytes, + .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, incoming_rows, + incoming_bytes); } 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 +118,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 +166,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 84840e98f1f799..2fa064c8a68e09 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" @@ -468,6 +470,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()); @@ -2348,6 +2352,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; @@ -2359,6 +2372,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; @@ -2379,6 +2397,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; } @@ -2502,45 +2524,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()); @@ -2549,8 +2533,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(); @@ -2563,12 +2559,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; } @@ -2577,14 +2580,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); } } @@ -2631,13 +2658,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 206f6adc8445d7..46c6cdf159cbf9 100644 --- a/be/src/exec/sink/writer/async_result_writer.cpp +++ b/be/src/exec/sink/writer/async_result_writer.cpp @@ -44,6 +44,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 @@ -55,9 +57,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(); } @@ -71,17 +76,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) { @@ -131,6 +150,12 @@ void AsyncResultWriter::process_block(RuntimeState* state, RuntimeProfile* opera } DCHECK(_dependency); + bool reservation_held_for_finalize = false; + Defer release_final_reservation {[&]() { + if (reservation_held_for_finalize) { + thread_context()->thread_mem_tracker_mgr->shrink_reserved(); + } + }}; while (_writer_status.ok()) { ThreadCpuStopWatch cpu_time_stop_watch; cpu_time_stop_watch.start(); @@ -159,24 +184,49 @@ 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) { + // Some writers finalize buffered data in close(), so the EOS reservation must outlive + // both finish() and close() instead of being released between the two callbacks. + reservation_held_for_finalize = 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; 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..0d4653400e6530 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,27 @@ 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 remove any published file. + 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 +160,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..bfcce657afeb1d 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,44 @@ 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); + _include_spill_merge_reservation(state, eos, &reservation); + return reservation; +} + +SorterReserveMemory VIcebergSortWriter::get_reserve_mem_size_components( + RuntimeState* state, bool eos, size_t incoming_rows, size_t incoming_bytes) const { + std::lock_guard lock(_sorter_mutex); + if (_sorter == nullptr) { + return {}; + } + auto reservation = + _sorter->get_reserve_mem_size_components(state, eos, incoming_rows, incoming_bytes); + _include_spill_merge_reservation(state, eos, &reservation); + return reservation; +} + +void VIcebergSortWriter::_include_spill_merge_reservation(RuntimeState* state, bool eos, + SorterReserveMemory* reservation) const { + 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); + } +} + Status VIcebergSortWriter::trigger_spill() { std::lock_guard lock(_sorter_mutex); if (_closed || _sorter == nullptr) { @@ -102,80 +141,35 @@ Status VIcebergSortWriter::close(const Status& status) { } Status VIcebergSortWriter::_close_locked(const Status& status) { - // Track the actual internal status of operations performed during close. - // This is important because if intermediate operations (like do_sort()) fail, - // we need to propagate the actual error status to the underlying partition writer's - // close() call, rather than the original status parameter which could be OK. - Status internal_status = Status::OK(); - // Track the close status of the underlying partition writer. - // If _iceberg_partition_writer->close() fails (e.g., Parquet file flush error), - // we must propagate this error to the caller to avoid silent data loss. - Status close_status = Status::OK(); - - // Defer ensures the underlying partition writer is always closed and - // spill streams are cleaned up, regardless of whether intermediate operations succeed. - // Uses internal_status to propagate any errors that occurred during close operations. - Defer defer {[&]() { - // If any intermediate operation failed, pass that error to the partition writer; - // otherwise, pass the original status from the caller. - close_status = - _iceberg_partition_writer->close(internal_status.ok() ? status : internal_status); - if (!close_status.ok()) { - LOG(WARNING) << fmt::format("_iceberg_partition_writer close failed, reason: {}", - close_status.to_string()); - } - _cleanup_spill_streams(); - }}; - - // If the original status is already an error or the query is cancelled, - // skip all close operations and propagate the original error - if (!status.ok() || _runtime_state->is_cancelled()) { - return status; - } - - // If sorter was never initialized (e.g., no data was written), nothing to do - if (_sorter == nullptr) { - return Status::OK(); - } - - // Check if there is any remaining data in the sorter (either unsorted or already sorted blocks) - if (!_sorter->merge_sort_state()->unsorted_block()->empty() || - !_sorter->merge_sort_state()->get_sorted_block().empty()) { - if (_sorted_spill_files.empty()) { - // No spill has occurred, all data is in memory. - // Sort the remaining data, prepare for reading, and write to file. - internal_status = _sorter->do_sort(); - if (!internal_status.ok()) { - return internal_status; - } - internal_status = _sorter->prepare_for_read(false); - if (!internal_status.ok()) { - return internal_status; + Status internal_status = status; + if (status.ok() && !_runtime_state->is_cancelled()) { + internal_status = Status::OK(); + if (_sorter != nullptr && (!_sorter->merge_sort_state()->unsorted_block()->empty() || + !_sorter->merge_sort_state()->get_sorted_block().empty())) { + if (_sorted_spill_files.empty()) { + internal_status = _sorter->do_sort(); + if (internal_status.ok()) { + internal_status = _sorter->prepare_for_read(false); + } + if (internal_status.ok()) { + internal_status = _write_sorted_data(); + } + } else { + internal_status = _do_spill(); } - internal_status = _write_sorted_data(); - return internal_status; } - - // Some data has already been spilled to disk. - // Spill the remaining in-memory data to a new spill stream. - internal_status = _do_spill(); - if (!internal_status.ok()) { - return internal_status; + if (internal_status.ok() && !_sorted_spill_files.empty()) { + internal_status = _combine_files_output(); } } - // Merge all spilled streams using multi-way merge sort and output final sorted data to files - if (!_sorted_spill_files.empty()) { - internal_status = _combine_files_output(); - if (!internal_status.ok()) { - return internal_status; - } + // Form the return value only after the underlying close runs; a deferred assignment is too late. + Status close_status = + _iceberg_partition_writer->close(internal_status.ok() ? status : internal_status); + _cleanup_spill_streams(); + if (!internal_status.ok()) { + return internal_status; } - - // Return close_status if internal operations succeeded but the underlying - // partition writer's close() failed (e.g., file flush error). - // This prevents silent data loss where the caller thinks the write succeeded - // but the file was not properly closed. return close_status; } diff --git a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h index e1e512f0a0cf79..b41c31828431f1 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h +++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h @@ -105,10 +105,19 @@ 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; + + SorterReserveMemory get_reserve_mem_size_components(RuntimeState* state, bool eos, + size_t incoming_rows, + size_t incoming_bytes) const; + // Called by the memory management system to trigger spilling data to disk Status trigger_spill(); private: + void _include_spill_merge_reservation(RuntimeState* state, bool eos, + SorterReserveMemory* reservation) const; + // Calculate average row size from the first non-empty block to determine // the optimal batch row count for spill operations void _update_spill_block_batch_row_count(const Block& block); 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 3519252d7ca940..5b7c31d6c0d9cb 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp @@ -31,6 +31,7 @@ #include "core/data_type/data_type_struct.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" @@ -49,12 +50,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; @@ -329,7 +333,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 @@ -346,6 +351,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()); @@ -355,6 +361,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); @@ -363,6 +370,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; } @@ -371,7 +379,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(); } @@ -389,6 +396,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()); @@ -398,6 +406,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); @@ -406,6 +415,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; } @@ -414,7 +424,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(); } @@ -466,10 +475,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(); @@ -478,8 +485,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( @@ -488,7 +495,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(); @@ -498,53 +504,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) { @@ -564,6 +570,7 @@ Status VIcebergTableWriter::close(Status status) { } } _partitions_to_writers.clear(); + _publish_active_writers(); } if (status.ok()) { SCOPED_TIMER(_operator_profile->total_time_counter()); @@ -579,7 +586,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; } @@ -590,11 +597,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 20d9562c35ee52..0d6bede9a1a783 100644 --- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h +++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h @@ -73,23 +73,17 @@ 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); FRIEND_TEST(VIcebergTableWriterTest, ResolvesNestedPartitionSource); + // 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, @@ -150,12 +144,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..40d7b38fc30236 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)); @@ -209,7 +224,6 @@ void VHivePartitionWriter::_add_s3_mpu_pending_upload_for_rollback() { if (!_build_s3_mpu_pending_upload(&s3_mpu_pending_upload)) { return; } - THivePartitionUpdate hive_partition_update; hive_partition_update.__set_name(_partition_name); hive_partition_update.__set_update_mode(_update_mode); 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..5f878ecd7279ea 100644 --- a/be/src/exec/sort/sorter.cpp +++ b/be/src/exec/sort/sorter.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,20 @@ #include "runtime/runtime_profile.h" #include "runtime/thread_context.h" +namespace { + +size_t saturating_add_size(size_t lhs, size_t rhs) { + return std::min(std::numeric_limits::max() - lhs, rhs) + lhs; +} + +size_t saturating_multiply_size(size_t lhs, size_t rhs) { + return lhs == 0 || rhs <= std::numeric_limits::max() / lhs + ? lhs * rhs + : std::numeric_limits::max(); +} + +} // namespace + namespace doris { class RowDescriptor; } // namespace doris @@ -202,35 +217,62 @@ 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 { + const auto rows = _state->unsorted_block()->rows(); + const auto bytes = _state->unsorted_block()->bytes(); + const auto bytes_per_row = rows == 0 ? 0 : bytes / rows; + return get_reserve_mem_size_components( + state, eos, state->batch_size(), + saturating_multiply_size(bytes_per_row, state->batch_size())); +} + +SorterReserveMemory FullSorter::get_reserve_mem_size_components(RuntimeState* state, bool eos, + size_t incoming_rows, + size_t incoming_bytes) const { + SorterReserveMemory reserve; const auto rows = _state->unsorted_block()->rows(); if (rows != 0) { const auto bytes = _state->unsorted_block()->bytes(); const auto allocated_bytes = _state->unsorted_block()->allocated_bytes(); - const auto bytes_per_row = bytes / rows; - const auto estimated_size_of_next_block = bytes_per_row * state->batch_size(); - auto new_block_bytes = estimated_size_of_next_block + bytes; - auto new_rows = rows + state->batch_size(); + auto new_block_bytes = saturating_add_size(bytes, incoming_bytes); + auto new_rows = saturating_add_size(rows, incoming_rows); // 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); + const auto growth_threshold = static_cast( + (static_cast(allocated_bytes) * 85 + 99) / 100); + const size_t growth_trigger_bytes = growth_threshold > bytes ? growth_threshold - bytes : 0; + if (incoming_rows > 0 && growth_trigger_bytes <= incoming_bytes) { + reserve.retained_growth = static_cast(std::min( + (static_cast(allocated_bytes) * 115 + 99) / 100, + std::numeric_limits::max())); + reserve.retained_growth_trigger_bytes = growth_trigger_bytes; } - auto sort = new_rows > _buffered_block_size || new_block_bytes > _buffered_block_bytes; + // Iceberg close forces every nonempty pending run to sort at EOS, even when the generic + // append thresholds are not reached, so admission must cover that final allocation too. + auto sort = (eos && new_rows > 0) || 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(); + // sort_block keeps the source columns live while materializing a fully permuted destination. + reserve.transient_workspace = + saturating_add_size(reserve.transient_workspace, new_block_bytes); // helping data structures used during sorting - size_to_reserve += new_rows * sizeof(IColumn::Permutation::value_type); + reserve.transient_workspace = saturating_add_size( + reserve.transient_workspace, + saturating_multiply_size(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 = saturating_add_size( + reserve.transient_workspace, + saturating_multiply_size(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..2a47438f0b3e6e 100644 --- a/be/src/exec/sort/sorter.h +++ b/be/src/exec/sort/sorter.h @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -39,6 +40,18 @@ #include "runtime/runtime_state.h" namespace doris { + +struct SorterReserveMemory { + size_t retained_growth = 0; + size_t retained_growth_trigger_bytes = 0; + size_t transient_workspace = 0; + + size_t total() const { + return retained_growth > std::numeric_limits::max() - transient_workspace + ? std::numeric_limits::max() + : retained_growth + transient_workspace; + } +}; class ObjectPool; class RowDescriptor; } // namespace doris @@ -194,6 +207,12 @@ 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; + + SorterReserveMemory get_reserve_mem_size_components(RuntimeState* state, bool eos, + size_t incoming_rows, + size_t incoming_bytes) 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..67cb690a04fdba 100644 --- a/be/src/io/fs/azure_obj_storage_client.cpp +++ b/be/src/io/fs/azure_obj_storage_client.cpp @@ -35,10 +35,12 @@ #include #include #include +#include #include #include #include #include +#include #include "common/exception.h" #include "common/logging.h" @@ -46,8 +48,8 @@ #include "cpp/obj_retry_strategy.h" #include "io/fs/obj_storage_client.h" #include "util/bvar_helper.h" -#include "util/coding.h" #include "util/s3_util.h" +#include "util/uid_util.h" using namespace Azure::Storage::Blobs; @@ -64,10 +66,16 @@ 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) { + // Keep the full upload UUID in every block ID so independent writers cannot stage the + // same block IDs even though Azure has no per-upload multipart namespace. + std::vector raw_id(upload_id.begin(), upload_id.end()); + auto part = static_cast(part_num); + for (size_t i = 0; i < sizeof(part); ++i) { + raw_id.push_back(static_cast(part >> (i * 8))); + } + Aws::Utils::ByteBuffer bytes(raw_id.data(), raw_id.size()); + return Aws::Utils::HashingUtils::Base64Encode(bytes); } // Rate limiting is applied by RateLimitedObjStorageClient, the decorator that @@ -79,6 +87,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 +206,12 @@ struct AzureBatchDeleter { std::vector> deferred_resps; }; -// Azure would do nothing ObjectStorageUploadResponse AzureObjStorageClient::create_multipart_upload( - const ObjectStoragePathOptions& opts) { + const ObjectStoragePathOptions&) { + // Azure has no multipart session; this local UUID only namespaces the writer's block IDs. return ObjectStorageUploadResponse { .resp = ObjectStorageResponse::OK(), + .upload_id = generate_uuid_string(), }; } @@ -216,33 +229,40 @@ ObjectStorageResponse AzureObjStorageClient::put_object(const ObjectStoragePathO ObjectStorageUploadResponse AzureObjStorageClient::upload_part(const ObjectStoragePathOptions& opts, std::string_view stream, int part_num) { + DCHECK(opts.upload_id.has_value()); auto client = _client->GetBlockBlobClient(opts.key); + 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); + 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_client = _client->GetBlockBlobClient(opts.key); 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); }); + 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); + }); return 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. + target_client.CommitBlockList(string_block_ids); }, opts, _tls_debug_context); } diff --git a/be/src/io/fs/azure_obj_storage_client.h b/be/src/io/fs/azure_obj_storage_client.h index 7d1cecc502e44d..6cf6493e082af8 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: diff --git a/be/src/io/fs/obj_storage_client.h b/be/src/io/fs/obj_storage_client.h index fa239ca3282e2a..2688ee70e02a4b 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; // token identifying this writer's parts }; struct ObjectCompleteMultiPart { @@ -86,7 +86,8 @@ 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 token may be provider-issued or local and + // identifies this writer's parts. // The input parameters should include the bucket and key for the object storage. virtual ObjectStorageUploadResponse create_multipart_upload( const ObjectStoragePathOptions& opts) = 0; diff --git a/be/src/io/fs/s3_file_writer.cpp b/be/src/io/fs/s3_file_writer.cpp index f8b836607a14a6..a85aa5ce405ee9 100644 --- a/be/src/io/fs/s3_file_writer.cpp +++ b/be/src/io/fs/s3_file_writer.cpp @@ -85,7 +85,8 @@ 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. + // Deferred uploads are reported to FE for cleanup. Uploads that never reach FE are left to + // the provider lifecycle policy, so destroying a writer must not mutate provider state here. if (state() == State::OPENED && !_failed) { s3_bytes_written_total << _bytes_appended; } diff --git a/be/src/io/fs/s3_file_writer.h b/be/src/io/fs/s3_file_writer.h index 5a8075e03cf404..83ec75c9184920 100644 --- a/be/src/io/fs/s3_file_writer.h +++ b/be/src/io/fs/s3_file_writer.h @@ -75,7 +75,6 @@ class S3FileWriter final : public FileWriter { private: Status _close_impl(); - Status _abort(); [[nodiscard]] std::string _dump_completed_part() const; void _wait_until_finish(std::string_view task_name); Status _complete(); diff --git a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp index 5de03560c289d2..f86e155f9acd44 100644 --- a/be/src/runtime/memory/thread_mem_tracker_mgr.cpp +++ b/be/src/runtime/memory/thread_mem_tracker_mgr.cpp @@ -27,6 +27,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 d9a915c439bc62..b4494997ca7cd1 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..380bc8f8f72081 100644 --- a/be/src/runtime/runtime_state.cpp +++ b/be/src/runtime/runtime_state.cpp @@ -52,12 +52,99 @@ #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) { + // Once an ACK can have been lost, a later rejection cannot prove FE never accepted the files. + _external_file_report_state->ownership_may_have_transferred = true; + return; + } + if (_external_file_report_state->ownership_may_have_transferred) { + 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 cf844dd8c45dcb..a3cfc5e4cad782 100644 --- a/be/src/runtime/runtime_state.h +++ b/be/src/runtime/runtime_state.h @@ -76,6 +76,21 @@ class RuntimeFilterConsumer; class RuntimeFilterProducer; class TaskExecutionContext; +// Keep RuntimeState self-contained without importing the full frontend Thrift service header. +class TReportExecStatusParams; + +class ExternalFileReportState { + friend class RuntimeState; + +private: + std::mutex mutex; + size_t iceberg_serialized_bytes = 0; + bool ownership_may_have_transferred = false; + 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 { @@ -525,14 +540,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 { @@ -978,6 +1006,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..3e004051220504 --- /dev/null +++ b/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp @@ -0,0 +1,146 @@ +// 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, BoundsRetainedGrowthByOneInputBlock) { + constexpr size_t MB = 1024 * 1024; + std::vector per_partition_reservations( + 128, {.retained_growth = 8 * MB, + .retained_growth_trigger_bytes = 8 * MB, + .transient_workspace = 4 * MB}); + + // Only one 8 MiB growth threshold can be crossed by this block. Treating the block as a full + // batch for every active partition would incorrectly reserve more than 1 GiB here. + EXPECT_EQ(12 * MB, bounded_iceberg_reserve_size(per_partition_reservations, 128, 8 * MB)); +} + +TEST(SpillIcebergTableSinkOperatorTest, RetainsNearCapacityGrowthAcrossAllPossiblePartitions) { + constexpr size_t MB = 1024 * 1024; + std::vector per_partition_reservations( + 4, {.retained_growth = 3 * MB, + .retained_growth_trigger_bytes = 0, + .transient_workspace = 2 * MB}); + + // A one-row append can grow every already-near-capacity sorter, so persistent growth remains + // cumulative even though the serially used workspace is shared. + EXPECT_EQ(14 * MB, bounded_iceberg_reserve_size(per_partition_reservations, 4, 1)); +} + +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_GE(iceberg_cold_writer_reserve_size(block, operator_floor), + 4 * block.allocated_bytes() + operator_floor); +} + +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/viceberg_merge_sink_test.cpp b/be/test/exec/sink/viceberg_merge_sink_test.cpp index eb7c0159d5fd38..7eac5823d49725 100644 --- a/be/test/exec/sink/viceberg_merge_sink_test.cpp +++ b/be/test/exec/sink/viceberg_merge_sink_test.cpp @@ -47,6 +47,16 @@ namespace doris { +class IcebergWriteMockRuntimeState : public MockRuntimeState { +public: + IcebergWriteMockRuntimeState() { + auto query_options = this->query_options(); + // Successful writer tests must model a coordinator that can accept file ownership. + query_options.__set_supports_external_file_report_ack(true); + set_query_options(query_options); + } +}; + class VIcebergMergeSinkTest : public testing::Test { protected: static std::string test_schema_json() { @@ -172,7 +182,7 @@ class VIcebergMergeSinkTest : public testing::Test { TEST_F(VIcebergMergeSinkTest, TestUpdateProducesDeleteAndInsert) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -202,7 +212,7 @@ TEST_F(VIcebergMergeSinkTest, TestUpdateProducesDeleteAndInsert) { TEST_F(VIcebergMergeSinkTest, TestMissingOperationColumn) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -225,7 +235,7 @@ TEST_F(VIcebergMergeSinkTest, TestMissingOperationColumn) { TEST_F(VIcebergMergeSinkTest, TestMissingRowIdColumn) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -248,7 +258,7 @@ TEST_F(VIcebergMergeSinkTest, TestMissingRowIdColumn) { TEST_F(VIcebergMergeSinkTest, TestUnknownOperation) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -275,7 +285,7 @@ TEST_F(VIcebergMergeSinkTest, TestUnknownOperation) { TEST_F(VIcebergMergeSinkTest, TestUpdateInsertAndDeleteOperations) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -305,7 +315,7 @@ TEST_F(VIcebergMergeSinkTest, TestUpdateInsertAndDeleteOperations) { TEST_F(VIcebergMergeSinkTest, TestSchemaMismatch) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -332,7 +342,7 @@ TEST_F(VIcebergMergeSinkTest, TestSchemaMismatch) { TEST_F(VIcebergMergeSinkTest, TestRejectsDuplicateMatchedTargetAcrossBlocks) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -361,7 +371,7 @@ TEST_F(VIcebergMergeSinkTest, TestRejectsDuplicateMatchedTargetAcrossBlocks) { TEST_F(VIcebergMergeSinkTest, TestUpdateSkipsCardinalityState) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -389,7 +399,7 @@ TEST_F(VIcebergMergeSinkTest, TestUpdateSkipsCardinalityState) { TEST_F(VIcebergMergeSinkTest, TestOldFePlanSkipsCardinalityState) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -413,7 +423,7 @@ TEST_F(VIcebergMergeSinkTest, TestOldFePlanSkipsCardinalityState) { TEST_F(VIcebergMergeSinkTest, TestRollingUpgradeSkipsCardinalityState) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; state.set_be_exec_version(SUPPORT_ICEBERG_MERGE_CARDINALITY_VERSION - 1); DataTypes types {std::make_shared(), @@ -437,7 +447,7 @@ TEST_F(VIcebergMergeSinkTest, TestRollingUpgradeSkipsCardinalityState) { TEST_F(VIcebergMergeSinkTest, TestErrorCloseRemovesRolledDataFiles) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -469,7 +479,7 @@ TEST_F(VIcebergMergeSinkTest, TestErrorCloseRemovesRolledDataFiles) { TEST_F(VIcebergMergeSinkTest, TestDeleteCloseFailureRemovesBothInnerSinkFiles) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -519,7 +529,7 @@ TEST_F(VIcebergMergeSinkTest, TestDeleteCloseFailureRemovesBothInnerSinkFiles) { TEST_F(VIcebergMergeSinkTest, TestMatchedRowIdsUseCompactRetainedState) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), @@ -548,7 +558,7 @@ TEST_F(VIcebergMergeSinkTest, TestMatchedRowIdsUseCompactRetainedState) { TEST_F(VIcebergMergeSinkTest, TestMatchedRowIdStateAcrossManyFilesAndWrites) { ObjectPool pool; - MockRuntimeState state; + IcebergWriteMockRuntimeState state; DataTypes types {std::make_shared(), std::make_shared(DataTypes {std::make_shared(), diff --git a/be/test/exec/sink/writer/async_result_writer_test.cpp b/be/test/exec/sink/writer/async_result_writer_test.cpp new file mode 100644 index 00000000000000..18339b818de3c8 --- /dev/null +++ b/be/test/exec/sink/writer/async_result_writer_test.cpp @@ -0,0 +1,185 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "exec/sink/writer/async_result_writer.h" + +#include + +#include "core/block/block.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_number.h" +#include "exec/pipeline/dependency.h" +#include "runtime/exec_env.h" +#include "runtime/fragment_mgr.h" +#include "runtime/memory/global_memory_arbitrator.h" +#include "runtime/memory/mem_tracker_limiter.h" +#include "runtime/thread_context.h" +#include "runtime/workload_management/resource_context.h" + +namespace doris { + +namespace { + +const VExprContextSPtrs EMPTY_OUTPUT_EXPRS; + +Block make_block() { + auto values = ColumnInt32::create(); + values->insert_value(1); + Block block; + block.insert({std::move(values), std::make_shared(), "value"}); + return block; +} + +class RecordingAsyncWriter final : public AsyncResultWriter { +public: + RecordingAsyncWriter(std::shared_ptr dependency, + std::shared_ptr finish_dependency, Status open_status) + : AsyncResultWriter(EMPTY_OUTPUT_EXPRS, std::move(dependency), + std::move(finish_dependency)), + _open_status(std::move(open_status)) {} + + Status open(RuntimeState*, RuntimeProfile*) override { return _open_status; } + + Status write(RuntimeState*, Block&) override { + reservation_seen_by_write = thread_context()->thread_mem_tracker_mgr->reserved_mem(); + return Status::OK(); + } + + Status finish(RuntimeState*) override { + reservation_seen_by_finish = thread_context()->thread_mem_tracker_mgr->reserved_mem(); + return Status::OK(); + } + + Status close(Status) override { + reservation_seen_by_close = thread_context()->thread_mem_tracker_mgr->reserved_mem(); + return Status::OK(); + } + + int64_t reservation_seen_by_write = 0; + int64_t reservation_seen_by_finish = 0; + int64_t reservation_seen_by_close = 0; + +private: + Status _open_status; +}; + +struct AsyncWriterHarness { + AsyncWriterHarness() + : dependency(std::make_shared(0, 0, "writer", true)), + finish_dependency(std::make_shared(0, 0, "finish", false)), + common_profile("CommonCounters"), + memory_usage(common_profile.AddHighWaterMarkCounter("MemoryUsage", TUnit::BYTES)) {} + + void prepare(AsyncResultWriter* writer) { + writer->_operator_profile = &operator_profile; + writer->_memory_used_counter = memory_usage; + } + + void process(AsyncResultWriter* writer) { writer->process_block(nullptr, &operator_profile); } + + std::shared_ptr dependency; + std::shared_ptr finish_dependency; + RuntimeProfile operator_profile {"operator"}; + RuntimeProfile common_profile; + RuntimeProfile::Counter* memory_usage; +}; + +} // namespace + +class AsyncResultWriterTest : public testing::Test { +protected: + void SetUp() override { + _exec_env = ExecEnv::GetInstance(); + if (_exec_env->fragment_mgr() == nullptr) { + _fragment_mgr = std::make_unique(_exec_env); + _exec_env->_fragment_mgr = _fragment_mgr.get(); + } + _tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "UT-AsyncResultWriterReservation"); + _resource_context = ResourceContext::create_shared(); + _resource_context->memory_context()->set_mem_tracker(_tracker); + thread_context()->attach_task(_resource_context); + } + + void TearDown() override { + thread_context()->thread_mem_tracker_mgr->shrink_reserved(); + thread_context()->detach_task(); + EXPECT_EQ(0, GlobalMemoryArbitrator::process_reserved_memory()); + if (_fragment_mgr != nullptr) { + _fragment_mgr->stop(); + _exec_env->_fragment_mgr = nullptr; + _fragment_mgr.reset(); + } + } + + std::shared_ptr _tracker; + std::shared_ptr _resource_context; + ExecEnv* _exec_env = nullptr; + std::unique_ptr _fragment_mgr; +}; + +TEST_F(AsyncResultWriterTest, TransfersQueuedReservationIntoActualWrite) { + AsyncWriterHarness harness; + RecordingAsyncWriter writer(harness.dependency, harness.finish_dependency, Status::OK()); + harness.prepare(&writer); + constexpr int64_t reservation = 4 * 1024 * 1024; + Block block = make_block(); + ASSERT_TRUE(thread_context()->thread_mem_tracker_mgr->try_reserve(reservation).ok()); + + ASSERT_TRUE(writer.sink(&block, true).ok()); + EXPECT_EQ(0, thread_context()->thread_mem_tracker_mgr->reserved_mem()); + harness.process(&writer); + + EXPECT_GT(writer.reservation_seen_by_write, 0); + EXPECT_LE(writer.reservation_seen_by_write, reservation); + EXPECT_EQ(0, thread_context()->thread_mem_tracker_mgr->reserved_mem()); +} + +TEST_F(AsyncResultWriterTest, RetainsEosReservationThroughActualClose) { + AsyncWriterHarness harness; + RecordingAsyncWriter writer(harness.dependency, harness.finish_dependency, Status::OK()); + harness.prepare(&writer); + constexpr int64_t reservation = 4 * 1024 * 1024; + ASSERT_TRUE(thread_context()->thread_mem_tracker_mgr->try_reserve(reservation).ok()); + Block block; + + ASSERT_TRUE(writer.sink(&block, true).ok()); + harness.process(&writer); + + EXPECT_EQ(reservation, writer.reservation_seen_by_finish); + EXPECT_EQ(reservation, writer.reservation_seen_by_close); + EXPECT_EQ(0, thread_context()->thread_mem_tracker_mgr->reserved_mem()); +} + +TEST_F(AsyncResultWriterTest, OpenFailureDrainsQueuedReservation) { + AsyncWriterHarness harness; + RecordingAsyncWriter writer(harness.dependency, harness.finish_dependency, + Status::IOError("injected open failure")); + harness.prepare(&writer); + constexpr int64_t reservation = 4 * 1024 * 1024; + Block block = make_block(); + ASSERT_TRUE(thread_context()->thread_mem_tracker_mgr->try_reserve(reservation).ok()); + + ASSERT_TRUE(writer.sink(&block, true).ok()); + harness.process(&writer); + + EXPECT_FALSE(writer.get_writer_status().ok()); + EXPECT_EQ(0, thread_context()->thread_mem_tracker_mgr->reserved_mem()); + EXPECT_EQ(0, GlobalMemoryArbitrator::process_reserved_memory()); +} + +} // namespace doris diff --git a/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp b/be/test/exec/sink/writer/iceberg/iceberg_partition_writer_test.cpp index d453177cf25044..d5a7e3f9ab3351 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,10 @@ #include #include "exec/sink/writer/iceberg/viceberg_partition_writer.h" +#include "exec/sink/writer/iceberg/viceberg_sort_writer.h" +#include "testutil/mock/mock_descriptors.h" +#include "testutil/mock/mock_runtime_state.h" +#include "testutil/mock/mock_slot_ref.h" namespace doris { @@ -27,13 +31,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 +73,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 +114,48 @@ 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); +} + +TEST_F(VIcebergPartitionWriterTest, EosReservationIncludesActualSpillFanIn) { + 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)); + VIcebergSortWriter sort_writer(partition_writer, TSortInfo(), 1024); + MockRuntimeState state; + ObjectPool pool; + auto row_desc = std::make_unique( + std::vector {std::make_shared()}, &pool); + auto ordering_expr_ctxs = + MockSlotRef::create_mock_contexts(0, std::make_shared()); + std::vector is_asc_order {true}; + std::vector nulls_first {false}; + sort_writer._sorter = FullSorter::create_unique(ordering_expr_ctxs, -1, 0, &pool, is_asc_order, + nulls_first, *row_desc, &state, nullptr); + sort_writer._sorted_spill_files.resize(12); + + const auto reservation = sort_writer.get_reserve_mem_size_components(&state, true, 0, 0); + + EXPECT_EQ(72 * 1024 * 1024, reservation.transient_workspace); +} + } // 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..ac6c100b37e143 --- /dev/null +++ b/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp @@ -0,0 +1,180 @@ +// 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_string.h" +#include "core/column/column_vector.h" +#include "core/data_type/data_type_number.h" +#include "core/data_type/data_type_string.h" +#include "exec/operator/iceberg_sorter_reserve_memory.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, ColdReserveCoversManyRealPartitionSelections) { + VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr); + auto values = ColumnString::create(); + for (size_t i = 0; i < 128; ++i) { + const std::string value(256 + i, static_cast('a' + i % 26)); + values->insert_data(value.data(), value.size()); + } + Block input; + input.insert({std::move(values), std::make_shared(), "value"}); + size_t selected_bytes = 0; + for (size_t row = 0; row < input.rows(); ++row) { + Block selected; + ASSERT_TRUE(select_block(&writer, input, {row}, &selected).ok()); + selected_bytes += selected.allocated_bytes(); + } + + const size_t reserve = iceberg_cold_writer_reserve_size(input, 0); + EXPECT_GE(reserve, + input.allocated_bytes() + 2 * selected_bytes + input.rows() * sizeof(size_t)); +} + +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..73862ac8b0580d --- /dev/null +++ b/be/test/exec/sink/writer/vhive_partition_writer_report_lifecycle_test.cpp @@ -0,0 +1,217 @@ +// 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::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 {}; + } +}; + +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, + io::ObjStorageType provider = io::ObjStorageType::AWS, std::string staged_block_id = {}) { + 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); + + S3ClientConf client_conf; + client_conf.provider = provider; + auto holder = std::make_shared(client_conf); + 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"; + if (!staged_block_id.empty()) { + file_writer->_completed_parts.push_back( + {.part_num = 1, .etag = std::move(staged_block_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, + PeriodicReportDefersMetadataAndFinalReportTransfersUploadIdentity) { + 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); + + 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); +} + +TEST(VHivePartitionWriterReportLifecycleTest, AzureFinalReportCarriesExactBlockIdentity) { + MockRuntimeState state; + VExprContextSPtrs output_exprs; + auto client = std::make_shared(); + auto writer = create_closed_hive_writer(&state, output_exprs, client, io::ObjStorageType::AZURE, + "exact-block-id"); + auto context = create_fragment_context(TUniqueId()); + auto final_request = report_request(&state, true); + TReportExecStatusParams final_params; + + 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()); + const auto& pending_uploads = final_params.hive_partition_updates[0].s3_mpu_pending_uploads; + ASSERT_EQ(1, pending_uploads.size()); + EXPECT_EQ("upload-id", pending_uploads[0].upload_id); + EXPECT_EQ("exact-block-id", pending_uploads[0].etags.at(1)); +} + +} // namespace doris diff --git a/be/test/exec/sort/full_sort_test.cpp b/be/test/exec/sort/full_sort_test.cpp index e182048c807dad..bd8f91b03cc863 100644 --- a/be/test/exec/sort/full_sort_test.cpp +++ b/be/test/exec/sort/full_sort_test.cpp @@ -22,6 +22,7 @@ #include #include +#include #include #include #include @@ -63,6 +64,13 @@ struct FullSorterTest : public testing::Test { std::vector nulls_first {false}; }; +TEST(SorterReserveMemoryTest, TotalSaturatesOnOverflow) { + SorterReserveMemory reservation {.retained_growth = std::numeric_limits::max() - 1, + .transient_workspace = 2}; + + EXPECT_EQ(std::numeric_limits::max(), reservation.total()); +} + TEST_F(FullSorterTest, test_full_sorter1) { sorter = FullSorter::create_unique(ordering_expr_ctxs, -1, 0, &pool, is_asc_order, nulls_first, *row_desc, &_state, nullptr); @@ -94,6 +102,20 @@ TEST_F(FullSorterTest, test_full_sorter2) { std::cout << sorter->get_reserve_mem_size(&_state, false) << std::endl; } +TEST_F(FullSorterTest, EosReservationIncludesForcedSortBelowAppendThresholds) { + sorter = FullSorter::create_unique(ordering_expr_ctxs, -1, 0, &pool, is_asc_order, nulls_first, + *row_desc, &_state, nullptr); + Block block = ColumnHelper::create_block({10, 9, 8, 7, 6, 5, 4, 3, 2, 1}); + const size_t buffered_bytes = block.bytes(); + const size_t buffered_rows = block.rows(); + ASSERT_TRUE(sorter->append_block(&block).ok()); + + const auto reservation = sorter->get_reserve_mem_size_components(&_state, true, 0, 0); + + EXPECT_GE(reservation.transient_workspace, + buffered_bytes + buffered_rows * sizeof(IColumn::Permutation::value_type)); +} + TEST_F(FullSorterTest, test_full_sorter3) { sorter = FullSorter::create_unique(ordering_expr_ctxs, 3, 3, &pool, is_asc_order, nulls_first, *row_desc, &_state, nullptr); @@ -113,4 +135,4 @@ TEST_F(FullSorterTest, test_full_sorter3) { EXPECT_EQ(sorter->_state->get_sorted_block()[1]->rows(), 4); } -} // namespace doris \ No newline at end of file +} // 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 375ff3ff57f8e2..2745e87a5dd6e9 100644 --- a/be/test/io/client/s3_file_system_test.cpp +++ b/be/test/io/client/s3_file_system_test.cpp @@ -928,9 +928,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"; @@ -1437,9 +1436,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"; @@ -1866,8 +1864,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..a45db87bfe0268 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,17 @@ #include +#include +#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 +40,41 @@ namespace doris { #ifdef USE_AZURE +TEST(AzureObjStorageClientMultipartHelperTest, full_upload_uuid_isolates_writer_blocks) { + constexpr std::string_view first_upload = "09492e3d-e231-4ed9-bf84-b6fc772cda54"; + constexpr std::string_view second_upload = "06996d15-1c2e-4ddd-8853-43816ea84a07"; + auto first_block = io::azure_multipart_block_id(first_upload, 1); + auto second_block = io::azure_multipart_block_id(second_upload, 1); + + EXPECT_NE(first_block, second_block); + EXPECT_EQ(first_block.size(), io::azure_multipart_block_id(first_upload, 999).size()); + auto decoded = Aws::Utils::HashingUtils::Base64Decode(first_block); + ASSERT_EQ(first_upload.size() + sizeof(uint32_t), decoded.GetLength()); + EXPECT_EQ(first_upload, + std::string_view(reinterpret_cast(decoded.GetUnderlyingData()), + first_upload.size())); + EXPECT_EQ(1, decoded.GetUnderlyingData()[first_upload.size()]); + EXPECT_EQ(0, decoded.GetUnderlyingData()[first_upload.size() + 1]); + EXPECT_EQ(0, decoded.GetUnderlyingData()[first_upload.size() + 2]); + EXPECT_EQ(0, decoded.GetUnderlyingData()[first_upload.size() + 3]); +} + +TEST(AzureObjStorageClientMultipartHelperTest, create_upload_is_provider_free) { + io::AzureObjStorageClient client( + std::shared_ptr {}); + + auto first = client.create_multipart_upload({}); + auto second = client.create_multipart_upload({}); + + ASSERT_EQ(ErrorCode::OK, first.resp.status.code); + ASSERT_EQ(ErrorCode::OK, second.resp.status.code); + ASSERT_TRUE(first.upload_id.has_value()); + ASSERT_TRUE(second.upload_id.has_value()); + EXPECT_EQ(36, first.upload_id->size()); + EXPECT_EQ(36, second.upload_id->size()); + EXPECT_NE(first.upload_id, second.upload_id); +} + using namespace Azure::Storage::Blobs; TEST(AzureObjStorageClientTlsHelperTest, detects_tls_ca_error) { @@ -156,6 +197,39 @@ TEST_F(AzureObjStorageClientTest, delete_objects_recursively) { EXPECT_EQ(response.status.code, ErrorCode::OK); EXPECT_EQ(files.size(), 0); } + +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_EQ(second_create.resp.status.code, ErrorCode::OK); + ASSERT_TRUE(first_create.upload_id.has_value()); + ASSERT_TRUE(second_create.upload_id.has_value()); + ASSERT_NE(first_create.upload_id, second_create.upload_id); + first.upload_id = first_create.upload_id; + second.upload_id = second_create.upload_id; + + auto first_part = obj_storage_client->upload_part(first, "first", 1); + auto second_part = obj_storage_client->upload_part(second, "second", 1); + ASSERT_EQ(first_part.resp.status.code, ErrorCode::OK); + ASSERT_EQ(second_part.resp.status.code, ErrorCode::OK); + ASSERT_NE(first_part.etag, second_part.etag); + ASSERT_EQ(obj_storage_client->complete_multipart_upload(first, {{.part_num = 1}}).status.code, + ErrorCode::OK); + ASSERT_NE(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), "first"); + 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..edaeb80bd2bbfc 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 @@ -455,7 +455,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 +463,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/runtime/memory/thread_mem_tracker_mgr_test.cpp b/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp index 716e5f2e9a4958..188f16e96a8ee6 100644 --- a/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp +++ b/be/test/runtime/memory/thread_mem_tracker_mgr_test.cpp @@ -331,6 +331,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..5a384378ec382b 100644 --- a/be/test/runtime/runtime_state_block_budget_test.cpp +++ b/be/test/runtime/runtime_state_block_budget_test.cpp @@ -18,12 +18,131 @@ #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(1, cleanup_count); +} + +TEST(RuntimeStateIcebergCommitDataTest, AmbiguousOwnershipCannotBecomeRejected) { + RuntimeState state; + int cleanup_count = 0; + state.add_rejected_external_file_report_cleanup([&] { ++cleanup_count; }); + + state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::AMBIGUOUS); + state.finalize_external_file_report_cleanup(ExternalFileReportOutcome::REJECTED); + + EXPECT_EQ(0, 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 375604ecd4a9e1..e7a95fe6999c19 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 @@ -288,13 +288,23 @@ private ScheduledFuture startCommitLockHeartbeat(long lockId) { } private void commitWhileTableLocked() { - validateWriteMetadataBeforePublication(); - // 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); - // Classification can perform metastore reads, so close that interval before any file or HMS mutation. - validateWriteMetadataBeforePublication(); + try { + // Object-store files remain unpublished until FE consumes one completion record per file. + validateObjectStoreCommitRecords(); + validateWriteMetadataBeforePublication(); + // Classification can perform metastore reads, so validate once more before any publication. + finishInsertTable(nameMapping); + validateWriteMetadataBeforePublication(); + } 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()) { @@ -436,6 +446,25 @@ private boolean isTargetTable(String dbName, String tableName) { && nameMapping.getRemoteTblName().equalsIgnoreCase(tableName); } + 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) { @@ -608,15 +637,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 586d38f2436186..4ae66aae99c63b 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 @@ -295,6 +295,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 5b5ded04976ddb..e3fe52ba2eebe0 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 @@ -268,10 +268,13 @@ public void testCommitRejectsSameShapeTableReplacementBeforePublishing() throws client.table = table(false, Collections.emptyMap(), 10, Collections.singletonList(col("c1", "int"))); client.currentNotificationEventId = 10; - 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("", TUpdateMode.APPEND, "s3://bucket/db/t", - Collections.singletonList("f1"), 100, 4))); + // A complete ownership record is required so this test reaches the generation fence instead of the + // earlier fail-closed multipart validation. + txn.addCommitData(serialize(puWithMpu("", TUpdateMode.APPEND, "s3://bucket/db/t", + "bucket", "db/t/f1", "upload-replaced", Collections.singletonMap(1, "etag-1")))); // Recreate the table with every digest-visible field unchanged. Only the HMS event stream can // distinguish this new object from the table whose schema and location were bound above. @@ -292,10 +295,12 @@ public void testCommitHoldsExclusiveTableLockThroughPublication() throws TExcept RecordingHmsClient client = new RecordingHmsClient(); client.table = table(false, Collections.emptyMap()); client.currentNotificationEventId = 10; - 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("", TUpdateMode.APPEND, "s3://bucket/db/t", - Collections.singletonList("f1"), 100, 4))); + // Keep multipart validation green so the assertion measures the lock around actual publication. + txn.addCommitData(serialize(puWithMpu("", TUpdateMode.APPEND, "s3://bucket/db/t", + "bucket", "db/t/f1", "upload-lock", Collections.singletonMap(1, "etag-1")))); txn.commit(); @@ -459,6 +464,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 @@ -503,15 +571,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 eb549709a0a937..f6f8da32875ba2 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 @@ -329,6 +329,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/IcebergConnectorTransaction.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergConnectorTransaction.java index 21b301a3416490..e2209ea0cbd6cf 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; @@ -413,14 +413,14 @@ 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). + // (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(); - if (ctx.isReadSnapshotResolved()) { - // An explicitly empty read stays null so validation covers a concurrent first append. - this.baseSnapshotId = pinnedReadSnapshot >= 0 ? Long.valueOf(pinnedReadSnapshot) : null; - } else { - this.baseSnapshotId = getSnapshotIdIfPresent(table); - } + // 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 = ctx.isReadSnapshotResolved() + ? Long.valueOf(pinnedReadSnapshot) : getSnapshotIdIfPresent(table); if (table instanceof HasTableOperations) { int formatVersion = ((HasTableOperations) table).operations().current().formatVersion(); if (formatVersion < 2) { @@ -430,7 +430,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); @@ -440,12 +439,53 @@ 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.isReadSnapshotResolved()) { + 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 public long getTransactionId() { return transactionId; @@ -705,6 +745,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); @@ -724,6 +766,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."); @@ -749,6 +796,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, @@ -758,6 +806,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 09f8fa44615093..bfabf1c986adba 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 @@ -936,7 +936,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 { @@ -1157,6 +1157,9 @@ private TableScan buildScan(Table table, IcebergTableHandle handle, Optional= 0 || ref != null; } + /** Whether the pin selects an Iceberg snapshot/ref rather than the explicit empty-table state. */ + public boolean hasSnapshotSelection() { + return snapshotId >= 0 || ref != null; + } + /** Whether snapshot resolution observed a table before its first snapshot was committed. */ public boolean isResolvedEmptySnapshot() { return snapshotResolved && snapshotId < 0 && ref == null; 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)) { + // Unknown creation time cannot prove the file predates every in-flight writer. + if (scope.owns(file.location()) && file.createdAtMillis() > 0 + && 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/IcebergConnectorTransactionTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTransactionTest.java index cdc1a341f40cc1..13ae750395437d 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 @@ -34,6 +34,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; @@ -41,10 +42,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; @@ -138,6 +141,12 @@ private static IcebergWriteContext overwriteToBranch(String branch) { WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.of(branch)); } + private static IcebergWriteContext overwriteCtxPinned(long readSnapshotId) { + // A resolved empty read also uses -1, so the explicit flag preserves it as an OCC fence. + return new IcebergWriteContext(WriteOperation.OVERWRITE, true, Collections.emptyMap(), Optional.empty(), + readSnapshotId, true); + } + private static IcebergWriteContext overwriteStaticCtx(Table table, Map staticValues) { IcebergWriteSchemaContext schemaContext = IcebergWriteSchemaContext.create(table, table.name(), Optional.empty(), false, false); @@ -159,12 +168,12 @@ private static IcebergWriteContext rewriteCtx() { private static IcebergWriteContext deleteCtxPinned(long readSnapshotId) { return new IcebergWriteContext( - WriteOperation.DELETE, false, Collections.emptyMap(), Optional.empty(), readSnapshotId); + WriteOperation.DELETE, false, Collections.emptyMap(), Optional.empty(), readSnapshotId, true); } private static IcebergWriteContext mergeCtxPinned(long readSnapshotId) { return new IcebergWriteContext( - WriteOperation.MERGE, false, Collections.emptyMap(), Optional.empty(), readSnapshotId); + WriteOperation.MERGE, false, Collections.emptyMap(), Optional.empty(), readSnapshotId, true); } /** @@ -517,6 +526,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(); @@ -592,6 +617,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(); @@ -642,6 +799,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(); @@ -682,6 +857,53 @@ public void overwriteStaticPartitionRejectsUnmatchedPartitionField() { "a stale static overwrite must fail at begin before it can degrade to an always-true filter"); } + @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(); @@ -876,6 +1098,23 @@ public void mergeFromResolvedEmptySnapshotRejectsConcurrentFirstAppend() throws committedFiles.get(0).path().toString()); } + @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 26fa913062afea..5766ebe8c71cc4 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 @@ -319,6 +319,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(); + // Keep construction behind the shared helper so catalog-property API migrations do not break this test. + IcebergScanPlanProvider provider = providerOver(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..4492d097b96dd8 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,17 @@ 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.assertFalse(empty.hasSnapshotPin()); + Assertions.assertFalse(empty.hasSnapshotSelection()); + Assertions.assertTrue(empty.isSnapshotResolved()); + 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..b1f5de5c588ad3 --- /dev/null +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/action/IcebergRemoveOrphanFilesActionTest.java @@ -0,0 +1,468 @@ +// 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.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + +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 keepsOrphansWhenFileCreationTimeIsUnknown(@TempDir Path temp) throws Exception { + Table table = createTable(temp.resolve("table"), Collections.emptyMap()); + Path zeroTimestamp = createOldFile(temp.resolve("table/data/zero.parquet")); + Path negativeTimestamp = createOldFile(temp.resolve("table/data/negative.parquet")); + RecordingFileIO recordingFileIO = new RecordingFileIO(table.io(), Collections.emptySet()) { + @Override + public Iterable listPrefix(String prefix) { + List files = StreamSupport.stream(super.listPrefix(prefix).spliterator(), false) + .map(file -> { + long createdAtMillis = file.createdAtMillis(); + if (IcebergRemoveOrphanFilesAction.sameFileIdentity( + file.location(), negativeTimestamp.toUri().toString())) { + createdAtMillis = -1; + } else if (IcebergRemoveOrphanFilesAction.sameFileIdentity( + file.location(), zeroTimestamp.toUri().toString())) { + createdAtMillis = 0; + } + return new FileInfo(file.location(), file.size(), createdAtMillis); + }) + .collect(Collectors.toList()); + return files; + } + }; + Table recordingTable = new BaseTable( + new StaticTableOperations(((HasTableOperations) table).operations().current(), recordingFileIO), + table.name()); + IcebergRemoveOrphanFilesAction action = action( + System.currentTimeMillis() - MIN_RETENTION_MS, false); + action.validate(); + + ConnectorProcedureResult result = action.execute(recordingTable, ActionTestTables.session("UTC")); + + Assertions.assertEquals("0", result.getRows().get(0).get(0)); + Assertions.assertTrue(Files.exists(zeroTimestamp)); + Assertions.assertTrue(Files.exists(negativeTimestamp)); + } + + @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 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 09c8bc2462835e..c81869b6b0e57f 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; @@ -645,13 +644,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, connectorColumns, null, writeOperation, - requireMergeCardinalityCheck, boundWriteMetadataIdentity); + requireMergeCardinalityCheck, boundWriteMetadataIdentity, metadata); } @Override @@ -710,20 +707,9 @@ 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)); - // Preserve the generation captured from the exact remote table load that supplied the bound schema. // A live lookup here would silently move the fence after a concurrent drop/recreate. String boundWriteMetadataIdentity = connectorTableSink.getBoundWriteMetadataIdentity(); - // 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 @@ -743,7 +729,7 @@ public PlanFragment visitPhysicalConnectorTableSink( PluginDrivenTableSink providerSink = new PluginDrivenTableSink(targetTable, writePlanProvider, connSession, providerTableHandle, connectorColumns, boundTargetColumns, writeSortInfo, writeOperation, false, - boundWriteMetadataIdentity); + boundWriteMetadataIdentity, 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 79e53cd3f02968..a337477c963b0c 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; @@ -59,6 +63,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; private final List boundTargetColumns; // The engine-built BE sort instruction for a connector that declares write-sort columns (iceberg @@ -114,7 +119,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); } /** @@ -127,7 +132,20 @@ public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, ConnectorTableHandle tableHandle, List connectorColumns, TSortInfo writeSortInfo, WriteOperation writeOperation, boolean requireMergeCardinalityCheck) { this(targetTable, writePlanProvider, connectorSession, tableHandle, connectorColumns, - connectorColumns, writeSortInfo, writeOperation, requireMergeCardinalityCheck); + writeSortInfo, writeOperation, requireMergeCardinalityCheck, null); + } + + /** + * Plan-provider mode with connector metadata used to resolve the exact branch-aware MVCC pin at bind time. + */ + public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, + ConnectorWritePlanProvider writePlanProvider, ConnectorSession connectorSession, + ConnectorTableHandle tableHandle, List connectorColumns, + TSortInfo writeSortInfo, WriteOperation writeOperation, boolean requireMergeCardinalityCheck, + ConnectorMetadata connectorMetadata) { + this(targetTable, writePlanProvider, connectorSession, tableHandle, connectorColumns, + connectorColumns, writeSortInfo, writeOperation, requireMergeCardinalityCheck, + null, connectorMetadata); } /** @@ -151,11 +169,27 @@ public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, List boundTargetColumns, TSortInfo writeSortInfo, WriteOperation writeOperation, boolean requireMergeCardinalityCheck, String boundWriteMetadataIdentity) { + this(targetTable, writePlanProvider, connectorSession, tableHandle, connectorColumns, + boundTargetColumns, writeSortInfo, writeOperation, requireMergeCardinalityCheck, + boundWriteMetadataIdentity, null); + } + + /** + * Complete plan-provider mode. The schema generation and branch-aware snapshot are separate invariants: + * the former must stay at bind time, while the latter is resolved for the exact write branch. + */ + public PluginDrivenTableSink(PluginDrivenExternalTable targetTable, + ConnectorWritePlanProvider writePlanProvider, ConnectorSession connectorSession, + ConnectorTableHandle tableHandle, List connectorColumns, + List boundTargetColumns, TSortInfo writeSortInfo, + WriteOperation writeOperation, boolean requireMergeCardinalityCheck, + String boundWriteMetadataIdentity, ConnectorMetadata connectorMetadata) { super(); this.targetTable = targetTable; this.writePlanProvider = writePlanProvider; this.connectorSession = connectorSession; this.tableHandle = tableHandle; + this.connectorMetadata = connectorMetadata; this.connectorColumns = connectorColumns; // Keep this immutable bind-time snapshot distinct from the write subset. Re-reading the live // table here would move the conflict-detection baseline and could miss ordinal schema drift. @@ -221,8 +255,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, boundTargetColumns, overwrite, writeContext, writeSortInfo, + boundTableHandle, connectorColumns, boundTargetColumns, overwrite, writeContext, writeSortInfo, boundWriteMetadataIdentity, 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..06006086d4dfaa 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 @@ -253,7 +253,8 @@ public class Coordinator implements CoordInterface { private String trackingUrl; private String firstErrorMsg; // related txnId and label of group commit - private long txnId; + // Final reports race with status readers, so the transaction identity must be safely published. + private volatile long txnId; private String label; // for export @@ -2558,7 +2559,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 +2579,104 @@ 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(); + } + // Keep this report's identity local so another report cannot redirect its commit data. + long reportTxnId = params.isSetTxnId() ? params.getTxnId() : txnId; + if (params.isSetTxnId()) { + txnId = reportTxnId; + } + 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(reportTxnId); + 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; } /* @@ -3060,7 +3083,9 @@ public static class PipelineExecContext { TPipelineFragmentParams rpcParams; PlanFragmentId fragmentId; boolean initiated; - boolean done; + // Non-final reports read this outside the monitor after updatePipelineStatus returns. + volatile boolean done; + boolean processingDoneReport; TNetworkAddress brpcAddress; TNetworkAddress address; @@ -3117,10 +3142,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 cc144759f24455..f0445ae9192e00 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 @@ -5377,6 +5377,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 ba275fd20246dc..0b5a3b1ecd4950 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..2e15875c73ffd0 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java @@ -0,0 +1,174 @@ +// 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 com.google.common.cache.Cache; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import java.lang.reflect.Field; +import java.lang.reflect.Modifier; +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 evictedAcceptanceTokenRejectsRetryAfterCoordinatorRemoval() throws Exception { + TUniqueId queryId = new TUniqueId(12345, 5); + 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; + acceptedExternalFileReports().invalidateAll(); + TReportExecStatusResult retry = report(params); + + Assertions.assertTrue(first.isExternalFileCommitDataAccepted()); + Assertions.assertFalse(retry.isExternalFileCommitDataAccepted()); + Assertions.assertEquals(TStatusCode.INTERNAL_ERROR, retry.getStatus().getStatusCode()); + Mockito.verify(coordinator, Mockito.times(1)).updateFragmentExecStatus(params); + } + + @Test + void coordinatorReportStateUsesCrossThreadVisibility() throws Exception { + Field done = Coordinator.PipelineExecContext.class.getDeclaredField("done"); + Field txnId = Coordinator.class.getDeclaredField("txnId"); + + Assertions.assertTrue(Modifier.isVolatile(done.getModifiers())); + Assertions.assertTrue(Modifier.isVolatile(txnId.getModifiers())); + } + + @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)); + } + + @SuppressWarnings("unchecked") + private static Cache acceptedExternalFileReports() throws Exception { + Field field = QeProcessorImpl.class.getDeclaredField("acceptedExternalFileReports"); + field.setAccessible(true); + return (Cache) field.get(QeProcessorImpl.INSTANCE); + } +} 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 f1f0bd14033c47..9febc45b150e6d 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 @@ -398,4 +398,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..837e739abd184b 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 @@ -48,6 +48,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 +57,7 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.UUID; /** * Azure Blob Storage implementation of {@link ObjStorage}. @@ -218,9 +220,8 @@ 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; + // Azure has no multipart session; this local UUID only namespaces the writer's block IDs. + return UUID.randomUUID().toString(); } @Override @@ -230,7 +231,7 @@ public UploadPartResult uploadPart(String remotePath, String uploadId, int partN AzureUri uri = AzureUri.parse(remotePath); BlockBlobClient blockBlobClient = getClient().getBlobContainerClient(uri.container()) .getBlobClient(uri.key()).getBlockBlobClient(); - String blockId = toBlockId(partNum); + String blockId = multipartBlockId(uploadId, partNum); blockBlobClient.stageBlock(blockId, body.content(), body.contentLength()); return new UploadPartResult(partNum, blockId); } catch (BlobStorageException e) { @@ -244,15 +245,20 @@ 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()); } - blockBlobClient.commitBlockList(blockIds); + // Put Block List is the atomic publication point and does not expose a staging blob to scans. + containerClient.getBlobClient(uri.key()).getBlockBlobClient().commitBlockList(blockIds); } catch (BlobStorageException e) { throw new IOException("completeMultipartUpload failed for " + remotePath + ": " + e.getMessage(), e); @@ -261,41 +267,8 @@ 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). - 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(); - } 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()); - } + // Azure cannot selectively remove one writer's uncommitted blocks. They are isolated by + // upload UUID and left for the service to garbage-collect without touching published data. } /** @@ -516,12 +489,15 @@ private static String requireProperty(String value, String key, String descripti return value; } - /** - * Converts a part number to a Base64-encoded block ID using little-endian byte order, - * consistent with the existing fe-core Azure multipart upload implementation. - */ - 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) { + byte[] uploadBytes = uploadId.getBytes(StandardCharsets.UTF_8); + // Keep the full upload UUID in every block ID so independent writers cannot stage the + // same block IDs even though Azure has no per-upload multipart namespace. + byte[] rawId = ByteBuffer.allocate(uploadBytes.length + Integer.BYTES) + .order(ByteOrder.LITTLE_ENDIAN) + .put(uploadBytes) + .putInt(partNum) + .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..2945b141cef296 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,25 +17,32 @@ 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.BlobStorageException; +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.nio.charset.StandardCharsets; import java.time.OffsetDateTime; import java.util.Arrays; +import java.util.Base64; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.UUID; /** * Unit tests for the cloud extension methods added to {@link AzureObjStorage}: @@ -375,51 +382,127 @@ void deleteObjectsByKeys_attachesPerKeyExceptionsAsSuppressed() throws Exception } // ------------------------------------------------------------------ - // F20 — abortMultipartUpload safe-noop / commit-empty behaviour + // F20 — multipart upload identity and safe no-op abort behaviour // ------------------------------------------------------------------ @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_embedsFullUploadUuid() { + String firstUpload = "09492e3d-e231-4ed9-bf84-b6fc772cda54"; + String secondUpload = "06996d15-1c2e-4ddd-8853-43816ea84a07"; + String firstBlock = AzureObjStorage.multipartBlockId(firstUpload, 1); + String secondBlock = AzureObjStorage.multipartBlockId(secondUpload, 1); + + Assertions.assertNotEquals(firstBlock, secondBlock); + Assertions.assertEquals(firstBlock.length(), + AzureObjStorage.multipartBlockId(firstUpload, 999).length()); + byte[] decoded = Base64.getDecoder().decode(firstBlock); + Assertions.assertEquals(firstUpload.length() + Integer.BYTES, decoded.length); + Assertions.assertArrayEquals(firstUpload.getBytes(StandardCharsets.UTF_8), + Arrays.copyOf(decoded, firstUpload.length())); + Assertions.assertArrayEquals(new byte[]{1, 0, 0, 0}, + Arrays.copyOfRange(decoded, firstUpload.length(), decoded.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_returnsLocalUuidWithoutTouchingProvider() throws Exception { + BlobServiceClient serviceClient = Mockito.mock(BlobServiceClient.class); + TestableAzureObjStorage storage = new TestableAzureObjStorage(buildBasicProps(), serviceClient); + + String uploadId = storage.initiateMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob"); + Assertions.assertEquals(uploadId, UUID.fromString(uploadId).toString()); + Mockito.verifyNoInteractions(serviceClient); + } + + @Test + void uploadPart_stagesBlockWithFullUploadUuid() 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); + TestableAzureObjStorage storage = new TestableAzureObjStorage(buildBasicProps(), serviceClient); + String uploadId = "09492e3d-e231-4ed9-bf84-b6fc772cda54"; + + UploadPartResult result = storage.uploadPart( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", uploadId, 1, + RequestBody.of(new ByteArrayInputStream(new byte[]{1}), 1)); + + String expectedBlockId = AzureObjStorage.multipartBlockId(uploadId, 1); + Assertions.assertEquals(expectedBlockId, result.etag()); + Mockito.verify(blockClient).stageBlock(Mockito.eq(expectedBlockId), + Mockito.any(java.io.InputStream.class), Mockito.eq(1L)); + } + @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); - storage.abortMultipartUpload( - "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", "uploadId"); + Assertions.assertThrows(IOException.class, () -> storage.completeMultipartUpload( + "wasb://mycontainer@myaccount.blob.core.windows.net/stage/blob", + "legacy-upload-id", Collections.singletonList(new UploadPartResult(1, "")))); - // 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(); } @Test - void abortMultipartUpload_commitsEmptyAndDeletesWhenNoCommittedBlob() throws Exception { + void completeMultipartUpload_rejectsMixedExactAndMissingBlockIds() 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); - 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 +512,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); } // ------------------------------------------------------------------ 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 0d0fef3eacce79..5d98075023cdfc 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 ec615dcac47bd8..1f05f5f2312fc4 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. diff --git a/regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/Suite.groovy b/regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/Suite.groovy index cedb00cc0e3e56..a363f523310561 100644 --- a/regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/Suite.groovy +++ b/regression-test/framework/src/main/groovy/org/apache/doris/regression/suite/Suite.groovy @@ -2323,10 +2323,16 @@ class Suite implements GroovyInterceptable { } } + static String buildJobNameQuery(String dbName, String mtmvName) { + return ("select Name from jobs('type'='mv') where MvDatabaseName = '${dbName}' " + + "and MvName = '${mtmvName}'") + } + String getJobName(String dbName, String mtmvName) { - String showMTMV = "select JobName from mv_infos('database'='${dbName}') where Name = '${mtmvName}'"; - logger.info(showMTMV) - List> result = sql(showMTMV) + // Job lookup must not materialize unrelated MVs whose external metadata may be unavailable. + String showJob = buildJobNameQuery(dbName, mtmvName) + logger.info(showJob) + List> result = sql(showJob) logger.info("result: " + result.toString()) if (result.isEmpty()) { Assert.fail(); diff --git a/regression-test/framework/src/test/groovy/org/apache/doris/regression/suite/SuiteJobLookupTest.groovy b/regression-test/framework/src/test/groovy/org/apache/doris/regression/suite/SuiteJobLookupTest.groovy new file mode 100644 index 00000000000000..e66b1dad09ef86 --- /dev/null +++ b/regression-test/framework/src/test/groovy/org/apache/doris/regression/suite/SuiteJobLookupTest.groovy @@ -0,0 +1,33 @@ +// 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.regression.suite + +import org.junit.jupiter.api.Test + +import static org.junit.jupiter.api.Assertions.assertEquals + +class SuiteJobLookupTest { + @Test + void jobLookupUsesJobMetadataWithoutMaterializingMvStatus() { + String query = Suite.buildJobNameQuery("db1", "mv1") + + assertEquals( + "select Name from jobs('type'='mv') where MvDatabaseName = 'db1' and MvName = 'mv1'", + query) + } +}