This is an automated email from the ASF dual-hosted git repository.
morningman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new be7243a4894 [improvement](be) Make Iceberg data file writes
synchronous (#69355)
be7243a4894 is described below
commit be7243a4894c0641216fb7a1555631f91511e5fb
Author: Mingyu Chen (Rayner) <[email protected]>
AuthorDate: Sun Oct 25 08:06:45 2026 +0000
[improvement](be) Make Iceberg data file writes synchronous (#69355)
### What problem does this PR solve?
Issue Number: None
Related PR: #67087
Problem Summary:
**In short.** This ports branch-4.1 #67087 to master. The normal and the
spill Iceberg table sinks stop being async writer sinks: they own the
Iceberg table writer and open, write and close it on the pipeline task,
scheduled on the blocking scheduler, the same way the Paimon sink
already works on master. Master's per-partition memory reservation and
spill revocation stay as they are.
**Background.** On master both Iceberg data file sinks
(`IcebergTableSinkLocalState`, and `SpillIcebergTableSinkLocalState` for
tables with a sort order) are `AsyncWriterSink`s. Every block is moved
into a queue, and a thread taken from the fragment manager's thread pool
opens the files, writes every block and closes the writer for the whole
statement, while the pipeline task only enqueues. To keep memory
admission correct with that hand-off, master moves each pipeline
reservation into the queue together with its block, holds the EOS
reservation in the writer thread until `close()` finishes, and makes the
spill sink admit a block only after the previous one has been fully
processed, because the reservation it computes samples sorter state that
the writer thread updates.
**The problem, and what it cost.**
| | Master today | After this PR |
|---|---|---|
| Who writes the files | A fragment-pool thread per sink instance, for
the whole statement | The pipeline task, on the blocking scheduler |
| Blocks in flight per instance | Up to 3 blocks parked in the async
queue | None, each block is written in `sink()` |
| Spill sink pipelining | None: each block waits until the writer thread
has processed the previous one | Same work, without the thread hop |
| Memory reservation | Transferred through the queue; the EOS
reservation is held by the writer thread through `finish()`/`close()` |
The EOS `sink()` call itself closes the writer under its reservation |
| Final spill merge output block | `batch_size()` rows, which for wide
rows is far larger than the spill buffer the merge reservation charges
for it | At most one spill batch of rows (≈ spill buffer size) |
**How this PR fixes it.**
- `VIcebergTableWriter` no longer derives from `AsyncResultWriter`; it
keeps the output expressions and operator profile itself. `open`,
`write`, `close` and `finish_deferred_file_cleanup` are virtual so unit
tests can replace the writer.
- `IcebergTableSinkLocalState` and `SpillIcebergTableSinkLocalState`
derive from `PipelineXSinkLocalState<FakeSharedState>`, are blockable,
create the writer in `init()`, clone the output expressions and open it
in `open()`, and write each non-empty block in `sink()`. Their `close()`
folds a cancellation seen before or during the writer close into the
final status, and calls `finish_deferred_file_cleanup()` with that
status, so a statement cancelled after the files were closed deletes
them instead of handing them to the final report.
- Both sinks close the writer in the EOS `sink()` call (idempotently
with `close()`, which closes it only when the task ends before EOS),
because the pipeline task holds the sink reservation only while `sink()`
runs. For the normal sink, closing flushes what every partition file
still buffers, such as the last half-filled Parquet row group; its EOS
call reserves the sink default
(`minimum_operator_memory_required_bytes()`), the amount master's async
path held until the writer closed. For the spill sink, the final sort
and spill merge run there; its EOS reservation already includes the
merge workspace of the largest partition
(`iceberg_spill_merge_workspace`) and now also covers the work that runs
inside `sink()`. Master's per-partition reservation, cold-writer
reservation and largest-partition-first revocation are kept. A mutex
guards the spill sink's writer pointer so a memory callback that arrives
after `close()` sees the released writer.
- The final spill merge of `VIcebergSortWriter` emits `min(batch_size,
spill batch rows)` rows per block, matching the per-block size the merge
reservation assumes.
- `AsyncResultWriter::wait_for_processing_before_next_sink()` and the
waiting mode of the async queue admission only existed for the Iceberg
spill sink and are removed. The other async writers (Hive, MaxCompute,
JDBC, TVF, OLAP, result file) always used the plain bounded queue, whose
behavior does not change.
- DELETE and MERGE sinks stay asynchronous. The MERGE sink keeps using
`VIcebergTableWriter` as its nested data file writer.
**Results.**
- Unit tests: all 329 tests of the Iceberg suites (`*Iceberg*`,
including the 18 cases of the new `IcebergTableSinkOperatorTest`),
`FullSorterTest`, `SorterReserveMemoryTest` and `AsyncResultWriterTest`
pass. Four mutations (the final merge batch back at `batch_size()` rows,
no writer close in the spill sink's EOS `sink()`, and two of the
cancellation checks around the writer close removed) make 7 of the
original 13 cases fail. Three mutations of the normal sink (no writer
close in its EOS `sink()`, a second writer close in `close()`, no
cancellation check at EOS) make 5, 5 and 1 of its cases fail.
- Regression test: the new `test_iceberg_write_sorted_spill` writes
60000 rows into a REST catalog table sorted by `id desc nulls last` and
partitioned by `region`, with forced spilling and
`spill_min_revocable_mem = 1`. On a local cluster the BE paused the
INSERT three times to revoke `SPILL_ICEBERG_TABLE_SINK_OPERATOR`; each
partition was still committed as one data file of 20000 rows, and Spark
reads back 60000 rows, the expected sum, and no row out of order inside
any file (`lag(id) over (partition by _file order by _pos)`). All 43
suites of `external_table_p0/iceberg/write` pass on the same cluster.
- Manual test: on a local `file://` Iceberg catalog the same
forced-spill INSERT wrote 12 spill files (4 per partition) before the
final merge, every data file read back in descending `id` order, and
UPDATE / DELETE / MERGE on the table returned the expected rows.
<details>
<summary>Classes and call graph</summary>
```
PipelineTask::execute (blocking scheduler, sink is_blockable() == true)
├─ IcebergTableSinkOperatorX::sink_impl
│ └─ IcebergTableSinkLocalState::sink
│ ├─ VIcebergTableWriter::write
│ └─ eos: _close_writer ──> VIcebergTableWriter::close
│ (flushes the files every partition
writer still buffers)
├─ SpillIcebergTableSinkOperatorX::get_reserve_mem_size (master model,
unchanged)
├─ SpillIcebergTableSinkOperatorX::sink_impl
│ └─ SpillIcebergTableSinkLocalState::sink
│ ├─ VIcebergTableWriter::write ──> VIcebergSortWriter::write (per
partition)
│ └─ eos: _close_writer ──> VIcebergTableWriter::close ──>
VIcebergSortWriter::close
│ (final sort / spill merge, capped
output batch)
└─ PipelineTask::close ──> *LocalState::close
└─ cancellation check ─ writer close (if not yet) ─ cancellation
check
─ VIcebergTableWriter::finish_deferred_file_cleanup(final status)
RevokableTask (while the task is paused)
└─ SpillIcebergTableSinkLocalState::revoke_memory
└─ under _writer_mutex: active_writers() snapshot ──> largest
VIcebergSortWriter::trigger_spill
```
</details>
Co-authored-by: daidai <[email protected]>
---
.../exec/operator/iceberg_table_sink_operator.cpp | 82 ++-
be/src/exec/operator/iceberg_table_sink_operator.h | 33 +-
be/src/exec/operator/operator.cpp | 2 -
.../operator/spill_iceberg_table_sink_operator.cpp | 107 +++-
.../operator/spill_iceberg_table_sink_operator.h | 25 +-
be/src/exec/sink/viceberg_merge_sink.cpp | 3 +-
be/src/exec/sink/writer/async_result_writer.cpp | 16 -
be/src/exec/sink/writer/async_result_writer.h | 8 +-
.../sink/writer/async_writer_queue_admission.h | 18 -
.../sink/writer/iceberg/viceberg_sort_writer.cpp | 8 +-
.../sink/writer/iceberg/viceberg_sort_writer.h | 14 +-
.../sink/writer/iceberg/viceberg_table_writer.cpp | 6 +-
.../sink/writer/iceberg/viceberg_table_writer.h | 24 +-
.../operator/iceberg_table_sink_operator_test.cpp | 591 +++++++++++++++++++++
.../spill_iceberg_table_sink_operator_test.cpp | 19 -
.../external_table_writer_partition_limit_test.cpp | 2 +-
.../writer/iceberg/iceberg_table_writer_test.cpp | 10 +-
.../writer/iceberg/viceberg_table_writer_test.cpp | 18 +-
be/test/exec/sort/full_sort_test.cpp | 4 +-
.../write/test_iceberg_write_sorted_spill.out | 16 +
.../write/test_iceberg_write_sorted_spill.groovy | 114 ++++
21 files changed, 977 insertions(+), 143 deletions(-)
diff --git a/be/src/exec/operator/iceberg_table_sink_operator.cpp
b/be/src/exec/operator/iceberg_table_sink_operator.cpp
index 8d2074e3a49..777aa1569b0 100644
--- a/be/src/exec/operator/iceberg_table_sink_operator.cpp
+++ b/be/src/exec/operator/iceberg_table_sink_operator.cpp
@@ -24,9 +24,87 @@ 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) {
+ DCHECK(_writer);
+ RETURN_IF_ERROR(_writer->write(state, *block));
+ }
+ if (!eos) {
+ return Status::OK();
+ }
+
+ // The pipeline task holds the sink reservation only while sink() runs.
Closing the writer
+ // flushes what each partition file still buffers, e.g. the last
half-filled Parquet row group,
+ // whose final pages grow the page and compression buffers. Done in
close(), that flush would
+ // run after the EOS reservation was released, so close the writer here.
+ Status close_status = Status::OK();
+ if (state->is_cancelled()) {
+ close_status = state->cancel_reason();
+ }
+ return _close_writer(close_status);
+}
+
+Status IcebergTableSinkLocalState::_close_writer(Status close_status) {
+ DCHECK(_writer);
+ if (!_writer_closed) {
+ _writer_close_status = _writer->close(close_status);
+ _writer_closed = true;
+ }
+ if (close_status.ok() && !_writer_close_status.ok()) {
+ close_status = _writer_close_status;
+ }
+ return close_status;
+}
+
+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();
+ }
+ // A task that reached EOS already closed the writer in sink(); an error
or cancellation before
+ // EOS closes it here with that status.
+ final_status = _close_writer(final_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 a91950af7a7..ab6cb1dc508 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"
@@ -24,21 +26,30 @@ 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;
+
+ // Writes run synchronously, so open, write and close must use the
blocking scheduler.
+ [[nodiscard]] bool is_blockable() const override { return true; }
+
+private:
+ Status _close_writer(Status close_status);
friend class IcebergTableSinkOperatorX;
+ friend class IcebergTableSinkOperatorTest;
+
+ VExprContextSPtrs _output_vexpr_ctxs;
+ std::unique_ptr<VIcebergTableWriter> _writer;
+ bool _writer_closed = false;
+ Status _writer_close_status;
};
class IcebergTableSinkOperatorX final : public
DataSinkOperatorX<IcebergTableSinkLocalState> {
@@ -73,16 +84,10 @@ 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;
ObjectPool* _pool = nullptr;
};
-/// Instantiated once in operator.cpp; suppresses per-TU implicit
instantiation.
-extern template class AsyncWriterSink<VIcebergTableWriter,
IcebergTableSinkOperatorX>;
-
} // namespace doris
diff --git a/be/src/exec/operator/operator.cpp
b/be/src/exec/operator/operator.cpp
index 1213acce8e2..a3f97fb873b 100644
--- a/be/src/exec/operator/operator.cpp
+++ b/be/src/exec/operator/operator.cpp
@@ -971,8 +971,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 18a5b95f643..87fae242e70 100644
--- a/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp
+++ b/be/src/exec/operator/spill_iceberg_table_sink_operator.cpp
@@ -46,15 +46,15 @@
SpillIcebergTableSinkLocalState::SpillIcebergTableSinkLocalState(DataSinkOperato
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);
_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();
}
@@ -62,7 +62,74 @@ 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());
+}
+
+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();
+ }
+
+ // Finish sorting and spill merging here so they run under the EOS
reservation.
+ Status close_status = Status::OK();
+ if (state->is_cancelled()) {
+ close_status = state->cancel_reason();
+ }
+ return _close_writer(close_status);
+}
+
+Status SpillIcebergTableSinkLocalState::_close_writer(Status close_status) {
+ DCHECK(_writer);
+ if (!_writer_closed) {
+ _writer_close_status = _writer->close(close_status);
+ _writer_closed = true;
+ }
+ 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;
+ // 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();
+ }
+ final_status = _close_writer(final_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();
+ }
+ DCHECK(_writer);
+ _writer->finish_deferred_file_cleanup(final_status);
+ {
+ std::lock_guard lock(_writer_mutex);
+ _writer.reset();
+ }
+
+ 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 {
@@ -71,9 +138,7 @@ bool SpillIcebergTableSinkLocalState::is_blockable() const {
size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState*
state, bool eos,
const Block*
block) {
- if (!_writer) {
- return 0;
- }
+ DCHECK(_writer);
std::vector<IcebergSorterReserveMemory> per_partition_reservations;
const size_t incoming_rows = block == nullptr ? 0 : block->rows();
const size_t incoming_bytes = block == nullptr ? 0 :
block->allocated_bytes();
@@ -90,7 +155,7 @@ size_t
SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state
}
}
// 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().
+ // The EOS block may contain rows, and its reservation also covers the
writer close in sink().
const size_t incoming_reserve =
block == nullptr ? state->minimum_operator_memory_required_bytes()
: iceberg_cold_writer_reserve_size(
@@ -100,12 +165,16 @@ size_t
SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state
}
size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState*
state) const {
- if (!_writer) {
- return 0;
+ // Retain the published container while the pipeline task may replace the
current snapshot.
+ std::shared_ptr<VIcebergTableWriter::ActiveWriterSnapshot> active_writers;
+ {
+ std::lock_guard lock(_writer_mutex);
+ if (!_writer) {
+ return 0;
+ }
+ active_writers = _writer->active_writers();
}
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();
@@ -116,13 +185,17 @@ size_t
SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState* sta
Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) {
RETURN_IF_CANCELLED(state);
- if (!_writer) {
- return Status::OK();
+ // Retain the published container while the pipeline task may replace the
current snapshot.
+ std::shared_ptr<VIcebergTableWriter::ActiveWriterSnapshot> active_writers;
+ {
+ std::lock_guard lock(_writer_mutex);
+ if (!_writer) {
+ return Status::OK();
+ }
+ active_writers = _writer->active_writers();
}
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();
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 e05d0c48164..f7a925af35e 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 <vector>
#include "exec/operator/iceberg_sorter_reserve_memory.h"
@@ -29,10 +30,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);
@@ -41,6 +41,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,
const Block* block);
@@ -48,8 +50,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 the published
writer snapshot whose
+ // sorter mutexes serialize 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
@@ -79,9 +92,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,7 +99,4 @@ private:
ObjectPool* _pool = nullptr;
};
-/// Instantiated once in operator.cpp; suppresses per-TU implicit
instantiation.
-extern template class AsyncWriterSink<VIcebergTableWriter,
SpillIcebergTableSinkOperatorX>;
-
} // namespace doris
diff --git a/be/src/exec/sink/viceberg_merge_sink.cpp
b/be/src/exec/sink/viceberg_merge_sink.cpp
index 4548f0449b6..c4eb9709e2e 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/async_result_writer.cpp
b/be/src/exec/sink/writer/async_result_writer.cpp
index 6397dc660e0..84f9fe18ff5 100644
--- a/be/src/exec/sink/writer/async_result_writer.cpp
+++ b/be/src/exec/sink/writer/async_result_writer.cpp
@@ -81,7 +81,6 @@ AsyncResultWriter::QueuedBlock
AsyncResultWriter::_get_block_from_queue() {
DCHECK(!_data_queue.empty());
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();
@@ -92,17 +91,6 @@ AsyncResultWriter::QueuedBlock
AsyncResultWriter::_get_block_from_queue() {
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) {
// Attention!!!
// AsyncResultWriter::open is called asynchronously,
@@ -196,7 +184,6 @@ void AsyncResultWriter::process_block(RuntimeState* state,
RuntimeProfile* opera
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() || _data_queue_is_available()) {
_dependency->set_ready();
@@ -211,11 +198,9 @@ void AsyncResultWriter::process_block(RuntimeState* state,
RuntimeProfile* opera
// 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();
}
{
@@ -225,7 +210,6 @@ void AsyncResultWriter::process_block(RuntimeState* state,
RuntimeProfile* opera
_memory_used_counter->update(-queued.block->allocated_bytes());
}
});
- _queue_admission.finish_processing();
_dependency->set_ready();
}
diff --git a/be/src/exec/sink/writer/async_result_writer.h
b/be/src/exec/sink/writer/async_result_writer.h
index fe851a9171a..8366e599f70 100644
--- a/be/src/exec/sink/writer/async_result_writer.h
+++ b/be/src/exec/sink/writer/async_result_writer.h
@@ -72,10 +72,6 @@ public:
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;
@@ -92,14 +88,13 @@ private:
void process_block(RuntimeState* state, RuntimeProfile* operator_profile);
[[nodiscard]] bool _data_queue_is_available() const {
- return _queue_admission.is_available(_data_queue.size());
+ return _data_queue.size() < ASYNC_WRITER_QUEUE_SIZE;
}
[[nodiscard]] bool _is_finished() const { return !_writer_status.ok() ||
_eos; }
void _set_ready_to_finish();
void _return_free_block(std::unique_ptr<Block>);
QueuedBlock _get_block_from_queue();
- void _notify_block_processed();
std::mutex _m;
std::condition_variable _cv;
@@ -107,7 +102,6 @@ private:
// Default value is ok
AtomicStatus _writer_status;
bool _eos = false;
- AsyncWriterQueueAdmission _queue_admission;
std::atomic_bool _low_memory_mode = false;
std::shared_ptr<Dependency> _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
index b5a73cb72aa..64020120ef0 100644
--- a/be/src/exec/sink/writer/async_writer_queue_admission.h
+++ b/be/src/exec/sink/writer/async_writer_queue_admission.h
@@ -23,24 +23,6 @@ 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 <typename Queue, typename BeforeRelease>
void drain_async_writer_queue(Queue& queue, BeforeRelease before_release) {
for (const auto& queued : queue) {
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 e99fabe9a73..c9a6426db7c 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp
+++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.cpp
@@ -388,8 +388,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 b41c3182843..0cc61148e45 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h
+++ b/be/src/exec/sink/writer/iceberg/viceberg_sort_writer.h
@@ -101,7 +101,7 @@ 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;
@@ -112,9 +112,11 @@ public:
size_t incoming_bytes)
const;
// Called by the memory management system to trigger spilling data to disk
- Status trigger_spill();
+ virtual Status trigger_spill();
private:
+ friend class IcebergTableSinkOperatorTest;
+
void _include_spill_merge_reservation(RuntimeState* state, bool eos,
SorterReserveMemory* reservation)
const;
@@ -161,6 +163,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();
@@ -179,12 +183,12 @@ private:
std::unique_ptr<FullSorter> _sorter;
std::unique_ptr<VSortedRunMerger> _merger;
- // Serialize all accesses to _sorter because async writes and revoke
spills run on
- // different thread pools but touch the same FullSorter instance.
+ // Serialize all accesses to _sorter because pipeline-task writes and
revoke spills run on
+ // different threads but touch the same FullSorter instance.
mutable std::mutex _sorter_mutex;
// Set to true once close() has finished tearing down the sorter /
underlying writer.
- // Late-arriving revoke spills (which run on a different thread than the
async writer)
+ // Late-arriving revoke spills (which run on a different thread than the
pipeline task)
// must become no-ops after close, otherwise they would write to a fresh
spill stream
// whose data never gets merged out (close has already produced the final
output and
// cleaned up spill files).
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 95f78cc0678..61cb105460d 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp
+++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.cpp
@@ -49,10 +49,8 @@
namespace doris {
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);
_active_writers.store(std::make_shared<ActiveWriterSnapshot>());
}
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 48b7d24e126..36eb5c08aa0 100644
--- a/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
+++ b/be/src/exec/sink/writer/iceberg/viceberg_table_writer.h
@@ -21,9 +21,9 @@
#include <gtest/gtest_prod.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"
@@ -38,6 +38,7 @@ class FileSystem;
}
class ObjectPool;
+class RowDescriptor;
class RuntimeState;
class IColumn;
@@ -45,29 +46,28 @@ 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);
+ VIcebergTableWriter(const TDataSink& t_sink, const VExprContextSPtrs&
output_exprs);
- ~VIcebergTableWriter() = default;
+ 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; }
@@ -83,8 +83,10 @@ private:
FRIEND_TEST(VIcebergTableWriterTest,
RejectsNestedRequiredNullBeforeWriterCreation);
// The lifecycle fixture inspects snapshots to verify that cross-thread
writer ownership stays stable.
friend class VIcebergTableWriterLifecycleTest;
+ friend class IcebergTableSinkOperatorTest;
- // The spill thread needs a stable view of every partition sorter, while
the async writer owns the map.
+ // Memory revocation needs a stable view of every partition sorter, while
the pipeline task owns
+ // the map.
doris::atomic_shared_ptr<ActiveWriterSnapshot> _active_writers;
class IcebergPartitionColumn {
public:
@@ -157,9 +159,11 @@ private:
void _cleanup_closed_files();
void _transfer_closed_files_to_report_cleanup();
+ const VExprContextSPtrs& _vec_output_expr_ctxs;
// Currently it is a copy, maybe it is better to use move semantics to
eliminate it.
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/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..9273040a8a9
--- /dev/null
+++ b/be/test/exec/operator/iceberg_table_sink_operator_test.cpp
@@ -0,0 +1,591 @@
+// 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 {
+
+constexpr size_t MB = 1024 * 1024;
+
+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) {
+ 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->_publish_active_writers();
+ }
+
+ 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, NormalSinkFinalizesEmptyEosAfterData) {
+ 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());
+ EXPECT_EQ(fake_state->close_count, 0);
+ // The pipeline task releases the EOS reservation as soon as sink()
returns, so the writer
+ // must already be closed, with its buffered files flushed, when the EOS
call returns.
+ 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, 1);
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+ EXPECT_TRUE(fake_state->close_inputs.front().ok());
+ EXPECT_EQ(fake_state->cleanup_count, 0);
+
+ 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, NormalSinkFinalizesNonEmptyEos) {
+ 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});
+ 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, NormalSinkEosObservesCancellation) {
+ 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 before EOS"));
+ Block empty = ColumnHelper::create_block<DataTypeInt32>({});
+ Status sink_status = local_state.sink(&state, &empty, true);
+ EXPECT_TRUE(sink_status.is<ErrorCode::CANCELLED>()) <<
sink_status.to_string();
+ ASSERT_EQ(fake_state->close_inputs.size(), 1);
+ EXPECT_TRUE(fake_state->close_inputs.front().is<ErrorCode::CANCELLED>());
+
+ Status close_status = local_state.close(&state, sink_status);
+ 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>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest, NormalSinkWriteErrorClosesWithError) {
+ 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->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,
NormalSinkCloseErrorPropagatesAndCleansUp) {
+ 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_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,
NormalSinkCancellationAfterEosDeletesDeferredFiles) {
+ 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 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 normal 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>());
+}
+
+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);
+ const size_t admission_floor =
state.minimum_operator_memory_required_bytes();
+
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true, nullptr),
admission_floor);
+
+ // The writer is closed inside the EOS sink() call, so the EOS reservation
must already cover
+ // the final spill merge of the largest partition.
+ auto empty_writer = std::make_shared<VIcebergSortWriter>(nullptr,
TSortInfo {}, 0);
+ auto spilled_writer = std::make_shared<VIcebergSortWriter>(nullptr,
TSortInfo {}, 0);
+ initialize_sort_writer_for_write(spilled_writer.get(), &state,
+ std::make_shared<DataTypeInt32>());
+ auto* spilled_writer_ptr = spilled_writer.get();
+ mark_spilled(spilled_writer.get());
+ set_partition_writers(&local_state, std::move(empty_writer),
std::move(spilled_writer));
+ // One spill run: one 8 MiB input buffer plus one 8 MiB output buffer.
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true, nullptr), 16 * MB
+ admission_floor);
+
+ mark_spilled(spilled_writer_ptr, 7);
+ // Eight runs fit in the 64 MiB merge limit: eight inputs plus one output
buffer.
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true, nullptr), 72 * MB
+ admission_floor);
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, false, nullptr),
admission_floor);
+
+ TQueryOptions query_options = state.query_options();
+ query_options.__set_spill_buffer_size_bytes(256 * MB);
+ query_options.__set_spill_sort_merge_mem_limit_bytes(MB);
+ state.set_query_options(query_options);
+ // The fan-in never drops below two inputs, even when one buffer exceeds
the merge limit.
+ EXPECT_EQ(local_state.get_reserve_mem_size(&state, true, nullptr), 768 *
MB + admission_floor);
+ 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(MB);
+ 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));
+ // Wide rows must keep the final merge output within the spill-buffer
estimate.
+ EXPECT_LT(expected_batch_rows, static_cast<size_t>(state.batch_size()));
+ EXPECT_EQ(final_merge_batch_row_count(writer), expected_batch_rows);
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SpillSinkAccountsAllPartitionsAndRevokesLargest) {
+ 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 = MB;
+ 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));
+
+ EXPECT_EQ(local_state.get_revocable_mem_size(&state),
+ large_state->data_size + small_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),
small_state->data_size);
+
+ ASSERT_TRUE(local_state.close(&state, Status::OK()).ok());
+ // Memory callbacks that arrive after close see the released writer and do
nothing.
+ EXPECT_EQ(local_state.get_revocable_mem_size(&state), 0);
+ ASSERT_TRUE(local_state.revoke_memory(&state).ok());
+ EXPECT_EQ(small_state->spill_count, 0);
+}
+
+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>());
+}
+
+TEST_F(IcebergTableSinkOperatorTest,
SpillSinkObservesCancellationDuringWriterClose) {
+ 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_hook = [&state]() {
+ state.cancel(Status::Cancelled("cancel during spill 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>());
+}
+
+} // namespace doris
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
index c5484468931..71ece22689b 100644
--- a/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp
+++ b/be/test/exec/operator/spill_iceberg_table_sink_operator_test.cpp
@@ -95,25 +95,6 @@ TEST(SpillIcebergTableSinkOperatorTest,
ReservesAllMergeInputsAndOutputAtEos) {
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 {
diff --git
a/be/test/exec/sink/writer/external_table_writer_partition_limit_test.cpp
b/be/test/exec/sink/writer/external_table_writer_partition_limit_test.cpp
index 26f3de403a6..de828950adf 100644
--- a/be/test/exec/sink/writer/external_table_writer_partition_limit_test.cpp
+++ b/be/test/exec/sink/writer/external_table_writer_partition_limit_test.cpp
@@ -95,7 +95,7 @@ TEST_P(ExternalTableWriterPartitionLimitTest,
IcebergRejectsNextPartitionBeforeO
sink.__set_iceberg_table_sink(TIcebergTableSink());
VExprContextSPtrs exprs;
iceberg::PartitionField field(1, 1000, "p", "identity");
- VIcebergTableWriter writer(sink, exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(sink, exprs);
// open() pins the table schema, and every written block is validated
against it.
std::vector<iceberg::NestedField> columns;
columns.emplace_back(true, 1, "p",
std::make_unique<iceberg::IntegerType>(), std::nullopt);
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
index ac6c100b37e..c19a00389ba 100644
--- a/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp
+++ b/be/test/exec/sink/writer/iceberg/iceberg_table_writer_test.cpp
@@ -91,7 +91,7 @@ protected:
};
TEST_F(VIcebergTableWriterLifecycleTest,
RejectsCoordinatorWithoutExternalFileReportAck) {
- VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr);
+ VIcebergTableWriter writer(make_sink(), {});
RuntimeState state;
RuntimeProfile profile("test");
@@ -102,7 +102,7 @@ TEST_F(VIcebergTableWriterLifecycleTest,
RejectsCoordinatorWithoutExternalFileRe
}
TEST_F(VIcebergTableWriterLifecycleTest, SelectBlockUsesRowPermutation) {
- VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr);
+ VIcebergTableWriter writer(make_sink(), {});
auto values = ColumnInt32::create();
values->insert_value(10);
values->insert_value(20);
@@ -121,7 +121,7 @@ TEST_F(VIcebergTableWriterLifecycleTest,
SelectBlockUsesRowPermutation) {
}
TEST_F(VIcebergTableWriterLifecycleTest,
ColdReserveCoversManyRealPartitionSelections) {
- VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr);
+ VIcebergTableWriter writer(make_sink(), {});
auto values = ColumnString::create();
for (size_t i = 0; i < 128; ++i) {
const std::string value(256 + i, static_cast<char>('a' + i % 26));
@@ -142,7 +142,7 @@ TEST_F(VIcebergTableWriterLifecycleTest,
ColdReserveCoversManyRealPartitionSelec
}
TEST_F(VIcebergTableWriterLifecycleTest,
ActiveWriterSnapshotContainsEveryOpenPartition) {
- VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr);
+ VIcebergTableWriter writer(make_sink(), {});
add_writer(&writer, "p=1");
add_writer(&writer, "p=2");
@@ -153,7 +153,7 @@ TEST_F(VIcebergTableWriterLifecycleTest,
ActiveWriterSnapshotContainsEveryOpenPa
}
TEST_F(VIcebergTableWriterLifecycleTest,
LoadedSnapshotRetainsWritersDuringConcurrentPublication) {
- VIcebergTableWriter writer(make_sink(), {}, nullptr, nullptr);
+ VIcebergTableWriter writer(make_sink(), {});
std::atomic<int> destroyed = 0;
add_writer(&writer, "p=1",
std::make_shared<FakePartitionWriter>(&destroyed));
publish_active_writers(&writer);
diff --git a/be/test/exec/sink/writer/iceberg/viceberg_table_writer_test.cpp
b/be/test/exec/sink/writer/iceberg/viceberg_table_writer_test.cpp
index 7ccb060ef5d..2f136adcdae 100644
--- a/be/test/exec/sink/writer/iceberg/viceberg_table_writer_test.cpp
+++ b/be/test/exec/sink/writer/iceberg/viceberg_table_writer_test.cpp
@@ -71,7 +71,7 @@ TEST(VIcebergTableWriterTest,
RejectsRequiredNullBeforeWriterCreation) {
TIcebergTableSink sink;
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
- VIcebergTableWriter writer(data_sink, {}, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, {});
writer._schema = schema;
auto type = make_nullable(std::make_shared<DataTypeInt32>());
@@ -102,7 +102,7 @@ TEST(VIcebergTableWriterTest,
RejectsNestedRequiredNullBeforeWriterCreation) {
TIcebergTableSink sink;
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
- VIcebergTableWriter writer(data_sink, {}, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, {});
writer._schema = schema;
auto child_type = make_nullable(std::make_shared<DataTypeInt32>());
@@ -144,7 +144,7 @@ TEST(VIcebergTableWriterTest,
TimestampIdentityPreservesUtcCommitAndNull) {
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
VExprContextSPtrs exprs;
- VIcebergTableWriter writer(data_sink, exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, exprs);
RuntimeState state;
state.set_timezone("America/New_York");
writer._state = &state;
@@ -228,7 +228,7 @@ TEST(VIcebergTableWriterTest,
StaticNullBinaryPartitionsKeepNullPathsAndCommitVa
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
VExprContextSPtrs exprs;
- VIcebergTableWriter writer(data_sink, exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, exprs);
writer._schema = schema;
auto type = make_nullable(std::make_shared<DataTypeVarbinary>());
writer._iceberg_partition_columns.emplace_back(
@@ -292,7 +292,7 @@ TEST(VIcebergTableWriterTest,
BinaryIdentityStaticAndDynamicRouting) {
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
VExprContextSPtrs exprs;
- VIcebergTableWriter writer(data_sink, exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, exprs);
writer._schema = schema;
writer._iceberg_partition_columns.emplace_back(
spec->fields()[0], TYPE_VARBINARY, 0, std::vector<size_t> {},
@@ -338,7 +338,7 @@ TEST(VIcebergTableWriterTest, RejectMissingPartitionSource)
{
TIcebergTableSink iceberg_sink;
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(iceberg_sink);
- VIcebergTableWriter writer(data_sink, {}, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, {});
writer._schema = schema;
writer._partition_spec = iceberg::PartitionSpecParser::from_json(schema,
spec_json);
@@ -370,7 +370,7 @@ TEST(VIcebergTableWriterTest,
VoidPartitionFieldNeedsNoSourceColumn) {
TIcebergTableSink iceberg_sink;
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(iceberg_sink);
- VIcebergTableWriter writer(data_sink, output_exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, output_exprs);
writer._schema = schema;
writer._partition_spec = iceberg::PartitionSpecParser::from_json(schema,
spec_json);
writer._iceberg_partition_columns = writer._to_iceberg_partition_columns();
@@ -413,7 +413,7 @@ TEST(VIcebergTableWriterTest,
ResolvesNestedPartitionSource) {
TIcebergTableSink iceberg_sink;
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(iceberg_sink);
- VIcebergTableWriter writer(data_sink, output_exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, output_exprs);
writer._schema = schema;
writer._partition_spec = iceberg::PartitionSpecParser::from_json(schema,
spec_json);
@@ -462,7 +462,7 @@ TEST(VIcebergTableWriterTest,
UuidStaticPartitionAcceptsCanonicalText) {
TDataSink data_sink;
data_sink.__set_iceberg_table_sink(sink);
VExprContextSPtrs exprs;
- VIcebergTableWriter writer(data_sink, exprs, nullptr, nullptr);
+ VIcebergTableWriter writer(data_sink, exprs);
writer._schema = schema;
writer._iceberg_partition_columns.emplace_back(
spec->fields()[0], TYPE_VARBINARY, 0, std::vector<size_t> {},
diff --git a/be/test/exec/sort/full_sort_test.cpp
b/be/test/exec/sort/full_sort_test.cpp
index bd8f91b03cc..963b1820a98 100644
--- a/be/test/exec/sort/full_sort_test.cpp
+++ b/be/test/exec/sort/full_sort_test.cpp
@@ -99,7 +99,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, EosReservationIncludesForcedSortBelowAppendThresholds) {
diff --git
a/regression-test/data/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.out
b/regression-test/data/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.out
new file mode 100644
index 00000000000..c948f4b72fc
--- /dev/null
+++
b/regression-test/data/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.out
@@ -0,0 +1,16 @@
+-- This file is automatically generated. You should know what you did if you
want to edit this
+-- !sorted_spill_summary --
+R0 20000 20000 0 59997 599970000 4276294
+R1 20000 20000 1 59998 599990000 4276298
+R2 20000 20000 2 59999 600010000 4276298
+
+-- !sorted_spill_rows --
+0 R0
payload-0-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+15838 R1
payload-15838-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+23757 R0
payload-23757-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+31676 R2
payload-31676-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+39595 R1
payload-39595-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+47514 R0
payload-47514-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+55433 R2
payload-55433-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+7919 R2
payload-7919-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+
diff --git
a/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.groovy
b/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.groovy
new file mode 100644
index 00000000000..b8749a4501d
--- /dev/null
+++
b/regression-test/suites/external_table_p0/iceberg/write/test_iceberg_write_sorted_spill.groovy
@@ -0,0 +1,114 @@
+// 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.
+
+// Sorted Iceberg writes run synchronously in the spill sink. Forced spilling
makes every
+// partition sorter write several spill runs, so the EOS sink call has to
merge them before the
+// data files are committed.
+suite("test_iceberg_write_sorted_spill",
+ "p0,external,iceberg,external_docker,external_docker_iceberg") {
+ String enabled = context.config.otherConfigs.get("enableIcebergTest")
+ if (enabled == null || !enabled.equalsIgnoreCase("true")) {
+ logger.info("disable iceberg test")
+ return
+ }
+
+ String restPort = context.config.otherConfigs.get("iceberg_rest_uri_port")
+ String minioPort = context.config.otherConfigs.get("iceberg_minio_port")
+ String externalEnvIp = context.config.otherConfigs.get("externalEnvIp")
+ String catalogName = "test_iceberg_write_sorted_spill"
+ String dbName = "iceberg_write_sorted_spill_db"
+
+ sql """drop catalog if exists ${catalogName}"""
+ sql """
+ create catalog ${catalogName} properties (
+ "type" = "iceberg",
+ "iceberg.catalog.type" = "rest",
+ "uri" = "http://${externalEnvIp}:${restPort}";,
+ "s3.access_key" = "admin",
+ "s3.secret_key" = "password",
+ "s3.endpoint" = "http://${externalEnvIp}:${minioPort}";,
+ "s3.region" = "us-east-1",
+ "meta.cache.iceberg.table.ttl-second" = "0",
+ "meta.cache.iceberg.schema.ttl-second" = "0"
+ )
+ """
+ sql """switch ${catalogName}"""
+ sql """drop database if exists ${dbName} force"""
+ sql """create database ${dbName}"""
+ sql """use ${dbName}"""
+
+ sql """drop table if exists sorted_spill"""
+ sql """
+ create table sorted_spill (
+ id bigint,
+ region string,
+ payload string
+ )
+ order by (id desc nulls last)
+ partition by list (region) ()
+ properties (
+ "format-version" = "2",
+ "write.format.default" = "parquet"
+ )
+ """
+
+ // Each partition sorter receives several MiB, while forced spilling
revokes the sink as soon
+ // as it holds 1 MiB of sortable data.
+ sql """set enable_spill = true"""
+ sql """set enable_force_spill = true"""
+ sql """set enable_reserve_memory = true"""
+ sql """set spill_min_revocable_mem = 1"""
+ sql """
+ insert into sorted_spill
+ select number,
+ concat('R', number % 3),
+ concat('payload-', number, '-', repeat('x', 200))
+ from numbers('number' = '60000')
+ """
+ sql """set enable_force_spill = false"""
+ sql """set enable_spill = false"""
+
+ order_qt_sorted_spill_summary """
+ select region, count(*), count(distinct id), min(id), max(id), sum(id),
+ sum(length(payload))
+ from sorted_spill
+ group by region
+ """
+ order_qt_sorted_spill_rows """
+ select id, region, payload
+ from sorted_spill
+ where id % 7919 = 0
+ """
+
+ // Every committed data file must keep the table order after the spill
runs are merged.
+ spark_iceberg """refresh table demo.${dbName}.sorted_spill"""
+ def sparkSummary = spark_iceberg """
+ select count(*), count(distinct _file), sum(id)
+ from demo.${dbName}.sorted_spill
+ """
+ assertEquals(60000L, sparkSummary[0][0] as long)
+ assertEquals(1799970000L, sparkSummary[0][2] as long)
+ def unsortedRows = spark_iceberg """
+ select count(*)
+ from (
+ select id, lag(id) over (partition by _file order by _pos) as
previous_id
+ from demo.${dbName}.sorted_spill
+ ) ordered_rows
+ where previous_id < id
+ """
+ assertEquals(0L, unsortedRows[0][0] as long)
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
