[feat](Nereids) support set var in hint when parse sql (#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
apache · Oct 23, 2024 · 98a9a45 · 98a9a45
1 parent a21bb80
commit 98a9a45
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 36 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,29 @@ public Map<String, Optional<String>> getParameters() {
         return parameters;
     }
 
+    /**
+     * set session variable in sql level
+     * @param context statement context
+     */
+    public void setVarOnceInSql(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 (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())));
+                    context.invalidCache(key);
+                } catch (Throwable t) {
+                    throw new AnalysisException("Can not set session variable '"
+                        + key + "' = '" + value.get() + "'", t);
+                }
+            }
+        }
+    }
+
     @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,8 +32,6 @@
 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;
@@ -58,7 +53,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 +76,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
@@ -27,9 +27,11 @@
 import org.apache.doris.common.FeConstants;
 import org.apache.doris.common.PatternMatcher;
 import org.apache.doris.common.PatternMatcherWrapper;
+import org.apache.doris.common.VariableAnnotation;
 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 +231,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);
+    }
 }