fix(pipeline): limit streaming aggregation memory without spill

be/src/exec/operator/streaming_aggregation_operator.cpp

@@ -290,8 +290,8 @@ bool StreamingAggLocalState::_should_not_do_pre_agg(size_t rows) {
     auto& p = Base::_parent->template cast<StreamingAggOperatorX>();
     bool ret_flag = false;
     const auto spill_streaming_agg_mem_limit = p._spill_streaming_agg_mem_limit;
-    const bool used_too_much_memory =
-            spill_streaming_agg_mem_limit > 0 && _memory_usage() > spill_streaming_agg_mem_limit;
+    const bool used_too_much_memory = spill_streaming_agg_mem_limit > 0 &&
+                                      _memory_usage() >= spill_streaming_agg_mem_limit;
     std::visit(
             Overload {
                     [&](std::monostate& arg) {
@@ -954,15 +954,13 @@ Status StreamingAggOperatorX::init(const TPlanNode& tnode, RuntimeState* state)
         _aggregate_evaluators.push_back(evaluator);
     }
 
-    if (state->enable_spill()) {
-        // If spill enabled, the streaming agg should not occupy too much memory.
-        _spill_streaming_agg_mem_limit =
-                state->query_options().__isset.spill_streaming_agg_mem_limit
-                        ? state->query_options().spill_streaming_agg_mem_limit
-                        : 0;
-    } else {
-        _spill_streaming_agg_mem_limit = 0;
-    }
+    // Streaming aggregation does not spill itself. Apply its memory limit independently of
+    // whether spill is enabled for the query. The limit is checked before processing each block,
+    // so memory can exceed it by at most the allocations made while processing one block.
+    _spill_streaming_agg_mem_limit =
+            state->query_options().__isset.spill_streaming_agg_mem_limit
+                    ? state->query_options().spill_streaming_agg_mem_limit
+                    : 0;
 
     const auto& agg_functions = tnode.agg_node.aggregate_functions;
     auto is_merge = std::any_of(agg_functions.cbegin(), agg_functions.cend(),

be/src/exec/operator/streaming_aggregation_operator.h

@@ -265,8 +265,8 @@ class StreamingAggOperatorX MOCK_REMOVE(final) : public StatefulOperatorX<Stream
     /// The total size of the row from the aggregate functions.
     size_t _total_size_of_aggregate_states = 0;
 
-    /// When spilling is enabled, the streaming agg should not occupy too much memory.
-    size_t _spill_streaming_agg_mem_limit;
+    /// Streaming aggregation switches to pass-through mode after reaching this memory limit.
+    size_t _spill_streaming_agg_mem_limit = 0;
     // group by k1,k2
     VExprContextSPtrs _probe_expr_ctxs;
     std::vector<AggFnEvaluator*> _aggregate_evaluators;

be/test/exec/operator/streaming_agg_operator_test.cpp

@@ -50,14 +50,12 @@ struct MockStreamingAggLocalState : public StreamingAggLocalState {
             : StreamingAggLocalState(state, parent) {}
 
     bool _should_not_do_pre_agg(size_t rows) override {
-        static_cast<void>(_should_expand_preagg_hash_tables()); // mock the function
-        static_cast<void>(_memory_usage());                     // mock the function
-        static_cast<void>(
-                StreamingAggLocalState::_should_not_do_pre_agg(rows)); // mock the function
-        return should_not_do_pre_agg;
+        const bool base_decision = StreamingAggLocalState::_should_not_do_pre_agg(rows);
+        return use_base_decision ? base_decision : should_not_do_pre_agg;
     }
 
     bool should_not_do_pre_agg = false;
+    bool use_base_decision = false;
 };
 
 class MockStreamingAggOperatorChildOperator : public OperatorXBase {
@@ -104,6 +102,25 @@ struct StreamingAggOperatorTest : public testing::Test {
     ObjectPool pool;
 };
 
+TEST_F(StreamingAggOperatorTest, memoryLimitWithoutSpill) {
+    constexpr int64_t memory_limit = 256 * 1024 * 1024;
+    state->set_enable_spill(false);
+    state->set_spill_streaming_agg_mem_limit(memory_limit);
+
+    TPlanNode tnode;
+    tnode.node_id = 0;
+    tnode.node_type = TPlanNodeType::AGGREGATION_NODE;
+    tnode.num_children = 1;
+    tnode.nereids_id = 0;
+    tnode.limit = -1;
+    tnode.agg_node.need_finalize = false;
+    tnode.agg_node.intermediate_tuple_id = 0;
+    tnode.agg_node.output_tuple_id = 0;
+
+    EXPECT_TRUE(op->init(tnode, state.get()).ok());
+    EXPECT_EQ(op->_spill_streaming_agg_mem_limit, memory_limit);
+}
+
 TEST_F(StreamingAggOperatorTest, test1) {
     op->_aggregate_evaluators.push_back(create_mock_agg_fn_evaluator(
             pool, MockSlotRef::create_mock_contexts(1, std::make_shared<DataTypeInt64>()), false,
@@ -216,7 +233,8 @@ TEST_F(StreamingAggOperatorTest, test2) {
     }
 
     {
-        local_state->should_not_do_pre_agg = true;
+        op->_spill_streaming_agg_mem_limit = 1;
+        local_state->use_base_decision = true;
         Block block {
                 ColumnHelper::create_column_with_name<DataTypeInt64>({2, 2, 2, 2, 4, 4}),
                 ColumnHelper::create_column_with_name<DataTypeInt64>({1, 1, 100, 100, 100, 1000})};

be/test/testutil/mock/mock_runtime_state.h

@@ -67,6 +67,10 @@ class MockRuntimeState : public RuntimeState {
 
     void set_enable_spill(bool enable) { _query_options.__set_enable_spill(enable); }
 
+    void set_spill_streaming_agg_mem_limit(int64_t limit) {
+        _query_options.__set_spill_streaming_agg_mem_limit(limit);
+    }
+
     void set_enable_strict_cast(bool enable) { _query_options.__set_enable_strict_cast(enable); }
 
     bool enable_local_exchange() const override { return true; }

fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java

@@ -3268,7 +3268,8 @@ public void setDetailShapePlanNodes(String detailShapePlanNodes) {
     @VarAttrDef.VarAttr(name = LOW_MEMORY_MODE_BUFFER_LIMIT, fuzzy = false)
     public long lowMemoryModeBufferLimit = 33554432;
 
-    // The memory limit of streaming agg when spilling is enabled
+    // The memory limit of streaming aggregation. The operator itself does not spill and switches
+    // to pass-through mode after reaching this limit, regardless of whether query spill is enabled.
     // NOTE: streaming agg operator will not spill to disk.
     @VarAttrDef.VarAttr(name = SPILL_STREAMING_AGG_MEM_LIMIT, fuzzy = false)
     public long spillStreamingAggMemLimit = 268435456; //256MB

gensrc/thrift/PaloInternalService.thrift

@@ -294,6 +294,8 @@ struct TQueryOptions {
 
   104: optional i64 min_revocable_mem = 0
 
+  // memory limit of streaming aggregation. When reached, the operator switches to
+  // pass-through mode instead of spilling, regardless of whether query spill is enabled.
   105: optional i64 spill_streaming_agg_mem_limit = 0;
 
   // max rows of each sub-queue in DataQueue.