This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new e409a25a914 branch-4.1:[improvement](be) Make Iceberg data file writes
synchronous (#67087)
e409a25a914 is described below
commit e409a25a9141c939b7ffd2e024329846f14a7973
Author: daidai <[email protected]>
AuthorDate: Wed Aug 26 10:35:41 2026 +0800
branch-4.1:[improvement](be) Make Iceberg data file writes synchronous
(#67087)
### What problem does this PR solve?
Problem Summary: Iceberg data-file sinks inherited AsyncResultWriter,
retaining queued Blocks and dispatching file I/O on a separate writer
thread. Make the normal and spill table sinks own VIcebergTableWriter
directly and run its lifecycle synchronously on the blocking pipeline
scheduler. Preserve spill/revoke behavior and the nested table-writer
compatibility used by MERGE; DELETE and MERGE top-level async paths
remain unchanged.
### Release note
Iceberg data-file writes now run synchronously on the blocking pipeline
scheduler without an intermediate async Block queue.
### Check List (For Author)
- Test <!-- At least one of them must be included. -->
- [ ] Regression test
- [ ] Unit Test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test or manual test. Explain why:
- [ ] This is a refactor/code format and no logic has been changed.
- [ ] Previous test can cover this change.
- [ ] No code files have been changed.
- [ ] Other reason <!-- Add your reason? -->
- Behavior changed:
- [ ] No.
- [ ] Yes. <!-- Explain the behavior change -->
- Does this need documentation?
- [ ] No.
- [ ] Yes. <!-- Add document PR link here. eg:
https://github.com/apache/doris-website/pull/1214 -->
### Check List (For Reviewer who merge this PR)
- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label <!-- Add branch pick label that this PR
should merge into -->
---
.../exec/operator/iceberg_table_sink_operator.cpp | 59 ++-
be/src/exec/operator/iceberg_table_sink_operator.h | 26 +-
be/src/exec/operator/operator.cpp | 2 -
.../operator/spill_iceberg_table_sink_operator.cpp | 151 +++++--
.../operator/spill_iceberg_table_sink_operator.h | 24 +-
be/src/exec/sink/viceberg_merge_sink.cpp | 3 +-
.../sink/writer/iceberg/viceberg_sort_writer.cpp | 27 +-
.../sink/writer/iceberg/viceberg_sort_writer.h | 8 +-
.../sink/writer/iceberg/viceberg_table_writer.cpp | 35 +-
.../sink/writer/iceberg/viceberg_table_writer.h | 40 +-
be/src/exec/sort/sorter.cpp | 3 +-
.../operator/iceberg_table_sink_operator_test.cpp | 450 +++++++++++++++++++++
be/test/exec/sort/full_sort_test.cpp | 6 +-
13 files changed, 757 insertions(+), 77 deletions(-)
diff --git a/be/src/exec/operator/iceberg_table_sink_operator.cpp
b/be/src/exec/operator/iceberg_table_sink_operator.cpp
index 22b7ddd1531..b241ae71ac1 100644
--- a/be/src/exec/operator/iceberg_table_sink_operator.cpp
+++ b/be/src/exec/operator/iceberg_table_sink_operator.cpp
@@ -25,9 +25,64 @@ Status IcebergTableSinkLocalState::init(RuntimeState* state,
LocalSinkStateInfo&
RETURN_IF_ERROR(Base::init(state, info));
SCOPED_TIMER(exec_time_counter());
SCOPED_TIMER(_init_timer);
- auto& p = _parent->cast<Parent>();
- RETURN_IF_ERROR(_writer->init_properties(p._pool, p._row_desc));
+ _writer = std::make_unique<VIcebergTableWriter>(info.tsink,
_output_vexpr_ctxs);
+ _writer->defer_file_cleanup_until_outer_close();
+ auto& parent = _parent->cast<Parent>();
+ RETURN_IF_ERROR(_writer->init_properties(parent._pool, parent._row_desc));
return Status::OK();
}
+Status IcebergTableSinkLocalState::open(RuntimeState* state) {
+ SCOPED_TIMER(exec_time_counter());
+ SCOPED_TIMER(_open_timer);
+ RETURN_IF_ERROR(Base::open(state));
+
+ auto& parent = _parent->cast<Parent>();
+ _output_vexpr_ctxs.resize(parent._output_vexpr_ctxs.size());
+ for (size_t i = 0; i < _output_vexpr_ctxs.size(); ++i) {
+ RETURN_IF_ERROR(parent._output_vexpr_ctxs[i]->clone(state,
_output_vexpr_ctxs[i]));
+ }
+ return _writer->open(state, operator_profile());
+}
+
+Status IcebergTableSinkLocalState::sink(RuntimeState* state, Block* block,
bool /*eos*/) {
+ if (block->rows() == 0) {
+ return Status::OK();
+ }
+ DCHECK(_writer);
+ return _writer->write(state, *block);
+}
+
+Status IcebergTableSinkLocalState::close(RuntimeState* state, Status
exec_status) {
+ if (_closed) {
+ return Status::OK();
+ }
+
+ SCOPED_TIMER(exec_time_counter());
+ SCOPED_TIMER(_close_timer);
+
+ DCHECK(_writer);
+ Status final_status = exec_status;
+ // Observe cancellation before close so a known-cancelled write is not
finalized as successful.
+ if (final_status.ok() && state->is_cancelled()) {
+ final_status = state->cancel_reason();
+ }
+ Status writer_status = _writer->close(final_status);
+ if (final_status.ok() && !writer_status.ok()) {
+ final_status = writer_status;
+ }
+ // close() may block on sort/merge/file I/O, so cancellation can arrive
while it is running.
+ if (final_status.ok() && state->is_cancelled()) {
+ final_status = state->cancel_reason();
+ }
+ _writer->finish_deferred_file_cleanup(final_status);
+ _writer.reset();
+
+ Status base_status = Base::close(state, final_status);
+ if (final_status.ok() && !base_status.ok()) {
+ final_status = base_status;
+ }
+ return final_status;
+}
+
} // namespace doris
diff --git a/be/src/exec/operator/iceberg_table_sink_operator.h
b/be/src/exec/operator/iceberg_table_sink_operator.h
index 1d5cfc9c25f..3a4aa07b2d5 100644
--- a/be/src/exec/operator/iceberg_table_sink_operator.h
+++ b/be/src/exec/operator/iceberg_table_sink_operator.h
@@ -17,6 +17,8 @@
#pragma once
+#include <memory>
+
#include "exec/operator/operator.h"
#include "exec/sink/writer/iceberg/viceberg_table_writer.h"
@@ -25,21 +27,26 @@ namespace doris {
class IcebergTableSinkOperatorX;
-class IcebergTableSinkLocalState final
- : public AsyncWriterSink<VIcebergTableWriter,
IcebergTableSinkOperatorX> {
+class IcebergTableSinkLocalState final : public
PipelineXSinkLocalState<FakeSharedState> {
public:
- using Base = AsyncWriterSink<VIcebergTableWriter,
IcebergTableSinkOperatorX>;
+ using Base = PipelineXSinkLocalState<FakeSharedState>;
using Parent = IcebergTableSinkOperatorX;
ENABLE_FACTORY_CREATOR(IcebergTableSinkLocalState);
IcebergTableSinkLocalState(DataSinkOperatorXBase* parent, RuntimeState*
state)
: Base(parent, state) {};
Status init(RuntimeState* state, LocalSinkStateInfo& info) override;
- Status open(RuntimeState* state) override {
- SCOPED_TIMER(exec_time_counter());
- SCOPED_TIMER(_open_timer);
- return Base::open(state);
- }
+ Status open(RuntimeState* state) override;
+ Status sink(RuntimeState* state, Block* block, bool eos);
+ Status close(RuntimeState* state, Status exec_status) override;
+
+ [[nodiscard]] bool is_blockable() const override { return true; }
+
+private:
friend class IcebergTableSinkOperatorX;
+ friend class IcebergTableSinkOperatorTest;
+
+ VExprContextSPtrs _output_vexpr_ctxs;
+ std::unique_ptr<VIcebergTableWriter> _writer;
};
class IcebergTableSinkOperatorX final : public
DataSinkOperatorX<IcebergTableSinkLocalState> {
@@ -74,9 +81,6 @@ public:
private:
friend class IcebergTableSinkLocalState;
- template <typename Writer, typename Parent>
- requires(std::is_base_of_v<AsyncResultWriter, Writer>)
- friend class AsyncWriterSink;
const RowDescriptor& _row_desc;
VExprContextSPtrs _output_vexpr_ctxs;
const std::vector<TExpr>& _t_output_expr;
diff --git a/be/src/exec/operator/operator.cpp
b/be/src/exec/operator/operator.cpp
index e078f4417fb..36449664875 100644
--- a/be/src/exec/operator/operator.cpp
+++ b/be/src/exec/operator/operator.cpp
@@ -938,8 +938,6 @@ template class AsyncWriterSink<doris::VJdbcTableWriter,
JdbcTableSinkOperatorX>;
template class AsyncWriterSink<doris::VTabletWriter, OlapTableSinkOperatorX>;
template class AsyncWriterSink<doris::VTabletWriterV2,
OlapTableSinkV2OperatorX>;
template class AsyncWriterSink<doris::VHiveTableWriter,
HiveTableSinkOperatorX>;
-template class AsyncWriterSink<doris::VIcebergTableWriter,
IcebergTableSinkOperatorX>;
-template class AsyncWriterSink<doris::VIcebergTableWriter,
SpillIcebergTableSinkOperatorX>;
template class AsyncWriterSink<doris::VIcebergDeleteSink,
IcebergDeleteSinkOperatorX>;
template class AsyncWriterSink<doris::VIcebergMergeSink,
IcebergMergeSinkOperatorX>;
template class AsyncWriterSink<doris::VMCTableWriter, MCTableSinkOperatorX>;
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 20b4eea954a..3177933da31 100644
--- a/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp
+++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp
@@ -22,6 +22,7 @@
#include "exec/operator/spill_utils.h"
#include "exec/sink/writer/iceberg/viceberg_sort_writer.h"
#include "exec/sink/writer/iceberg/viceberg_table_writer.h"
+#include "exec/spill/spill_file.h"
namespace doris {
#include "common/compile_check_begin.h"
@@ -36,9 +37,11 @@ Status SpillIcebergTableSinkLocalState::init(RuntimeState*
state, LocalSinkState
SCOPED_TIMER(_init_timer);
_init_spill_counters();
+ _writer = std::make_unique<VIcebergTableWriter>(info.tsink,
_output_vexpr_ctxs);
+ _writer->defer_file_cleanup_until_outer_close();
- auto& p = _parent->cast<Parent>();
- RETURN_IF_ERROR(_writer->init_properties(p._pool, p._row_desc));
+ auto& parent = _parent->cast<Parent>();
+ RETURN_IF_ERROR(_writer->init_properties(parent._pool, parent._row_desc));
return Status::OK();
}
@@ -46,53 +49,149 @@ Status SpillIcebergTableSinkLocalState::open(RuntimeState*
state) {
SCOPED_TIMER(Base::exec_time_counter());
SCOPED_TIMER(Base::_open_timer);
RETURN_IF_ERROR(Base::open(state));
- return Status::OK();
+
+ auto& parent = _parent->cast<Parent>();
+ _output_vexpr_ctxs.resize(parent._output_vexpr_ctxs.size());
+ for (size_t i = 0; i < _output_vexpr_ctxs.size(); ++i) {
+ RETURN_IF_ERROR(parent._output_vexpr_ctxs[i]->clone(state,
_output_vexpr_ctxs[i]));
+ }
+ return _writer->open(state, operator_profile());
}
-bool SpillIcebergTableSinkLocalState::is_blockable() const {
- return true;
+Status SpillIcebergTableSinkLocalState::sink(RuntimeState* state, Block*
block, bool eos) {
+ if (block->rows() > 0) {
+ DCHECK(_writer);
+ RETURN_IF_ERROR(_writer->write(state, *block));
+ }
+ if (!eos) {
+ return Status::OK();
+ }
+
+ Status close_status = Status::OK();
+ if (state->is_cancelled()) {
+ close_status = state->cancel_reason();
+ }
+ return _close_writer(close_status);
}
-size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState*
state, bool eos) {
- if (!_writer) {
- return 0;
+Status SpillIcebergTableSinkLocalState::_close_writer(Status close_status) {
+ DCHECK(_writer);
+ if (!_writer_closed) {
+ _writer_close_status = _writer->close(close_status);
+ _writer_closed = true;
}
- auto current_writer = _writer->current_writer();
- auto* sort_writer =
dynamic_cast<VIcebergSortWriter*>(current_writer.get());
- if (!sort_writer) {
- return 0;
+ if (close_status.ok() && !_writer_close_status.ok()) {
+ close_status = _writer_close_status;
+ }
+ return close_status;
+}
+
+Status SpillIcebergTableSinkLocalState::close(RuntimeState* state, Status
exec_status) {
+ if (_closed) {
+ return Status::OK();
+ }
+
+ SCOPED_TIMER(exec_time_counter());
+ SCOPED_TIMER(_close_timer);
+
+ Status final_status = exec_status;
+ if (final_status.ok() && state->is_cancelled()) {
+ final_status = state->cancel_reason();
+ }
+ final_status = _close_writer(final_status);
+ DCHECK(_writer);
+ _writer->finish_deferred_file_cleanup(final_status);
+ {
+ std::lock_guard lock(_writer_mutex);
+ _writer.reset();
}
- return sort_writer->get_reserve_mem_size(state, eos);
+ Status base_status = Base::close(state, final_status);
+ if (final_status.ok() && !base_status.ok()) {
+ final_status = base_status;
+ }
+ return final_status;
+}
+
+bool SpillIcebergTableSinkLocalState::is_blockable() const {
+ return true;
+}
+
+size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState*
state, bool eos) {
+ DCHECK(_writer);
+ return _writer->get_reserve_mem_size(state, eos);
}
size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState*
state) const {
- if (!_writer) {
- return 0;
+ std::shared_ptr<const VIcebergTableWriter::PartitionWriterSnapshot>
partition_writers;
+ std::shared_ptr<IPartitionWriterBase> current_writer;
+ {
+ std::lock_guard lock(_writer_mutex);
+ if (!_writer) {
+ return 0;
+ }
+ partition_writers = _writer->partition_writers_snapshot();
+ if (!partition_writers) {
+ current_writer = _writer->current_writer();
+ }
}
- auto current_writer = _writer->current_writer();
- auto* sort_writer =
dynamic_cast<VIcebergSortWriter*>(current_writer.get());
- if (!sort_writer) {
- return 0;
+
+ if (partition_writers) {
+ size_t revocable_size = 0;
+ for (const auto& sort_writer : *partition_writers) {
+ size_t writer_size = sort_writer->data_size();
+ if (writer_size >= SpillFile::MIN_SPILL_WRITE_BATCH_MEM) {
+ revocable_size += writer_size;
+ }
+ }
+ return revocable_size;
}
+ if (!current_writer) {
+ return 0;
+ }
+ auto* sort_writer =
dynamic_cast<VIcebergSortWriter*>(current_writer.get());
+ DORIS_CHECK(sort_writer != nullptr);
return sort_writer->data_size();
}
Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) {
RETURN_IF_CANCELLED(state);
- if (!_writer) {
- return Status::OK();
+ std::shared_ptr<const VIcebergTableWriter::PartitionWriterSnapshot>
partition_writers;
+ std::shared_ptr<IPartitionWriterBase> current_writer;
+ {
+ std::lock_guard lock(_writer_mutex);
+ if (!_writer) {
+ return Status::OK();
+ }
+ partition_writers = _writer->partition_writers_snapshot();
+ if (!partition_writers) {
+ current_writer = _writer->current_writer();
+ }
}
- auto current_writer = _writer->current_writer();
- auto* sort_writer =
dynamic_cast<VIcebergSortWriter*>(current_writer.get());
- if (!sort_writer) {
+
+ std::vector<std::shared_ptr<VIcebergSortWriter>> writers_to_spill;
+ if (partition_writers) {
+ for (const auto& sort_writer : *partition_writers) {
+ if (sort_writer->data_size() >=
SpillFile::MIN_SPILL_WRITE_BATCH_MEM) {
+ writers_to_spill.emplace_back(sort_writer);
+ }
+ }
+ } else if (current_writer) {
+ auto sort_writer =
std::dynamic_pointer_cast<VIcebergSortWriter>(current_writer);
+ DORIS_CHECK(sort_writer != nullptr);
+ writers_to_spill.emplace_back(std::move(sort_writer));
+ }
+ if (writers_to_spill.empty()) {
return Status::OK();
}
- auto exception_catch_func = [current_writer, sort_writer]() {
+ auto exception_catch_func = [writers = std::move(writers_to_spill)]() {
auto status = [&]() {
- RETURN_IF_CATCH_EXCEPTION({ return sort_writer->trigger_spill();
});
+ for (const auto& sort_writer : writers) {
+ RETURN_IF_CATCH_EXCEPTION({
RETURN_IF_ERROR(sort_writer->trigger_spill()); });
+ }
+ return Status::OK();
}();
return status;
};
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 7e6a037d2f5..693a7f33005 100644
--- a/be/src/exec/operator/spill_iceberg_table_sink_operator.h
+++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.h
@@ -18,6 +18,7 @@
#pragma once
#include <memory>
+#include <mutex>
#include "exec/operator/operator.h"
#include "exec/sink/writer/iceberg/viceberg_table_writer.h"
@@ -28,10 +29,9 @@ namespace doris {
class SpillIcebergTableSinkLocalState;
class SpillIcebergTableSinkOperatorX;
-class SpillIcebergTableSinkLocalState final
- : public AsyncWriterSink<VIcebergTableWriter,
SpillIcebergTableSinkOperatorX> {
+class SpillIcebergTableSinkLocalState final : public
PipelineXSinkLocalState<FakeSharedState> {
public:
- using Base = AsyncWriterSink<VIcebergTableWriter,
SpillIcebergTableSinkOperatorX>;
+ using Base = PipelineXSinkLocalState<FakeSharedState>;
using Parent = SpillIcebergTableSinkOperatorX;
ENABLE_FACTORY_CREATOR(SpillIcebergTableSinkLocalState);
@@ -40,6 +40,8 @@ public:
Status init(RuntimeState* state, LocalSinkStateInfo& info) override;
Status open(RuntimeState* state) override;
+ Status sink(RuntimeState* state, Block* block, bool eos);
+ Status close(RuntimeState* state, Status exec_status) override;
bool is_blockable() const override;
[[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos);
@@ -47,8 +49,19 @@ public:
size_t get_revocable_mem_size(RuntimeState* state) const;
private:
+ Status _close_writer(Status close_status);
void _init_spill_counters();
friend class SpillIcebergTableSinkOperatorX;
+ friend class IcebergTableSinkOperatorTest;
+
+ VExprContextSPtrs _output_vexpr_ctxs;
+ // Protects table-writer pointer access/reset against concurrent
workload-memory callbacks.
+ // Blocking I/O stays outside this lock; a callback keeps a child-writer
shared_ptr whose
+ // sorter mutex serializes it with write/close.
+ mutable std::mutex _writer_mutex;
+ std::unique_ptr<VIcebergTableWriter> _writer;
+ bool _writer_closed = false;
+ Status _writer_close_status;
};
class SpillIcebergTableSinkOperatorX final
@@ -78,9 +91,6 @@ public:
private:
friend class SpillIcebergTableSinkLocalState;
- template <typename Writer, typename Parent>
- requires(std::is_base_of_v<AsyncResultWriter, Writer>)
- friend class AsyncWriterSink;
const RowDescriptor& _row_desc;
VExprContextSPtrs _output_vexpr_ctxs;
@@ -89,4 +99,4 @@ private:
};
#include "common/compile_check_end.h"
-} // namespace doris
\ No newline at end of file
+} // namespace doris
diff --git a/be/src/exec/sink/viceberg_merge_sink.cpp
b/be/src/exec/sink/viceberg_merge_sink.cpp
index 5008b217228..0023b37e6d4 100644
--- a/be/src/exec/sink/viceberg_merge_sink.cpp
+++ b/be/src/exec/sink/viceberg_merge_sink.cpp
@@ -56,8 +56,7 @@ Status VIcebergMergeSink::init_properties(ObjectPool* pool,
const RowDescriptor&
RETURN_IF_ERROR(_build_inner_sinks());
if (_writes_data_files) {
- _table_writer = std::make_unique<VIcebergTableWriter>(_table_sink,
_table_output_expr_ctxs,
- nullptr,
nullptr);
+ _table_writer = std::make_unique<VIcebergTableWriter>(_table_sink,
_table_output_expr_ctxs);
_table_writer->defer_file_cleanup_until_outer_close();
RETURN_IF_ERROR(_table_writer->init_properties(pool, row_desc));
}
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 db453c511fa..b296fa9c504 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp
+++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp
@@ -58,9 +58,10 @@ Status VIcebergSortWriter::open(RuntimeState* state,
RuntimeProfile* profile,
Status VIcebergSortWriter::write(Block& block) {
std::lock_guard<std::mutex> lock(_sorter_mutex);
+ // FullSorter consumes the input block, so derive the spill batch size
before append_block().
+ _update_spill_block_batch_row_count(block);
// Append incoming block data to the sorter's internal buffer
RETURN_IF_ERROR(_sorter->append_block(&block));
- _update_spill_block_batch_row_count(block);
// When accumulated data size reaches the target file size threshold,
// sort the data in memory and flush it directly to a Parquet/ORC file.
@@ -82,7 +83,21 @@ size_t VIcebergSortWriter::data_size() const {
size_t VIcebergSortWriter::get_reserve_mem_size(RuntimeState* state, bool eos)
const {
std::lock_guard<std::mutex> lock(_sorter_mutex);
- return _sorter == nullptr ? 0 : _sorter->get_reserve_mem_size(state, eos);
+ size_t reserve_size = _sorter == nullptr ? 0 :
_sorter->get_reserve_mem_size(state, eos);
+ if (eos && !_sorted_spill_files.empty()) {
+ const size_t buffer_size =
static_cast<size_t>(state->spill_buffer_size_bytes());
+ const size_t merge_limit =
static_cast<size_t>(state->spill_sort_merge_mem_limit_bytes());
+ const size_t max_fan_in = std::max<size_t>(2, merge_limit /
buffer_size);
+ // Reservation is computed before sink(), so a non-empty EOS can add
one final spill run.
+ const size_t selected_streams = std::min(_sorted_spill_files.size(),
max_fan_in - 1) + 1;
+ // Every selected cursor eagerly owns one input block. A multiway
merge also builds a
+ // separate output block while those inputs remain live.
+ const size_t merge_buffer_count = selected_streams + (selected_streams
> 1 ? 1 : 0);
+ DORIS_CHECK(buffer_size <= std::numeric_limits<size_t>::max() /
merge_buffer_count);
+ const size_t merge_reserve_size = merge_buffer_count * buffer_size;
+ reserve_size = std::max({reserve_size, merge_limit,
merge_reserve_size});
+ }
+ return reserve_size;
}
Status VIcebergSortWriter::trigger_spill() {
@@ -390,8 +405,12 @@ Status VIcebergSortWriter::_create_merger(bool
is_final_merge, size_t batch_size
}
Status VIcebergSortWriter::_create_final_merger() {
- // Final merger uses the runtime batch size and merges all remaining
streams
- return _create_merger(true, _runtime_state->batch_size(), 1);
+ // Keep the final output within both the normal runtime batch and the
spill-buffer estimate.
+ return _create_merger(true, _final_merge_batch_row_count(), 1);
+}
+
+size_t VIcebergSortWriter::_final_merge_batch_row_count() const {
+ return std::min<size_t>(_runtime_state->batch_size(),
_spill_block_batch_row_count);
}
void VIcebergSortWriter::_cleanup_spill_streams() {
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 e1e512f0a0c..9222aed9c82 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h
+++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h
@@ -101,14 +101,16 @@ public:
inline size_t written_len() const override { return
_iceberg_partition_writer->written_len(); }
- size_t data_size() const;
+ virtual size_t data_size() const;
size_t get_reserve_mem_size(RuntimeState* state, bool eos) const;
// Called by the memory management system to trigger spilling data to disk
- Status trigger_spill();
+ virtual Status trigger_spill();
private:
+ friend class IcebergTableSinkOperatorTest;
+
// 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);
@@ -152,6 +154,8 @@ private:
// Create the final merger that merges all remaining spill streams
Status _create_final_merger();
+ size_t _final_merge_batch_row_count() const;
+
// Release all spill stream resources (both pending and currently merging)
void _cleanup_spill_streams();
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 b9eeeac38a7..1c8a9523f05 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp
+++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp
@@ -40,10 +40,8 @@ namespace doris {
#include "common/compile_check_begin.h"
VIcebergTableWriter::VIcebergTableWriter(const TDataSink& t_sink,
- const VExprContextSPtrs&
output_expr_ctxs,
- std::shared_ptr<Dependency> dep,
- std::shared_ptr<Dependency> fin_dep)
- : AsyncResultWriter(output_expr_ctxs, dep, fin_dep), _t_sink(t_sink) {
+ const VExprContextSPtrs&
output_expr_ctxs)
+ : _vec_output_expr_ctxs(output_expr_ctxs), _t_sink(t_sink) {
DCHECK(_t_sink.__isset.iceberg_table_sink);
}
@@ -221,6 +219,34 @@ Status VIcebergTableWriter::write_prepared_block(Block&
block) {
return _write_prepared_block(block);
}
+size_t VIcebergTableWriter::get_reserve_mem_size(RuntimeState* state, bool
eos) const {
+ size_t reserve_size = state->minimum_operator_memory_required_bytes();
+ if (!eos) {
+ auto current_writer = _current_writer.load();
+ if (!current_writer) {
+ return reserve_size;
+ }
+ auto* sort_writer =
dynamic_cast<VIcebergSortWriter*>(current_writer.get());
+ DORIS_CHECK(sort_writer != nullptr);
+ return std::max(reserve_size, sort_writer->get_reserve_mem_size(state,
false));
+ }
+
+ auto partition_writer_snapshot =
std::make_shared<PartitionWriterSnapshot>();
+ partition_writer_snapshot->reserve(_partitions_to_writers.size());
+ // close() finalizes partition writers sequentially, so reserve the
largest close peak rather
+ // than the sum. Publish the same owning set for workload
accounting/revoke if reservation
+ // fails. The admission floor covers a first non-empty EOS before its
first writer exists.
+ for (const auto& [_, writer] : _partitions_to_writers) {
+ auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get());
+ DORIS_CHECK(sort_writer != nullptr);
+ partition_writer_snapshot->emplace_back(
+ std::static_pointer_cast<VIcebergSortWriter>(writer));
+ reserve_size = std::max(reserve_size,
sort_writer->get_reserve_mem_size(state, true));
+ }
+ _partition_writer_snapshot.store(std::move(partition_writer_snapshot));
+ return reserve_size;
+}
+
Status VIcebergTableWriter::_process_row_lineage_columns(Block& block) {
if (_write_type != TIcebergWriteType::INSERT) {
return Status::OK();
@@ -475,6 +501,7 @@ Status VIcebergTableWriter::close(Status status) {
}
}
_partitions_to_writers.clear();
+
_partition_writer_snapshot.store(std::make_shared<PartitionWriterSnapshot>());
}
if (status.ok()) {
SCOPED_TIMER(_operator_profile->total_time_counter());
diff --git a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
index 070019d85db..b618a949412 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
+++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
@@ -20,9 +20,9 @@
#include <gen_cpp/DataSinks_types.h>
#include "common/atomic_shared_ptr.h"
+#include "common/status.h"
#include "core/block/block.h"
#include "core/column/column.h"
-#include "exec/sink/writer/async_result_writer.h"
#include "exec/sink/writer/iceberg/partition_data.h"
#include "exec/sink/writer/iceberg/partition_transformers.h"
#include "exprs/vexpr_fwd.h"
@@ -44,48 +44,58 @@ class IPartitionWriterBase;
class VIcebergSortWriter;
struct ColumnWithTypeAndName;
-class VIcebergTableWriter final : public AsyncResultWriter {
+class VIcebergTableWriter {
public:
- VIcebergTableWriter(const TDataSink& t_sink, const VExprContextSPtrs&
output_exprs,
- std::shared_ptr<Dependency> dep,
std::shared_ptr<Dependency> fin_dep);
+ using PartitionWriterSnapshot =
std::vector<std::shared_ptr<VIcebergSortWriter>>;
- ~VIcebergTableWriter() = default;
+ VIcebergTableWriter(const TDataSink& t_sink, const VExprContextSPtrs&
output_exprs);
+
+ virtual ~VIcebergTableWriter() = default;
Status init_properties(ObjectPool* pool, const RowDescriptor& row_desc) {
_row_desc = &row_desc;
return Status::OK();
}
- Status open(RuntimeState* state, RuntimeProfile* profile) override;
+ virtual Status open(RuntimeState* state, RuntimeProfile* profile);
- Status write(RuntimeState* state, Block& block) override;
+ virtual Status write(RuntimeState* state, Block& block);
Status write_prepared_block(Block& block);
- Status close(Status) override;
+ virtual Status close(Status);
void defer_file_cleanup_until_outer_close() {
_defer_file_cleanup_until_outer_close = true; }
- void finish_deferred_file_cleanup(Status outer_status);
+ virtual void finish_deferred_file_cleanup(Status outer_status);
bool is_rewrite_compaction() const { return _write_type ==
TIcebergWriteType::REWRITE; }
TIcebergWriteType::type write_type() const { return _write_type; }
+ size_t get_reserve_mem_size(RuntimeState* state, bool eos) const;
+
+ std::shared_ptr<const PartitionWriterSnapshot>
partition_writers_snapshot() const {
+ return _partition_writer_snapshot.load();
+ }
+
// 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.
+ // Returns a snapshot by value. The spill/revoke path may inspect the
current writer
+ // independently of the pipeline task, so callers must hold their own copy
while operating
+ // on it instead of dereferencing the underlying member directly.
std::shared_ptr<IPartitionWriterBase> current_writer() const { return
_current_writer.load(); }
private:
+ friend class IcebergTableSinkOperatorTest;
+
// 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.
+ // Wrapped in atomic_shared_ptr because revoke_memory /
get_revocable_mem_size may run
+ // independently of the pipeline task that assigns to it.
doris::atomic_shared_ptr<IPartitionWriterBase> _current_writer;
+ mutable doris::atomic_shared_ptr<const PartitionWriterSnapshot>
_partition_writer_snapshot;
class IcebergPartitionColumn {
public:
IcebergPartitionColumn(const iceberg::PartitionField& field,
@@ -148,8 +158,10 @@ private:
void _cleanup_closed_files();
// Currently it is a copy, maybe it is better to use move semantics to
eliminate it.
+ const VExprContextSPtrs& _vec_output_expr_ctxs;
TDataSink _t_sink;
RuntimeState* _state = nullptr;
+ RuntimeProfile* _operator_profile = nullptr;
// Target file size in bytes for controlling when to split files
int64_t _target_file_size_bytes = 0;
diff --git a/be/src/exec/sort/sorter.cpp b/be/src/exec/sort/sorter.cpp
index 26deeef3a85..b06c7426cb9 100644
--- a/be/src/exec/sort/sorter.cpp
+++ b/be/src/exec/sort/sorter.cpp
@@ -197,7 +197,8 @@ size_t FullSorter::get_reserve_mem_size(RuntimeState*
state, bool eos) const {
if ((new_block_bytes * 100 / allocated_bytes) >= 85) {
size_to_reserve += (size_t)(allocated_bytes * 1.15);
}
- auto sort = new_rows > _buffered_block_size || new_block_bytes >
_buffered_block_bytes;
+ auto sort =
+ eos || 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
diff --git a/be/test/exec/operator/iceberg_table_sink_operator_test.cpp
b/be/test/exec/operator/iceberg_table_sink_operator_test.cpp
new file mode 100644
index 00000000000..6c3e48d4576
--- /dev/null
+++ b/be/test/exec/operator/iceberg_table_sink_operator_test.cpp
@@ -0,0 +1,450 @@
+// 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/operator/iceberg_table_sink_operator.h"
+
+#include <gtest/gtest.h>
+
+#include <functional>
+#include <limits>
+#include <memory>
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "common/object_pool.h"
+#include "core/data_type/data_type_number.h"
+#include "core/data_type/data_type_string.h"
+#include "exec/operator/spill_iceberg_table_sink_operator.h"
+#include "exec/sink/writer/async_result_writer.h"
+#include "exec/sink/writer/iceberg/viceberg_sort_writer.h"
+#include "runtime/runtime_profile.h"
+#include "testutil/column_helper.h"
+#include "testutil/mock/mock_descriptors.h"
+#include "testutil/mock/mock_runtime_state.h"
+#include "testutil/mock/mock_slot_ref.h"
+
+namespace doris {
+
+static_assert(!std::is_base_of_v<AsyncResultWriter, VIcebergTableWriter>);
+
+namespace {
+
+TDataSink make_iceberg_table_sink() {
+ TIcebergTableSink iceberg_sink;
+ TDataSink sink;
+ sink.__set_type(TDataSinkType::ICEBERG_TABLE_SINK);
+ sink.__set_iceberg_table_sink(iceberg_sink);
+ return sink;
+}
+
+struct FakeWriterState {
+ Status open_result;
+ Status write_result;
+ Status close_result;
+ int open_count = 0;
+ int write_count = 0;
+ int close_count = 0;
+ int cleanup_count = 0;
+ size_t written_rows = 0;
+ std::vector<Status> close_inputs;
+ Status cleanup_status;
+ std::function<void()> close_hook;
+};
+
+class FakeIcebergTableWriter final : public VIcebergTableWriter {
+public:
+ FakeIcebergTableWriter(const VExprContextSPtrs& output_exprs,
+ std::shared_ptr<FakeWriterState> state)
+ : VIcebergTableWriter(make_iceberg_table_sink(), output_exprs),
+ _fake_state(std::move(state)) {}
+
+ Status open(RuntimeState*, RuntimeProfile*) override {
+ ++_fake_state->open_count;
+ return _fake_state->open_result;
+ }
+
+ Status write(RuntimeState*, Block& block) override {
+ ++_fake_state->write_count;
+ _fake_state->written_rows += block.rows();
+ return _fake_state->write_result;
+ }
+
+ Status close(Status status) override {
+ ++_fake_state->close_count;
+ _fake_state->close_inputs.emplace_back(status);
+ if (_fake_state->close_hook) {
+ _fake_state->close_hook();
+ }
+ return _fake_state->close_result;
+ }
+
+ void finish_deferred_file_cleanup(Status status) override {
+ ++_fake_state->cleanup_count;
+ _fake_state->cleanup_status = status;
+ }
+
+private:
+ std::shared_ptr<FakeWriterState> _fake_state;
+};
+
+struct FakeSortWriterState {
+ size_t data_size = 0;
+ int spill_count = 0;
+};
+
+class FakeRevocableIcebergSortWriter final : public VIcebergSortWriter {
+public:
+ explicit
FakeRevocableIcebergSortWriter(std::shared_ptr<FakeSortWriterState> state)
+ : VIcebergSortWriter(nullptr, TSortInfo {}, 0),
_fake_state(std::move(state)) {}
+
+ size_t data_size() const override { return _fake_state->data_size; }
+
+ Status trigger_spill() override {
+ ++_fake_state->spill_count;
+ _fake_state->data_size = 0;
+ return Status::OK();
+ }
+
+private:
+ std::shared_ptr<FakeSortWriterState> _fake_state;
+};
+
+} // namespace
+
+class IcebergTableSinkOperatorTest : public testing::Test {
+protected:
+ template <typename Parent, typename LocalState>
+ void initialize_local_state(Parent* parent, LocalState* local_state,
MockRuntimeState* state,
+ std::shared_ptr<FakeWriterState>* fake_state) {
+ TDataSink sink = make_iceberg_table_sink();
+ ASSERT_TRUE(parent->init(sink).ok());
+
+ auto shared_state = parent->create_shared_state();
+ LocalSinkStateInfo info {.task_idx = 0,
+ .parent_profile = &_parent_profile,
+ .sender_id = 0,
+ .shared_state = shared_state.get(),
+ .shared_state_map = {},
+ .tsink = sink};
+ ASSERT_TRUE(local_state->init(state, info).ok());
+
+ *fake_state = std::make_shared<FakeWriterState>();
+ auto writer = std::make_unique<FakeIcebergTableWriter>(_output_exprs,
*fake_state);
+ local_state->_writer = std::move(writer);
+ ASSERT_TRUE(local_state->open(state).ok());
+ }
+
+ void set_partition_writers(SpillIcebergTableSinkLocalState* local_state,
+ std::shared_ptr<VIcebergSortWriter> first,
+ std::shared_ptr<VIcebergSortWriter> second) {
+ auto current_writer = second;
+ local_state->_writer->_partitions_to_writers.emplace("first",
std::move(first));
+ local_state->_writer->_partitions_to_writers.emplace("second",
std::move(second));
+ local_state->_writer->_current_writer.store(std::move(current_writer));
+ }
+
+ void mark_spilled(VIcebergSortWriter* writer, size_t count = 1) {
+ for (size_t i = 0; i < count; ++i) {
+ writer->_sorted_spill_files.emplace_back(nullptr);
+ }
+ }
+
+ void initialize_sort_writer_for_write(VIcebergSortWriter* writer,
RuntimeState* state,
+ const DataTypePtr& data_type) {
+ _sort_row_desc =
+ std::make_unique<MockRowDescriptor>(std::vector<DataTypePtr>
{data_type}, &_pool);
+ writer->_runtime_state = state;
+ writer->_ordering_expr_ctxs = MockSlotRef::create_mock_contexts(0,
data_type);
+ writer->_sort_info.is_asc_order = {true};
+ writer->_sort_info.nulls_first = {false};
+ writer->_sorter = FullSorter::create_unique(
+ writer->_ordering_expr_ctxs, -1, 0, &writer->_pool,
writer->_sort_info.is_asc_order,
+ writer->_sort_info.nulls_first, *_sort_row_desc, state,
nullptr);
+ writer->_sorter->set_enable_spill();
+ writer->_target_file_size_bytes = std::numeric_limits<int64_t>::max();
+ }
+
+ std::pair<size_t, size_t> spill_batch_state(const VIcebergSortWriter&
writer) const {
+ return {writer._avg_row_bytes, writer._spill_block_batch_row_count};
+ }
+
+ size_t final_merge_batch_row_count(const VIcebergSortWriter& writer) const
{
+ return writer._final_merge_batch_row_count();
+ }
+
+ ObjectPool _pool;
+ RowDescriptor _row_desc;
+ std::unique_ptr<MockRowDescriptor> _sort_row_desc;
+ VExprContextSPtrs _output_exprs;
+ RuntimeProfile _parent_profile {"IcebergTableSinkOperatorTest"};
+};
+
+TEST_F(IcebergTableSinkOperatorTest,
SyncWritersUseBlockingSchedulerWithoutDependencies) {
+ IcebergTableSinkLocalState table_sink(nullptr, nullptr);
+ SpillIcebergTableSinkLocalState spill_sink(nullptr, nullptr);
+
+ EXPECT_TRUE(table_sink.is_blockable());
+ EXPECT_TRUE(spill_sink.is_blockable());
+ EXPECT_TRUE(table_sink.dependencies().empty());
+ EXPECT_TRUE(spill_sink.dependencies().empty());
+ EXPECT_EQ(table_sink.finishdependency(), nullptr);
+ EXPECT_EQ(spill_sink.finishdependency(), nullptr);
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
NormalSinkWritesDataAndClosesAfterEmptyEos) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ IcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ IcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ Block data = ColumnHelper::create_block<DataTypeInt32>({1, 2, 3});
+ Block empty = ColumnHelper::create_block<DataTypeInt32>({});
+ ASSERT_TRUE(local_state.sink(&state, &data, false).ok());
+ ASSERT_TRUE(local_state.sink(&state, &empty, true).ok());
+ EXPECT_EQ(fake_state->write_count, 1);
+ EXPECT_EQ(fake_state->written_rows, 3);
+ EXPECT_EQ(fake_state->close_count, 0);
+
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+ EXPECT_EQ(fake_state->close_count, 1);
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+ EXPECT_TRUE(fake_state->close_inputs.front().ok());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, NormalSinkConvertsOkCloseToCancellation) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ IcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ IcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ state.cancel(Status::Cancelled("cancel normal Iceberg sink"));
+ Status close_status = local_state.close(&state, Status::OK());
+ EXPECT_TRUE(close_status.is<ErrorCode::CANCELLED>()) <<
close_status.to_string();
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+ EXPECT_TRUE(fake_state->close_inputs.front().is<ErrorCode::CANCELLED>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
NormalSinkObservesCancellationDuringWriterClose) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ IcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ IcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+ fake_state->close_hook = [&state]() {
+ state.cancel(Status::Cancelled("cancel during normal writer close"));
+ };
+
+ Status close_status = local_state.close(&state, Status::OK());
+ EXPECT_TRUE(close_status.is<ErrorCode::CANCELLED>()) <<
close_status.to_string();
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+ EXPECT_TRUE(fake_state->close_inputs.front().ok());
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+ EXPECT_TRUE(fake_state->cleanup_status.is<ErrorCode::CANCELLED>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, SpillSinkFinalizesNonEmptyEos) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ Block data = ColumnHelper::create_block<DataTypeInt32>({1, 2, 3});
+ ASSERT_TRUE(local_state.sink(&state, &data, true).ok());
+ EXPECT_EQ(fake_state->write_count, 1);
+ EXPECT_EQ(fake_state->close_count, 1);
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+ EXPECT_EQ(fake_state->close_count, 1);
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+ EXPECT_TRUE(fake_state->cleanup_status.ok());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, SpillSinkFinalizesEmptyEosAfterData) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ Block data = ColumnHelper::create_block<DataTypeInt32>({1, 2, 3});
+ Block empty = ColumnHelper::create_block<DataTypeInt32>({});
+ ASSERT_TRUE(local_state.sink(&state, &data, false).ok());
+ EXPECT_EQ(fake_state->close_count, 0);
+ ASSERT_TRUE(local_state.sink(&state, &empty, true).ok());
+ EXPECT_EQ(fake_state->close_count, 1);
+
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+ EXPECT_EQ(fake_state->close_count, 1);
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SpillSinkReservesEosAdmissionAndLargestPartitionMerge) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true),
+ state.minimum_operator_memory_required_bytes());
+
+ auto empty_writer = std::make_shared<VIcebergSortWriter>(nullptr,
TSortInfo {}, 0);
+ auto spilled_writer = std::make_shared<VIcebergSortWriter>(nullptr,
TSortInfo {}, 0);
+ auto* spilled_writer_ptr = spilled_writer.get();
+ mark_spilled(spilled_writer.get());
+ EXPECT_EQ(spilled_writer->get_reserve_mem_size(&state, true),
+ static_cast<size_t>(state.spill_sort_merge_mem_limit_bytes()));
+ set_partition_writers(&local_state, std::move(empty_writer),
std::move(spilled_writer));
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true),
+ static_cast<size_t>(state.spill_sort_merge_mem_limit_bytes()));
+
+ mark_spilled(spilled_writer_ptr, 7);
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true), 72 * 1024 *
1024);
+
+ TQueryOptions query_options = state.query_options();
+ query_options.__set_spill_buffer_size_bytes(256 * 1024 * 1024);
+ query_options.__set_spill_sort_merge_mem_limit_bytes(1024 * 1024);
+ state.set_query_options(query_options);
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true), 768 * 1024 *
1024);
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SortWriterSamplesSpillBatchBeforeConsumingInput) {
+ MockRuntimeState state;
+ TQueryOptions query_options = state.query_options();
+ query_options.__set_spill_buffer_size_bytes(1024 * 1024);
+ state.set_query_options(query_options);
+
+ auto data_type = std::make_shared<DataTypeString>();
+ VIcebergSortWriter writer(nullptr, TSortInfo {},
std::numeric_limits<int64_t>::max());
+ initialize_sort_writer_for_write(&writer, &state, data_type);
+
+ std::string wide_value(64 * 1024, 'x');
+ Block block = ColumnHelper::create_block<DataTypeString>(
+ {wide_value, wide_value, wide_value, wide_value});
+ const size_t expected_avg_row_bytes = block.bytes() / block.rows();
+ const size_t expected_batch_rows =
+ (state.spill_buffer_size_bytes() + expected_avg_row_bytes - 1) /
expected_avg_row_bytes;
+
+ ASSERT_TRUE(writer.write(block).ok());
+ EXPECT_EQ(block.rows(), 0);
+ EXPECT_EQ(spill_batch_state(writer),
+ std::make_pair(expected_avg_row_bytes, expected_batch_rows));
+ EXPECT_EQ(final_merge_batch_row_count(writer),
+ std::min<size_t>(state.batch_size(), expected_batch_rows));
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SpillSinkAccountsAndRevokesAllEosPartitions) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ auto large_state = std::make_shared<FakeSortWriterState>();
+ large_state->data_size = 1024 * 1024;
+ auto small_state = std::make_shared<FakeSortWriterState>();
+ small_state->data_size = 128 * 1024;
+ set_partition_writers(&local_state,
+
std::make_shared<FakeRevocableIcebergSortWriter>(large_state),
+
std::make_shared<FakeRevocableIcebergSortWriter>(small_state));
+ static_cast<void>(local_state.get_reserve_mem_size(&state, true));
+
+ EXPECT_EQ(local_state.get_revocable_mem_size(&state),
large_state->data_size);
+ ASSERT_TRUE(local_state.revoke_memory(&state).ok());
+ EXPECT_EQ(large_state->spill_count, 1);
+ EXPECT_EQ(small_state->spill_count, 0);
+ EXPECT_EQ(local_state.get_revocable_mem_size(&state), 0);
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, SpillSinkWriteErrorClosesWithError) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+ fake_state->write_result = Status::InternalError("injected write failure");
+
+ Block data = ColumnHelper::create_block<DataTypeInt32>({1});
+ Status sink_status = local_state.sink(&state, &data, true);
+ EXPECT_TRUE(sink_status.is<ErrorCode::INTERNAL_ERROR>()) <<
sink_status.to_string();
+ EXPECT_EQ(fake_state->close_count, 0);
+
+ Status close_status = local_state.close(&state, sink_status);
+ EXPECT_TRUE(close_status.is<ErrorCode::INTERNAL_ERROR>()) <<
close_status.to_string();
+ EXPECT_EQ(fake_state->close_count, 1);
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+
EXPECT_TRUE(fake_state->close_inputs.front().is<ErrorCode::INTERNAL_ERROR>());
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+ EXPECT_TRUE(fake_state->cleanup_status.is<ErrorCode::INTERNAL_ERROR>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, SpillSinkCloseErrorPropagatesAndCleansUp)
{
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+ fake_state->close_result = Status::InternalError("injected close failure");
+
+ Block empty = ColumnHelper::create_block<DataTypeInt32>({});
+ Status sink_status = local_state.sink(&state, &empty, true);
+ EXPECT_TRUE(sink_status.is<ErrorCode::INTERNAL_ERROR>()) <<
sink_status.to_string();
+ EXPECT_EQ(fake_state->close_count, 1);
+
+ Status close_status = local_state.close(&state, sink_status);
+ EXPECT_TRUE(close_status.is<ErrorCode::INTERNAL_ERROR>()) <<
close_status.to_string();
+ EXPECT_EQ(fake_state->close_count, 1);
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+ EXPECT_TRUE(fake_state->cleanup_status.is<ErrorCode::INTERNAL_ERROR>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SpillSinkCancellationAfterEosDeletesDeferredFiles) {
+ MockRuntimeState state;
+ std::vector<TExpr> thrift_exprs;
+ SpillIcebergTableSinkOperatorX parent(&_pool, 1, _row_desc, thrift_exprs);
+ SpillIcebergTableSinkLocalState local_state(&parent, &state);
+ std::shared_ptr<FakeWriterState> fake_state;
+ initialize_local_state(&parent, &local_state, &state, &fake_state);
+
+ Block empty = ColumnHelper::create_block<DataTypeInt32>({});
+ ASSERT_TRUE(local_state.sink(&state, &empty, true).ok());
+ EXPECT_EQ(fake_state->close_count, 1);
+
+ state.cancel(Status::Cancelled("cancel spill Iceberg sink after EOS"));
+ Status close_status = local_state.close(&state, Status::OK());
+ EXPECT_TRUE(close_status.is<ErrorCode::CANCELLED>()) <<
close_status.to_string();
+ EXPECT_EQ(fake_state->close_count, 1);
+ EXPECT_EQ(fake_state->cleanup_count, 1);
+ EXPECT_TRUE(fake_state->cleanup_status.is<ErrorCode::CANCELLED>());
+}
+
+} // namespace doris
diff --git a/be/test/exec/sort/full_sort_test.cpp
b/be/test/exec/sort/full_sort_test.cpp
index e182048c807..b7b695fa278 100644
--- a/be/test/exec/sort/full_sort_test.cpp
+++ b/be/test/exec/sort/full_sort_test.cpp
@@ -91,7 +91,9 @@ TEST_F(FullSorterTest, test_full_sorter2) {
EXPECT_TRUE(sorter->append_block(&block).ok());
}
- std::cout << sorter->get_reserve_mem_size(&_state, false) << std::endl;
+ size_t non_eos_reserve_size = sorter->get_reserve_mem_size(&_state, false);
+ size_t eos_reserve_size = sorter->get_reserve_mem_size(&_state, true);
+ EXPECT_GT(eos_reserve_size, non_eos_reserve_size);
}
TEST_F(FullSorterTest, test_full_sorter3) {
@@ -113,4 +115,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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
