Skip to content

[fix](iceberg) Harden write lifecycle and cleanup #66348

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Gabriel39 wants to merge 2 commits into
base: master
Choose a base branch
from
+1,233 −205
Open

[fix](iceberg) Harden write lifecycle and cleanup #66348

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6ebd318
[fix](iceberg) Harden write lifecycle and cleanup
Gabriel39 Aug 1, 2026
1241244
[fix](iceberg) Address write safety review findings
Gabriel39 Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 44 additions & 25 deletions be/src/exec/operator/spill_iceberg_table_sink_operator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,22 @@

#include "exec/operator/spill_iceberg_table_sink_operator.h"

#include <algorithm>

#include "common/status.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 bounded_iceberg_reserve_size(const std::vector<size_t>& per_partition_reservations) {
return per_partition_reservations.empty()
? 0
: *std::max_element(per_partition_reservations.begin(),
per_partition_reservations.end());
}

SpillIcebergTableSinkLocalState::SpillIcebergTableSinkLocalState(DataSinkOperatorXBase* parent,
RuntimeState* state)
: Base(parent, state) {}
Expand Down Expand Up @@ -55,47 +64,57 @@ size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state
if (!_writer) {
return 0;
}
auto current_writer = _writer->current_writer();
auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(current_writer.get());
if (!sort_writer) {
return 0;
std::vector<size_t> per_partition_reservations;
auto active_writers = _writer->active_writers();
per_partition_reservations.reserve(active_writers->size());
for (const auto& writer : *active_writers) {
if (auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get())) {
per_partition_reservations.push_back(sort_writer->get_reserve_mem_size(state, eos));
}
}

return sort_writer->get_reserve_mem_size(state, eos);
// One input block is partitioned among writers and consumed serially, so their full-batch estimates overlap.
return bounded_iceberg_reserve_size(per_partition_reservations);
}

size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState* state) const {
if (!_writer) {
return 0;
}
auto current_writer = _writer->current_writer();
auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(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<VIcebergSortWriter*>(writer.get())) {
revocable_size += sort_writer->data_size();
}
}

return sort_writer->data_size();
return revocable_size;
}

Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) {
RETURN_IF_CANCELLED(state);
if (!_writer) {
return Status::OK();
}
auto current_writer = _writer->current_writer();
auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(current_writer.get());
if (!sort_writer) {
return Status::OK();
std::shared_ptr<IPartitionWriterBase> 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<VIcebergSortWriter*>(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<VIcebergSortWriter*>(largest_writer.get());
RETURN_IF_CATCH_EXCEPTION({ RETURN_IF_ERROR(sort_writer->trigger_spill()); });
}
return Status::OK();
}

SpillIcebergTableSinkOperatorX::SpillIcebergTableSinkOperatorX(
Expand Down
5 changes: 4 additions & 1 deletion be/src/exec/operator/spill_iceberg_table_sink_operator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#pragma once

#include <memory>
#include <vector>

#include "exec/operator/operator.h"
#include "exec/sink/writer/iceberg/viceberg_table_writer.h"
Expand All @@ -27,6 +28,8 @@ namespace doris {
class SpillIcebergTableSinkLocalState;
class SpillIcebergTableSinkOperatorX;

size_t bounded_iceberg_reserve_size(const std::vector<size_t>& per_partition_reservations);

class SpillIcebergTableSinkLocalState final
: public AsyncWriterSink<VIcebergTableWriter, SpillIcebergTableSinkOperatorX> {
public:
Expand Down Expand Up @@ -87,4 +90,4 @@ class SpillIcebergTableSinkOperatorX final
ObjectPool* _pool = nullptr;
};

} // namespace doris
} // namespace doris
12 changes: 6 additions & 6 deletions be/src/exec/pipeline/pipeline_fragment_context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,6 +467,8 @@ Status PipelineFragmentContext::_build_pipeline_tasks_for_instance(
_params.query_options, _query_ctx->query_globals, _exec_env, _query_ctx.get());
{
// Initialize runtime state for this task
task_runtime_state->set_iceberg_commit_data_budget(
_runtime_state->iceberg_commit_data_budget());
task_runtime_state->set_query_mem_tracker(_query_ctx->query_mem_tracker());

task_runtime_state->set_task_execution_context(shared_from_this());
Expand Down Expand Up @@ -2514,16 +2516,14 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r
}
}
}
if (auto icd = req.runtime_state->iceberg_commit_datas(); !icd.empty()) {
req.runtime_state->append_iceberg_commit_datas(&params.iceberg_commit_datas);
if (!params.iceberg_commit_datas.empty()) {
params.__isset.iceberg_commit_datas = true;
params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(), icd.begin(),
icd.end());
} else if (!req.runtime_states.empty()) {
for (auto* rs : req.runtime_states) {
if (auto rs_icd = rs->iceberg_commit_datas(); !rs_icd.empty()) {
rs->append_iceberg_commit_datas(&params.iceberg_commit_datas);

@github-actions github-actions Bot Aug 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Enforce one budget for the aggregated report

The new guard is per task RuntimeState, but this loop concatenates every task's accepted vector into one TReportExecStatusParams. With a 100 MiB limit, two parallel sink tasks can each accept about 60 MiB and then send a roughly 120 MiB RPC, which FE rejects. By this point successful writers have cleared their local cleanup lists, so the failed report also strands the uncommitted objects. Please enforce a shared fragment/report budget before cleanup ownership is released and add a multi-task aggregate-limit test.

if (!params.iceberg_commit_datas.empty()) {
params.__isset.iceberg_commit_datas = true;
params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(),
rs_icd.begin(), rs_icd.end());
}
}
}
Expand Down
8 changes: 6 additions & 2 deletions be/src/exec/sink/viceberg_delete_sink.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,8 +283,12 @@ Status VIcebergDeleteSink::close(Status close_status) {
_delete_file_count);

if (_state != nullptr) {
for (const auto& commit_data : _commit_data_list) {
_state->add_iceberg_commit_datas(commit_data);
for (auto& commit_data : _commit_data_list) {
Status report_status = _state->add_iceberg_commit_datas(std::move(commit_data));
if (!report_status.ok()) {
_cleanup_created_files();
return report_status;
}
}
}

Expand Down
33 changes: 26 additions & 7 deletions be/src/exec/sink/writer/iceberg/viceberg_partition_writer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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,
Expand All @@ -103,19 +108,28 @@ Status VIcebergPartitionWriter::open(RuntimeState* state, RuntimeProfile* profil
_file_format_transformer = std::make_unique<VParquetTransformer>(
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<VOrcTransformer>(
state, _file_writer.get(), _write_output_expr_ctxs, "", _write_column_names, false,
_compress_type, &_schema, _fs);
return _file_format_transformer->open();
open_status = _file_format_transformer->open();
break;
}
default: {
return Status::InternalError("Unsupported file format type {}",
to_string(_file_format_type));
open_status = Status::InternalError("Unsupported file format type {}",
to_string(_file_format_type));
break;
}
}
if (!open_status.ok()) {
// A transformer failure happens after object creation, so abort multipart state before deleting the path.
WARN_IF_ERROR(_file_writer->abort(), "failed to abort Iceberg file after open error");
WARN_IF_ERROR(_fs->delete_file(_path), "failed to delete Iceberg file after open error");
}
return open_status;
}

Status VIcebergPartitionWriter::close(const Status& status) {
Expand Down Expand Up @@ -147,7 +161,12 @@ Status VIcebergPartitionWriter::close(const Status& status) {
}
return commit_status;
}
_state->add_iceberg_commit_datas(commit_data);
Status report_status = _state->add_iceberg_commit_datas(std::move(commit_data));
if (!report_status.ok()) {
// A closed object that cannot be reported can never be committed, so remove it immediately.
WARN_IF_ERROR(_fs->delete_file(_path), "failed to delete unreportable Iceberg file");
return report_status;
}
if (_closed_file_callback) {
_closed_file_callback(_fs, _path);
}
Expand Down
91 changes: 23 additions & 68 deletions be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,80 +102,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;
}

Expand Down
Loading
Loading