[feat](Nereids) support set var in hint when parse sql (#41331) (#42335)

pick: #41331 set var hint need to be enable to use before analyze, so it need to be set when parsing sql now it would set twice when parse and begin of analyze ## Proposed changes Issue Number: close #xxx
apache · Oct 25, 2024 · c8b1f75 · c8b1f75
1 parent be56e40
commit c8b1f75
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 38 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java
@@ -1834,7 +1834,9 @@ private LogicalPlan withSelectHint(LogicalPlan logicalPlan, List<ParserRuleConte
                                 parameters.put(parameterName, value);
                             }
                         }
-                        hints.put(hintName, new SelectHintSetVar(hintName, parameters));
+                        SelectHintSetVar setVar = new SelectHintSetVar(hintName, parameters);
+                        setVar.setVarOnceInSql(ConnectContext.get().getStatementContext());
+                        hints.put(hintName, setVar);
                         break;
                     case "leading":
                         List<String> leadingParameters = new ArrayList<String>();

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/SelectHintSetVar.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/SelectHintSetVar.java
@@ -17,6 +17,13 @@
 
 package org.apache.doris.nereids.properties;
 
+import org.apache.doris.analysis.SetVar;
+import org.apache.doris.analysis.StringLiteral;
+import org.apache.doris.nereids.StatementContext;
+import org.apache.doris.nereids.exceptions.AnalysisException;
+import org.apache.doris.qe.SessionVariable;
+import org.apache.doris.qe.VariableMgr;
+
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -38,6 +45,42 @@ public Map<String, Optional<String>> getParameters() {
         return parameters;
     }
 
+    /**
+     * set session variable in sql level
+     * @param context statement context
+     */
+    public void setVarOnceInSql(StatementContext context) {
+        if (context == null) {
+            return;
+        }
+        SessionVariable sessionVariable = context.getConnectContext().getSessionVariable();
+        // set temporary session value, and then revert value in the 'finally block' of StmtExecutor#execute
+        sessionVariable.setIsSingleSetVar(true);
+        for (Map.Entry<String, Optional<String>> kv : getParameters().entrySet()) {
+            String key = kv.getKey();
+            Optional<String> value = kv.getValue();
+            if (value.isPresent()) {
+                try {
+                    VariableMgr.setVar(sessionVariable, new SetVar(key, new StringLiteral(value.get())));
+                } catch (Throwable t) {
+                    throw new AnalysisException("Can not set session variable '"
+                        + key + "' = '" + value.get() + "'", t);
+                }
+            }
+        }
+        // if sv set enable_nereids_planner=true and hint set enable_nereids_planner=false, we should set
+        // enable_fallback_to_original_planner=true and revert it after executing.
+        // throw exception to fall back to original planner
+        if (!sessionVariable.isEnableNereidsPlanner()) {
+            try {
+                sessionVariable.enableFallbackToOriginalPlannerOnce();
+            } catch (Throwable t) {
+                throw new AnalysisException("failed to set fallback to original planner to true", t);
+            }
+            throw new AnalysisException("The nereids is disabled in this sql, fallback to original planner");
+        }
+    }
+
     @Override
     public String toString() {
         String kvString = parameters

diff --git a/...ore/src/main/java/org/apache/doris/nereids/rules/analysis/EliminateLogicalSelectHint.java b/...ore/src/main/java/org/apache/doris/nereids/rules/analysis/EliminateLogicalSelectHint.java
@@ -17,12 +17,9 @@
 
 package org.apache.doris.nereids.rules.analysis;
 
-import org.apache.doris.analysis.SetVar;
-import org.apache.doris.analysis.StringLiteral;
 import org.apache.doris.common.DdlException;
 import org.apache.doris.nereids.CascadesContext;
 import org.apache.doris.nereids.StatementContext;
-import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.hint.Hint;
 import org.apache.doris.nereids.hint.LeadingHint;
 import org.apache.doris.nereids.hint.OrderedHint;
@@ -35,15 +32,12 @@
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalSelectHint;
 import org.apache.doris.qe.ConnectContext;
-import org.apache.doris.qe.SessionVariable;
-import org.apache.doris.qe.VariableMgr;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.Map;
 import java.util.Map.Entry;
-import java.util.Optional;
 
 /**
  * eliminate logical select hint and set them to cascade context
@@ -58,7 +52,7 @@ public Rule build() {
             for (Entry<String, SelectHint> hint : selectHintPlan.getHints().entrySet()) {
                 String hintName = hint.getKey();
                 if (hintName.equalsIgnoreCase("SET_VAR")) {
-                    setVar((SelectHintSetVar) hint.getValue(), ctx.statementContext);
+                    ((SelectHintSetVar) hint.getValue()).setVarOnceInSql(ctx.statementContext);
                 } else if (hintName.equalsIgnoreCase("ORDERED")) {
                     try {
                         ctx.cascadesContext.getConnectContext().getSessionVariable()
@@ -81,35 +75,6 @@ public Rule build() {
         }).toRule(RuleType.ELIMINATE_LOGICAL_SELECT_HINT);
     }
 
-    private void setVar(SelectHintSetVar selectHint, StatementContext context) {
-        SessionVariable sessionVariable = context.getConnectContext().getSessionVariable();
-        // set temporary session value, and then revert value in the 'finally block' of StmtExecutor#execute
-        sessionVariable.setIsSingleSetVar(true);
-        for (Entry<String, Optional<String>> kv : selectHint.getParameters().entrySet()) {
-            String key = kv.getKey();
-            Optional<String> value = kv.getValue();
-            if (value.isPresent()) {
-                try {
-                    VariableMgr.setVar(sessionVariable, new SetVar(key, new StringLiteral(value.get())));
-                } catch (Throwable t) {
-                    throw new AnalysisException("Can not set session variable '"
-                        + key + "' = '" + value.get() + "'", t);
-                }
-            }
-        }
-        // if sv set enable_nereids_planner=true and hint set enable_nereids_planner=false, we should set
-        // enable_fallback_to_original_planner=true and revert it after executing.
-        // throw exception to fall back to original planner
-        if (!sessionVariable.isEnableNereidsPlanner()) {
-            try {
-                sessionVariable.enableFallbackToOriginalPlannerOnce();
-            } catch (Throwable t) {
-                throw new AnalysisException("failed to set fallback to original planner to true", t);
-            }
-            throw new AnalysisException("The nereids is disabled in this sql, fallback to original planner");
-        }
-    }
-
     private void extractLeading(SelectHintLeading selectHint, CascadesContext context,
                                     StatementContext statementContext, Map<String, SelectHint> hints) {
         LeadingHint hint = new LeadingHint("Leading", selectHint.getParameters(), selectHint.toString());

diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SessionVariablesTest.java
@@ -30,6 +30,7 @@
 import org.apache.doris.common.util.ProfileManager;
 import org.apache.doris.common.util.RuntimeProfile;
 import org.apache.doris.load.ExportJob;
+import org.apache.doris.nereids.parser.NereidsParser;
 import org.apache.doris.task.ExportExportingTask;
 import org.apache.doris.thrift.TQueryOptions;
 import org.apache.doris.utframe.TestWithFeService;
@@ -229,4 +230,11 @@ public void testDisableProfile() {
         }
 
     }
+
+    @Test
+    public void testSetVarInHint() {
+        String sql = "insert into test_t1 select /*+ set_var(enable_nereids_dml_with_pipeline=false)*/ * from test_t1 where enable_nereids_dml_with_pipeline=true";
+        new NereidsParser().parseSQL(sql);
+        Assertions.assertEquals(false, connectContext.getSessionVariable().enableNereidsDmlWithPipeline);
+    }
 }
diff --git a/regression-test/suites/nereids_syntax_p0/system_var.groovy b/regression-test/suites/nereids_syntax_p0/system_var.groovy
@@ -38,7 +38,7 @@ suite("nereids_sys_var") {
     // set an invalid parameter, and throw an exception
     test {
         sql "select /*+SET_VAR(runtime_filter_type=10000)*/ * from supplier limit 10"
-        exception "Unexpected exception: Can not set session variable"
+        exception "errCode"
     }
 
     sql "select @@session.time_zone"