diff --git a/be/src/exec/operator/streaming_aggregation_operator.cpp b/be/src/exec/operator/streaming_aggregation_operator.cpp index f7549c1572423f..5c34be04aeaa8c 100644 --- a/be/src/exec/operator/streaming_aggregation_operator.cpp +++ b/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(); 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(), diff --git a/be/src/exec/operator/streaming_aggregation_operator.h b/be/src/exec/operator/streaming_aggregation_operator.h index 48fae130213e2f..6f065f11bcdfd7 100644 --- a/be/src/exec/operator/streaming_aggregation_operator.h +++ b/be/src/exec/operator/streaming_aggregation_operator.h @@ -265,8 +265,8 @@ class StreamingAggOperatorX MOCK_REMOVE(final) : public StatefulOperatorX _aggregate_evaluators; diff --git a/be/test/exec/operator/streaming_agg_operator_test.cpp b/be/test/exec/operator/streaming_agg_operator_test.cpp index 7e2fbb6a20b16a..9df1f60d93c4ca 100644 --- a/be/test/exec/operator/streaming_agg_operator_test.cpp +++ b/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(_should_expand_preagg_hash_tables()); // mock the function - static_cast(_memory_usage()); // mock the function - static_cast( - 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()), 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({2, 2, 2, 2, 4, 4}), ColumnHelper::create_column_with_name({1, 1, 100, 100, 100, 1000})}; diff --git a/be/test/testutil/mock/mock_runtime_state.h b/be/test/testutil/mock/mock_runtime_state.h index 82c7d788943e4f..fd7e8882384c46 100644 --- a/be/test/testutil/mock/mock_runtime_state.h +++ b/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; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index c5b503e9f144b0..349d29d9359a34 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/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 diff --git a/gensrc/thrift/PaloInternalService.thrift b/gensrc/thrift/PaloInternalService.thrift index 07944139285228..7281199748a947 100644 --- a/gensrc/thrift/PaloInternalService.thrift +++ b/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.