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

be/src/exec/operator/iceberg_sorter_reserve_memory.h

@@ -0,0 +1,132 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#pragma once
+
+#include <algorithm>
+#include <limits>
+#include <vector>
+
+namespace doris {
+
+class Block;
+
+struct IcebergSorterReserveMemory {
+    size_t retained_growth = 0;
+    size_t retained_growth_trigger_bytes = 0;
+    size_t transient_workspace = 0;
+};
+
+inline size_t iceberg_saturating_add(size_t lhs, size_t rhs) {
+    return std::min(std::numeric_limits<size_t>::max() - lhs, rhs) + lhs;
+}
+
+inline size_t bounded_iceberg_reserve_size(
+        const std::vector<IcebergSorterReserveMemory>& per_partition_reservations,
+        size_t incoming_rows = std::numeric_limits<size_t>::max(),
+        size_t incoming_bytes = std::numeric_limits<size_t>::max()) {
+    size_t transient_workspace = 0;
+    for (const auto& reservation : per_partition_reservations) {
+        transient_workspace = std::max(transient_workspace, reservation.transient_workspace);
+    }
+
+    std::vector<const IcebergSorterReserveMemory*> growth_candidates;
+    growth_candidates.reserve(per_partition_reservations.size());
+    for (const auto& reservation : per_partition_reservations) {
+        if (reservation.retained_growth > 0) {
+            growth_candidates.push_back(&reservation);
+        }
+    }
+
+    std::sort(growth_candidates.begin(), growth_candidates.end(),
+              [](const auto* lhs, const auto* rhs) {
+                  return lhs->retained_growth > rhs->retained_growth;
+              });
+    size_t row_bound = 0;
+    for (size_t i = 0; i < std::min(incoming_rows, growth_candidates.size()); ++i) {
+        row_bound = iceberg_saturating_add(row_bound, growth_candidates[i]->retained_growth);
+    }
+
+    size_t byte_bound = 0;
+    std::vector<const IcebergSorterReserveMemory*> positive_trigger_candidates;
+    positive_trigger_candidates.reserve(growth_candidates.size());
+    for (const auto* reservation : growth_candidates) {
+        if (reservation->retained_growth_trigger_bytes == 0) {
+            byte_bound = iceberg_saturating_add(byte_bound, reservation->retained_growth);
+        } else {
+            positive_trigger_candidates.push_back(reservation);
+        }
+    }
+    std::sort(positive_trigger_candidates.begin(), positive_trigger_candidates.end(),
+              [](const auto* lhs, const auto* rhs) {
+                  return static_cast<unsigned __int128>(lhs->retained_growth) *
+                                 rhs->retained_growth_trigger_bytes >
+                         static_cast<unsigned __int128>(rhs->retained_growth) *
+                                 lhs->retained_growth_trigger_bytes;
+              });
+    size_t remaining_bytes = incoming_bytes;
+    for (const auto* reservation : positive_trigger_candidates) {
+        if (reservation->retained_growth_trigger_bytes <= remaining_bytes) {
+            byte_bound = iceberg_saturating_add(byte_bound, reservation->retained_growth);
+            remaining_bytes -= reservation->retained_growth_trigger_bytes;
+            continue;
+        }
+        const auto numerator =
+                static_cast<unsigned __int128>(reservation->retained_growth) * remaining_bytes +
+                reservation->retained_growth_trigger_bytes - 1;
+        const auto fractional_growth =
+                std::min<unsigned __int128>(numerator / reservation->retained_growth_trigger_bytes,
+                                            std::numeric_limits<size_t>::max());
+        byte_bound = iceberg_saturating_add(byte_bound, static_cast<size_t>(fractional_growth));
+        break;
+    }
+
+    // A block's rows and bytes are divided across partition sorters. The two fractional-relaxation
+    // bounds avoid charging the complete input block to every active partition while remaining safe.
+    const size_t retained_growth = std::min(row_bound, byte_bound);
+    return iceberg_saturating_add(retained_growth, transient_workspace);
+}
+
+inline size_t iceberg_reserve_size(
+        const std::vector<IcebergSorterReserveMemory>& per_partition_reservations,
+        size_t incoming_block_reserve, size_t incoming_rows = std::numeric_limits<size_t>::max(),
+        size_t incoming_bytes = std::numeric_limits<size_t>::max()) {
+    size_t sorter_reserve =
+            bounded_iceberg_reserve_size(per_partition_reservations, incoming_rows, incoming_bytes);
+    // The incoming block creates cold partition writers before they can appear in the published snapshot.
+    return iceberg_saturating_add(sorter_reserve, incoming_block_reserve);
+}
+
+size_t iceberg_cold_writer_reserve_size(const Block& block, size_t writer_workspace_bytes);
+
+inline size_t iceberg_spill_merge_workspace(size_t spill_file_count, size_t spill_buffer_bytes,
+                                            size_t merge_limit_bytes) {
+    if (spill_file_count == 0 || spill_buffer_bytes == 0) {
+        return 0;
+    }
+    const size_t max_fan_in = std::max<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 dispatch_copies = block_bytes > std::numeric_limits<size_t>::max() / 4
+                                           ? std::numeric_limits<size_t>::max()
+                                           : block_bytes * 4;
+    size_t reserve = iceberg_saturating_add(writer_workspace_bytes, dispatch_copies);
+    // A transform, selected blocks, and retained sorter copies can coexist; one extra block-sized
+    // allowance covers allocator fragmentation when a block is split into many tiny partitions.
+    return iceberg_saturating_add(reserve, row_index_bytes);
+}
+
 SpillIcebergTableSinkLocalState::SpillIcebergTableSinkLocalState(DataSinkOperatorXBase* parent,
                                                                  RuntimeState* state)
         : Base(parent, state) {}
 
 Status SpillIcebergTableSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& info) {
     RETURN_IF_ERROR(Base::init(state, info));
+    // Admission samples async sorter state, so the next block must wait until the prior append publishes it.
+    _writer->wait_for_processing_before_next_sink();
     SCOPED_TIMER(exec_time_counter());
     SCOPED_TIMER(_init_timer);
 
@@ -51,51 +68,75 @@ 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;
+    const size_t incoming_rows = block == nullptr ? 0 : block->rows();
+    const size_t incoming_bytes = block == nullptr ? 0 : block->allocated_bytes();
+    auto active_writers = _writer->active_writers();
+    per_partition_reservations.reserve(active_writers->size());
+    for (const auto& writer : *active_writers) {
+        if (auto* sort_writer = dynamic_cast<VIcebergSortWriter*>(writer.get())) {
+            auto reservation = sort_writer->get_reserve_mem_size_components(
+                    state, eos, incoming_rows, incoming_bytes);
+            per_partition_reservations.push_back(
+                    {.retained_growth = reservation.retained_growth,
+                     .retained_growth_trigger_bytes = reservation.retained_growth_trigger_bytes,
+                     .transient_workspace = reservation.transient_workspace});
+        }
     }
-
-    return sort_writer->get_reserve_mem_size(state, eos);
+    // Column growth remains in every touched sorter, while sorting workspace is reused by serial dispatch.
+    // The final queued item may contain rows and also owns the reservation used by async finish().
+    const size_t incoming_reserve =
+            block == nullptr ? state->minimum_operator_memory_required_bytes()
+                             : iceberg_cold_writer_reserve_size(
+                                       *block, state->minimum_operator_memory_required_bytes());
+    return iceberg_reserve_size(per_partition_reservations, incoming_reserve, incoming_rows,
+                                incoming_bytes);
 }
 
 size_t SpillIcebergTableSinkLocalState::get_revocable_mem_size(RuntimeState* state) const {
     if (!_writer) {
         return 0;
     }
-    auto current_writer = _writer->current_writer();
-    auto* sort_writer = dynamic_cast<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 +166,10 @@ Status SpillIcebergTableSinkOperatorX::sink_impl(RuntimeState* state, Block* in_
     return local_state.sink(state, in_block, eos);
 }
 
-size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos) {
+size_t SpillIcebergTableSinkOperatorX::get_reserve_mem_size(RuntimeState* state, bool eos,
+                                                            const Block* block) {
     auto& local_state = get_local_state(state);
-    return local_state.get_reserve_mem_size(state, eos);
+    return local_state.get_reserve_mem_size(state, eos, block);
 }
 
 size_t SpillIcebergTableSinkOperatorX::revocable_mem_size(RuntimeState* state) const {

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