[Enhancement] ignore the exec_mem_limit in pipeline execution (#34120)

Signed-off-by: Murphy <[email protected]>
Signed-off-by: Murphy <[email protected]>
(cherry picked from commit 915d762)

Author: murphyatwork

be/src/exec/pipeline/fragment_executor.cpp

@@ -218,7 +218,6 @@ Status FragmentExecutor::_prepare_runtime_state(ExecEnv* exec_env, const Unified
     const auto& query_globals = request.common().query_globals;
     const auto& query_options = request.common().query_options;
     const auto& t_desc_tbl = request.common().desc_tbl;
-    const int32_t degree_of_parallelism = _calc_dop(exec_env, request);
     auto& wg = _wg;
 
     _fragment_ctx->set_runtime_state(
@@ -228,17 +227,17 @@ Status FragmentExecutor::_prepare_runtime_state(ExecEnv* exec_env, const Unified
     runtime_state->set_fragment_ctx(_fragment_ctx.get());
     runtime_state->set_query_ctx(_query_ctx);
 
+    // Only consider the `query_mem_limit` variable
+    // If query_mem_limit is <= 0, it would set to -1, which means no limit
     auto* parent_mem_tracker = wg->mem_tracker();
-    auto per_instance_mem_limit = query_options.__isset.mem_limit ? query_options.mem_limit : -1;
-    auto option_query_mem_limit = query_options.__isset.query_mem_limit ? query_options.query_mem_limit : -1;
-    int64_t query_mem_limit = _query_ctx->compute_query_mem_limit(parent_mem_tracker->limit(), per_instance_mem_limit,
-                                                                  degree_of_parallelism, option_query_mem_limit);
+    int64_t option_query_mem_limit = query_options.__isset.query_mem_limit ? query_options.query_mem_limit : -1;
+    if (option_query_mem_limit <= 0) option_query_mem_limit = -1;
     int64_t big_query_mem_limit = wg->use_big_query_mem_limit() ? wg->big_query_mem_limit() : -1;
     int64_t spill_mem_limit_bytes = -1;
     if (query_options.__isset.enable_spill && query_options.enable_spill == true) {
-        spill_mem_limit_bytes = query_mem_limit * query_options.spill_mem_limit_threshold;
+        spill_mem_limit_bytes = option_query_mem_limit * query_options.spill_mem_limit_threshold;
     }
-    _query_ctx->init_mem_tracker(query_mem_limit, parent_mem_tracker, big_query_mem_limit, spill_mem_limit_bytes,
+    _query_ctx->init_mem_tracker(option_query_mem_limit, parent_mem_tracker, big_query_mem_limit, spill_mem_limit_bytes,
                                  wg.get());
 
     auto query_mem_tracker = _query_ctx->mem_tracker();

be/src/exec/pipeline/query_context.cpp

@@ -98,31 +98,6 @@ void QueryContext::cancel(const Status& status) {
     _fragment_mgr->cancel(status);
 }
 
-int64_t QueryContext::compute_query_mem_limit(int64_t parent_mem_limit, int64_t per_instance_mem_limit,
-                                              size_t pipeline_dop, int64_t option_query_mem_limit) {
-    // no mem_limit
-    if (per_instance_mem_limit <= 0 && option_query_mem_limit <= 0) {
-        return -1;
-    }
-
-    int64_t mem_limit;
-    if (option_query_mem_limit > 0) {
-        mem_limit = option_query_mem_limit;
-    } else {
-        mem_limit = per_instance_mem_limit;
-        // query's mem_limit = per-instance mem_limit * num_instances * pipeline_dop
-        static constexpr int64_t MEM_LIMIT_MAX = std::numeric_limits<int64_t>::max();
-        if (MEM_LIMIT_MAX / total_fragments() / pipeline_dop > mem_limit) {
-            mem_limit *= static_cast<int64_t>(total_fragments()) * pipeline_dop;
-        } else {
-            mem_limit = MEM_LIMIT_MAX;
-        }
-    }
-
-    // query's mem_limit never exceeds its parent's limit if it exists
-    return parent_mem_limit == -1 ? mem_limit : std::min(parent_mem_limit, mem_limit);
-}
-
 void QueryContext::init_mem_tracker(int64_t query_mem_limit, MemTracker* parent, int64_t big_query_mem_limit,
                                     int64_t spill_mem_limit, workgroup::WorkGroup* wg) {
     std::call_once(_init_mem_tracker_once, [=]() {

be/src/exec/pipeline/query_context.h

@@ -126,10 +126,7 @@ class QueryContext : public std::enable_shared_from_this<QueryContext> {
         DCHECK(_desc_tbl != nullptr);
         return _desc_tbl;
     }
-    // If option_query_mem_limit > 0, use it directly.
-    // Otherwise, use per_instance_mem_limit * num_fragments * pipeline_dop.
-    int64_t compute_query_mem_limit(int64_t parent_mem_limit, int64_t per_instance_mem_limit, size_t pipeline_dop,
-                                    int64_t option_query_mem_limit);
+
     size_t total_fragments() { return _total_fragments; }
     /// Initialize the mem_tracker of this query.
     /// Positive `big_query_mem_limit` and non-null `wg` indicate

be/src/exec/scan_node.cpp

@@ -34,6 +34,7 @@
 
 #include "exec/scan_node.h"
 
+#include "exec/pipeline/query_context.h"
 #include "exec/pipeline/scan/morsel.h"
 
 namespace starrocks {
@@ -59,7 +60,14 @@ Status ScanNode::init(const TPlanNode& tnode, RuntimeState* state) {
     if (options.__isset.scan_use_query_mem_ratio) {
         mem_ratio = options.scan_use_query_mem_ratio;
     }
-    _mem_limit = state->query_mem_tracker_ptr()->limit() * mem_ratio;
+    if (runtime_state()->query_ctx()) {
+        // Used in pipeline-engine
+        _mem_limit = state->query_ctx()->get_static_query_mem_limit() * mem_ratio;
+    } else if (runtime_state()->query_mem_tracker_ptr()) {
+        // Fallback in non-pipeline
+        _mem_limit = state->query_mem_tracker_ptr()->limit() * mem_ratio;
+    }
+
     return Status::OK();
 }
 

be/test/CMakeLists.txt

@@ -39,7 +39,6 @@ set(EXEC_FILES
         ./exec/pipeline/pipeline_file_scan_node_test.cpp
         ./exec/pipeline/pipeline_test_base.cpp
         ./exec/pipeline/query_context_manger_test.cpp
-        ./exec/pipeline/query_context_test.cpp
         ./exec/pipeline/table_function_operator_test.cpp
         ./exec/query_cache/query_cache_test.cpp
         ./exec/query_cache/transform_operator.cpp

fe/fe-core/src/main/java/com/starrocks/qe/SessionVariable.java

@@ -97,11 +97,8 @@ public class SessionVariable implements Serializable, Writable, Cloneable {
 
     /**
      * The mem limit of query on BE. It takes effects only when enabling pipeline engine.
-     * - If `query_mem_limit` > 0, use it to limit the memory of a query.
-     * The memory a query able to be used is just `query_mem_limit`.
-     * - Otherwise, use `exec_mem_limit` to limit the memory of a query.
-     * The memory a query able to be used is `exec_mem_limit * num_fragments * pipeline_dop`.
-     * To maintain compatibility, the default value is 0.
+     * If `query_mem_limit` > 0, use it to limit the memory of a query.
+     * Otherwise, no limitation
      */
     public static final String QUERY_MEM_LIMIT = "query_mem_limit";
 
@@ -671,7 +668,9 @@ public static MaterializedViewRewriteMode parse(String str) {
     @VariableMgr.VarAttr(name = ENABLE_SHARED_SCAN)
     private boolean enableSharedScan = false;
 
-    // max memory used on every backend.
+    // max memory used on each fragment instance
+    // NOTE: only used for non-pipeline engine and stream_load
+    // The pipeline engine uses the query_mem_limit
     public static final long DEFAULT_EXEC_MEM_LIMIT = 2147483648L;
     @VariableMgr.VarAttr(name = EXEC_MEM_LIMIT, flag = VariableMgr.INVISIBLE)
     public long maxExecMemByte = DEFAULT_EXEC_MEM_LIMIT;
@@ -2857,9 +2856,7 @@ public void setSkewJoinRandRange(int skewJoinRandRange) {
     public TQueryOptions toThrift() {
         TQueryOptions tResult = new TQueryOptions();
         tResult.setMem_limit(maxExecMemByte);
-        if (queryMemLimit > 0) {
-            tResult.setQuery_mem_limit(queryMemLimit);
-        }
+        tResult.setQuery_mem_limit(queryMemLimit);
 
         // Avoid integer overflow
         tResult.setQuery_timeout(Math.min(Integer.MAX_VALUE / 1000, queryTimeoutS));

fe/fe-core/src/test/java/com/starrocks/qe/scheduler/JobSpecTest.java

@@ -219,7 +219,7 @@ public void testFromBrokerLoadJobSpec() throws Exception {
         jobSpec = coordinator.getJobSpec();
         Assert.assertEquals(loadMemLimit, jobSpec.getQueryOptions().getLoad_mem_limit());
         Assert.assertTrue(jobSpec.getQueryOptions().isSetMem_limit());
-        Assert.assertFalse(jobSpec.getQueryOptions().isSetQuery_mem_limit());
+        Assert.assertTrue(jobSpec.getQueryOptions().isSetQuery_mem_limit());
     }
 
     @Test
@@ -289,7 +289,7 @@ public void testFromStreamLoadJobSpec() throws Exception {
         jobSpec = coordinator.getJobSpec();
         Assert.assertEquals(loadMemLimit, jobSpec.getQueryOptions().getLoad_mem_limit());
         Assert.assertTrue(jobSpec.getQueryOptions().isSetMem_limit());
-        Assert.assertFalse(jobSpec.getQueryOptions().isSetQuery_mem_limit());
+        Assert.assertTrue(jobSpec.getQueryOptions().isSetQuery_mem_limit());
     }
 
     @Test