[fix](iceberg) Harden external-write lifecycle, OCC, memory, and file ownership

be/src/exec/operator/iceberg_sorter_reserve_memory.h

@@ -0,0 +1,74 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <algorithm>
+#include <limits>
+#include <vector>
+
+namespace doris {
+
+class Block;
+
+struct IcebergSorterReserveMemory {
+    size_t retained_growth = 0;
+    size_t transient_workspace = 0;
+};
+
+inline size_t bounded_iceberg_reserve_size(
+        const std::vector<IcebergSorterReserveMemory>& per_partition_reservations) {
+    size_t retained_growth = 0;
+    size_t transient_workspace = 0;
+    for (const auto& reservation : per_partition_reservations) {
+        retained_growth = std::min(std::numeric_limits<size_t>::max() - retained_growth,
+                                   reservation.retained_growth) +
+                          retained_growth;
+        transient_workspace = std::max(transient_workspace, reservation.transient_workspace);
+    }
+    return std::min(std::numeric_limits<size_t>::max() - retained_growth, transient_workspace) +
+           retained_growth;
+}
+
+inline size_t iceberg_reserve_size(
+        const std::vector<IcebergSorterReserveMemory>& per_partition_reservations,
+        size_t incoming_block_bytes) {
+    size_t sorter_reserve = bounded_iceberg_reserve_size(per_partition_reservations);
+    // The incoming block creates cold partition writers before they can appear in the published snapshot.
+    return std::min(std::numeric_limits<size_t>::max() - sorter_reserve, incoming_block_bytes) +
+           sorter_reserve;
+}
+
+size_t iceberg_cold_writer_reserve_size(const Block& block, size_t writer_workspace_bytes);
+
+inline size_t iceberg_spill_merge_workspace(size_t spill_file_count, size_t spill_buffer_bytes,
+                                            size_t merge_limit_bytes) {
+    if (spill_file_count == 0 || spill_buffer_bytes == 0) {
+        return 0;
+    }
+    const size_t max_fan_in = std::max<size_t>(2, merge_limit_bytes / spill_buffer_bytes);
+    const size_t input_count = std::min(spill_file_count, max_fan_in);
+    const size_t max_size = std::numeric_limits<size_t>::max();
+    const size_t input_bytes = input_count > max_size / spill_buffer_bytes
+                                       ? max_size
+                                       : input_count * spill_buffer_bytes;
+    // VSortedRunMerger materializes one block per input cursor plus the block being emitted.
+    return input_bytes > max_size - spill_buffer_bytes ? max_size
+                                                       : input_bytes + spill_buffer_bytes;
+}
+
+} // namespace doris

be/src/exec/operator/operator.h

@@ -631,6 +631,10 @@ class DataSinkOperatorXBase : public OperatorBase {
     [[nodiscard]] virtual size_t get_reserve_mem_size(RuntimeState* state, bool eos) {
         return state->minimum_operator_memory_required_bytes();
     }
+    [[nodiscard]] virtual size_t get_reserve_mem_size(RuntimeState* state, bool eos,
+                                                      const Block* block) {
+        return get_reserve_mem_size(state, eos);
+    }
     bool is_blockable(RuntimeState* state) const override {
         return state->get_sink_local_state()->is_blockable();
     }

be/src/exec/operator/spill_iceberg_table_sink_operator.cpp

@@ -18,18 +18,35 @@
 #include "exec/operator/spill_iceberg_table_sink_operator.h"
 
 #include "common/status.h"
+#include "core/block/block.h"
 #include "exec/operator/iceberg_table_sink_operator.h"
 #include "exec/sink/writer/iceberg/viceberg_sort_writer.h"
 #include "exec/sink/writer/iceberg/viceberg_table_writer.h"
 
 namespace doris {
 
+size_t iceberg_cold_writer_reserve_size(const Block& block, size_t writer_workspace_bytes) {
+    const size_t block_bytes = block.allocated_bytes();
+    const size_t row_index_bytes =
+            std::min(std::numeric_limits<size_t>::max() / sizeof(size_t), block.rows()) *
+            sizeof(size_t);
+    const size_t selected_and_retained_bytes =
+            std::min(std::numeric_limits<size_t>::max() / 2, block_bytes) * 2;
+    size_t reserve = std::min(std::numeric_limits<size_t>::max() - writer_workspace_bytes,
+                              selected_and_retained_bytes) +
+                     writer_workspace_bytes;
+    // Cold dispatch may allocate a selected block and a retained sorter copy before publication.
+    return std::min(std::numeric_limits<size_t>::max() - reserve, row_index_bytes) + reserve;
+}
+
 SpillIcebergTableSinkLocalState::SpillIcebergTableSinkLocalState(DataSinkOperatorXBase* parent,
                                                                  RuntimeState* state)
         : Base(parent, state) {}
 
 Status SpillIcebergTableSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& info) {
     RETURN_IF_ERROR(Base::init(state, info));
+    // Admission samples async sorter state, so the next block must wait until the prior append publishes it.
+    _writer->wait_for_processing_before_next_sink();
     SCOPED_TIMER(exec_time_counter());
     SCOPED_TIMER(_init_timer);
 
@@ -51,51 +68,70 @@ bool SpillIcebergTableSinkLocalState::is_blockable() const {
     return true;
 }
 
-size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state, bool eos) {
+size_t SpillIcebergTableSinkLocalState::get_reserve_mem_size(RuntimeState* state, bool eos,
+                                                             const Block* block) {
     if (!_writer) {
         return 0;
     }
-    auto current_writer = _writer->current_writer();
-    auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(current_writer.get());
-    if (!sort_writer) {
-        return 0;
+    std::vector<IcebergSorterReserveMemory> per_partition_reservations;
+    auto active_writers = _writer->active_writers();
+    per_partition_reservations.reserve(active_writers->size());
+    for (const auto& writer : *active_writers) {
+        if (auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get())) {
+            auto reservation = sort_writer->get_reserve_mem_size_components(state, eos);
+            per_partition_reservations.push_back(
+                    {.retained_growth = reservation.retained_growth,
+                     .transient_workspace = reservation.transient_workspace});
+        }
     }
-
-    return sort_writer->get_reserve_mem_size(state, eos);
+    // Column growth remains in every touched sorter, while sorting workspace is reused by serial dispatch.
+    // The final queued item may contain rows and also owns the reservation used by async finish().
+    const size_t incoming_reserve =
+            block == nullptr ? state->minimum_operator_memory_required_bytes()
+                             : iceberg_cold_writer_reserve_size(
+                                       *block, state->minimum_operator_memory_required_bytes());
+    return iceberg_reserve_size(per_partition_reservations, incoming_reserve);
 }
 
 size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState* state) const {
     if (!_writer) {
         return 0;
     }
-    auto current_writer = _writer->current_writer();
-    auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(current_writer.get());
-    if (!sort_writer) {
-        return 0;
+    size_t revocable_size = 0;
+    // Retain the published container while the async writer may replace the current snapshot.
+    auto active_writers = _writer->active_writers();
+    for (const auto& writer : *active_writers) {
+        if (auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get())) {
+            revocable_size += sort_writer->data_size();
+        }
     }
-
-    return sort_writer->data_size();
+    return revocable_size;
 }
 
 Status SpillIcebergTableSinkLocalState::revoke_memory(RuntimeState* state) {
     RETURN_IF_CANCELLED(state);
     if (!_writer) {
         return Status::OK();
     }
-    auto current_writer = _writer->current_writer();
-    auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(current_writer.get());
-    if (!sort_writer) {
-        return Status::OK();
+    std::shared_ptr<IPartitionWriterBase> largest_writer;
+    size_t largest_size = 0;
+    // Retain the published container while the async writer may replace the current snapshot.
+    auto active_writers = _writer->active_writers();
+    for (const auto& writer : *active_writers) {
+        if (auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get())) {
+            size_t size = sort_writer->data_size();
+            if (size > largest_size) {
+                largest_size = size;
+                largest_writer = writer;
+            }
+        }
     }
-
-    auto exception_catch_func = [current_writer, sort_writer]() {
-        auto status = [&]() {
-            RETURN_IF_CATCH_EXCEPTION({ return sort_writer->trigger_spill(); });
-        }();
-        return status;
-    };
-
-    return exception_catch_func();
+    if (largest_writer != nullptr) {
+        // Repeated revocation drains the largest partition first without launching O(P) spill jobs at once.
+        auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(largest_writer.get());
+        RETURN_IF_CATCH_EXCEPTION({ RETURN_IF_ERROR(sort_writer->trigger_spill()); });
+    }
+    return Status::OK();
 }
 
 SpillIcebergTableSinkOperatorX::SpillIcebergTableSinkOperatorX(
@@ -125,9 +161,10 @@ Status SpillIcebergTableSinkOperatorX::sink_impl(RuntimeState* state, Block* in_
     return local_state.sink(state, in_block, eos);
 }
 
-size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos) {
+size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos,
+                                                            const Block* block) {
     auto& local_state = get_local_state(state);
-    return local_state.get_reserve_mem_size(state, eos);
+    return local_state.get_reserve_mem_size(state, eos, block);
 }
 
 size_t SpillIcebergTableSinkOperatorX::revocable_mem_size(RuntimeState* state) const {

be/src/exec/operator/spill_iceberg_table_sink_operator.h

@@ -18,7 +18,9 @@
 #pragma once
 
 #include <memory>
+#include <vector>
 
+#include "exec/operator/iceberg_sorter_reserve_memory.h"
 #include "exec/operator/operator.h"
 #include "exec/sink/writer/iceberg/viceberg_table_writer.h"
 
@@ -41,7 +43,7 @@ class SpillIcebergTableSinkLocalState final
     Status open(RuntimeState* state) override;
 
     bool is_blockable() const override;
-    [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos);
+    [[nodiscard]] size_t get_reserve_mem_size(RuntimeState* state, bool eos, const Block* block);
     Status revoke_memory(RuntimeState* state);
     size_t get_revocable_mem_size(RuntimeState* state) const;
 
@@ -65,7 +67,7 @@ class SpillIcebergTableSinkOperatorX final
 
     Status sink_impl(RuntimeState* state, Block* in_block, bool eos) override;
 
-    size_t get_reserve_mem_size(RuntimeState* state, bool eos) override;
+    size_t get_reserve_mem_size(RuntimeState* state, bool eos, const Block* block) override;
 
     size_t revocable_mem_size(RuntimeState* state) const override;
 
@@ -87,4 +89,4 @@ class SpillIcebergTableSinkOperatorX final
     ObjectPool* _pool = nullptr;
 };
 
-} // namespace doris
+} // namespace doris