[Feature]Support SELECT Optimizer Hints SET_VAR

docs/en/administrator-guide/variables.md

@@ -85,6 +85,18 @@ SET exec_mem_limit = 10 * 1024 * 1024 * 1024;
 SET forward_to_master = concat('tr', 'u', 'e');
 ```
 
+### Set variables in the query statement
+
+In some scenarios, we may need to set variables specifically for certain queries.
+The SET_VAR hint sets the session value of a system variable temporarily (for the duration of a single statement). Examples:
+
+```
+SELECT /*+ SET_VAR(exec_mem_limit = 8589934592) */ name FROM people ORDER BY name;
+SELECT /*+ SET_VAR(query_timeout = 1) */ sleep(3);
+```
+
+Note that the comment must start with /*+ and can only follow the SELECT.
+
 ## Supported variables
 
 * `SQL_AUTO_IS_NULL`

docs/zh-CN/administrator-guide/variables.md

@@ -85,6 +85,18 @@ SET exec_mem_limit = 10 * 1024 * 1024 * 1024;
 SET forward_to_master = concat('tr', 'u', 'e');
 ```
 
+### 在查询语句中设置变量
+
+在一些场景中，我们可能需要对某些查询有针对性的设置变量。
+通过使用SET_VAR提示可以在查询中设置会话变量（在单个语句内生效）。例子：
+
+```
+SELECT /*+ SET_VAR(exec_mem_limit = 8589934592) */ name FROM people ORDER BY name;
+SELECT /*+ SET_VAR(query_timeout = 1) */ sleep(3);
+```
+
+注意注释必须以/*+ 开头，并且只能跟随在SELECT之后。
+
 ## 支持的变量
 
 * `SQL_AUTO_IS_NULL`

fe/fe-core/src/main/cup/sql_parser.cup

@@ -20,6 +20,7 @@ package org.apache.doris.analysis;
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -255,7 +256,7 @@ terminal String KW_ADD, KW_ADMIN, KW_AFTER, KW_AGGREGATE, KW_ALL, KW_ALTER, KW_A
     KW_RANDOM, KW_RANGE, KW_READ, KW_RECOVER, KW_REGEXP, KW_RELEASE, KW_RENAME,
     KW_REPAIR, KW_REPEATABLE, KW_REPOSITORY, KW_REPOSITORIES, KW_REPLACE, KW_REPLACE_IF_NOT_NULL, KW_REPLICA, KW_RESOURCE, KW_RESOURCES, KW_RESTORE, KW_RETURNS, KW_RESUME, KW_REVOKE,
     KW_RIGHT, KW_ROLE, KW_ROLES, KW_ROLLBACK, KW_ROLLUP, KW_ROUTINE, KW_ROW, KW_ROWS,
-    KW_SCHEMA, KW_SCHEMAS, KW_SECOND, KW_SELECT, KW_SEMI, KW_SERIALIZABLE, KW_SESSION, KW_SET, KW_SETS, KW_SHOW, KW_SIGNED,
+    KW_SCHEMA, KW_SCHEMAS, KW_SECOND, KW_SELECT, KW_SEMI, KW_SERIALIZABLE, KW_SESSION, KW_SET, KW_SETS, KW_SET_VAR, KW_SHOW, KW_SIGNED,
     KW_SMALLINT, KW_SNAPSHOT, KW_SONAME, KW_SPLIT, KW_START, KW_STATUS, KW_STOP, KW_STORAGE, KW_STRING,
     KW_SUM, KW_SUPERUSER, KW_SYNC, KW_SYSTEM,
     KW_TABLE, KW_TABLES, KW_TABLET, KW_TASK, KW_TEMPORARY, KW_TERMINATED, KW_THAN, KW_TIME, KW_THEN, KW_TIMESTAMP, KW_TINYINT,
@@ -380,6 +381,7 @@ nonterminal InlineViewRef inline_view_ref;
 nonterminal JoinOperator join_operator;
 nonterminal ArrayList<String> opt_plan_hints;
 nonterminal ArrayList<String> opt_sort_hints;
+nonterminal HashMap<String, String> opt_select_hints;
 nonterminal Expr sign_chain_expr;
 nonterminal Qualifier opt_set_qualifier;
 nonterminal Operation set_op;
@@ -3298,21 +3300,37 @@ expr_or_default ::=
     ;
 
 select_clause ::=
-    KW_SELECT select_list:l
+    KW_SELECT opt_select_hints:hints select_list:l
     {:
+        l.setOptHints(hints);
         RESULT = l;
     :}
-    | KW_SELECT KW_ALL select_list:l
+    | KW_SELECT opt_select_hints:hints KW_ALL select_list:l
     {:
+        l.setOptHints(hints);
         RESULT = l;
     :}
-    | KW_SELECT KW_DISTINCT select_list:l
+    | KW_SELECT opt_select_hints:hints KW_DISTINCT select_list:l
     {:
+        l.setOptHints(hints);
         l.setIsDistinct(true);
         RESULT = l;
     :}
     ;
 
+opt_select_hints ::=
+    COMMENTED_PLAN_HINT_START KW_SET_VAR LPAREN ident_or_text:k EQUAL literal:v RPAREN COMMENTED_PLAN_HINT_END
+    {:
+        HashMap<String, String> map = new HashMap<String, String>();
+        map.put(k, v.getStringValue());
+        RESULT = map;
+    :}
+    | /* empty */
+    {:
+        RESULT = null;
+    :}
+    ;
+
 select_list ::=
     select_sublist:list
     {:

fe/fe-core/src/main/java/org/apache/doris/analysis/SelectList.java

@@ -24,12 +24,14 @@
 import com.google.common.collect.Lists;
 
 import java.util.List;
+import java.util.Map;
 
 /**
  * Select list items plus distinct clause.
  */
 public class SelectList {
     private boolean isDistinct;
+    private Map<String, String> optHints;
 
     // ///////////////////////////////////////
     // BEGIN: Members that need to be reset()
@@ -72,7 +74,15 @@ public boolean isDistinct() {
     public void setIsDistinct(boolean value) {
         isDistinct = value;
     }
-
+
+    public Map<String, String> getOptHints() {
+        return optHints;
+    }
+
+    public void setOptHints(Map<String, String> optHints) {
+        this.optHints = optHints;
+    }
+
     public void reset() {
         for (SelectListItem item : items) {
             if (!item.isStar()) {

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

@@ -36,6 +36,8 @@
 import org.apache.doris.analysis.StmtRewriter;
 import org.apache.doris.analysis.UnsupportedStmt;
 import org.apache.doris.analysis.UseStmt;
+import org.apache.doris.analysis.SetVar;
+import org.apache.doris.analysis.StringLiteral;
 import org.apache.doris.catalog.Catalog;
 import org.apache.doris.catalog.Column;
 import org.apache.doris.catalog.Database;
@@ -238,8 +240,19 @@ public void execute() throws Exception {
         context.setQueryId(new TUniqueId(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits()));
 
         try {
+            // support select hint e.g. select /*+ SET_VAR(query_timeout=1) */ sleep(3);
+            SessionVariable sessionVariable = context.getSessionVariable();
+            if (parsedStmt != null && parsedStmt instanceof SelectStmt) {
+                SelectStmt selectStmt = (SelectStmt) parsedStmt;
+                Map<String, String> optHints = selectStmt.getSelectList().getOptHints();
+                if(optHints != null) {
+                    for (String key : optHints.keySet()) {
+                        VariableMgr.setVar(sessionVariable, new SetVar(key, new StringLiteral(optHints.get(key))));
+                    }
+                }
+            }
             // analyze this query
-            analyze(context.getSessionVariable().toThrift());
+            analyze(sessionVariable.toThrift());
 
             if (isForwardToMaster()) {
                 forwardToMaster();

fe/fe-core/src/main/jflex/sql_scanner.flex

@@ -327,6 +327,7 @@ import org.apache.doris.qe.SqlModeHelper;
         keywordMap.put("serializable", new Integer(SqlParserSymbols.KW_SERIALIZABLE));
         keywordMap.put("session", new Integer(SqlParserSymbols.KW_SESSION));
         keywordMap.put("set", new Integer(SqlParserSymbols.KW_SET));
+        keywordMap.put("set_var", new Integer(SqlParserSymbols.KW_SET_VAR));
         keywordMap.put("sets", new Integer(SqlParserSymbols.KW_SETS));
         keywordMap.put("show", new Integer(SqlParserSymbols.KW_SHOW));
         keywordMap.put("signed", new Integer(SqlParserSymbols.KW_SIGNED));

fe/fe-core/src/test/java/org/apache/doris/analysis/SelectStmtTest.java

@@ -19,6 +19,8 @@
 
 import org.apache.doris.common.AnalysisException;
 import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.planner.Planner;
+import org.apache.doris.qe.VariableMgr;
 import org.apache.doris.utframe.DorisAssert;
 import org.apache.doris.utframe.UtFrameUtils;
 
@@ -453,4 +455,26 @@ public void testDeleteSign() throws Exception {
         String sql4 = " SELECT * FROM table1 table2";
         Assert.assertTrue(dorisAssert.query(sql4).explainQuery().contains("`table2`.`__DORIS_DELETE_SIGN__` = 0"));
     }
+
+    @Test
+    public void testSelectHintSetVar() throws Exception {
+        String sql = "SELECT sleep(3);";
+        Planner planner = dorisAssert.query(sql).internalExecuteOneAndGetPlan();
+        Assert.assertEquals(VariableMgr.getDefaultSessionVariable().getQueryTimeoutS(),
+                planner.getPlannerContext().getQueryOptions().query_timeout);
+
+        sql = "SELECT /*+ SET_VAR(query_timeout = 1) */ sleep(3);";
+        planner = dorisAssert.query(sql).internalExecuteOneAndGetPlan();
+        Assert.assertEquals(1, planner.getPlannerContext().getQueryOptions().query_timeout);
+
+        sql = "select * from db1.partition_table where datekey=20200726";
+        planner = dorisAssert.query(sql).internalExecuteOneAndGetPlan();
+        Assert.assertEquals(VariableMgr.getDefaultSessionVariable().getMaxExecMemByte(),
+                planner.getPlannerContext().getQueryOptions().mem_limit);
+
+        sql = "select /*+ SET_VAR(exec_mem_limit = 8589934592) */ poi_id, count(*) from db1.partition_table " +
+                "where datekey=20200726 group by 1";
+        planner = dorisAssert.query(sql).internalExecuteOneAndGetPlan();
+        Assert.assertEquals(8589934592L, planner.getPlannerContext().getQueryOptions().mem_limit);
+    }
 }

fe/fe-core/src/test/java/org/apache/doris/utframe/DorisAssert.java

@@ -23,10 +23,14 @@
 import org.apache.doris.analysis.CreateMaterializedViewStmt;
 import org.apache.doris.analysis.CreateTableStmt;
 import org.apache.doris.analysis.DropTableStmt;
+import org.apache.doris.analysis.SqlParser;
+import org.apache.doris.analysis.SqlScanner;
+import org.apache.doris.analysis.StatementBase;
 import org.apache.doris.catalog.Catalog;
 import org.apache.doris.cluster.ClusterNamespace;
 import org.apache.doris.common.AnalysisException;
 import org.apache.doris.common.Config;
+import org.apache.doris.common.util.SqlParserUtils;
 import org.apache.doris.planner.Planner;
 import org.apache.doris.qe.ConnectContext;
 import org.apache.doris.qe.QueryState;
@@ -38,6 +42,8 @@
 import org.junit.Assert;
 
 import java.io.IOException;
+import java.io.StringReader;
+import java.util.List;
 import java.util.Map;
 import java.util.stream.Stream;
 
@@ -162,5 +168,15 @@ private String internalExecute(String sql) throws Exception {
             System.out.println(explainString);
             return explainString;
         }
+
+        public Planner internalExecuteOneAndGetPlan() throws Exception{
+            SqlScanner input = new SqlScanner(new StringReader(sql), ctx.getSessionVariable().getSqlMode());
+            SqlParser parser = new SqlParser(input);
+            List<StatementBase> stmts =  SqlParserUtils.getMultiStmts(parser);
+            StmtExecutor stmtExecutor = new StmtExecutor(connectContext, stmts.get(0));
+            stmtExecutor.execute();
+
+            return stmtExecutor.planner();
+        }
     }
 }