From 95322e2ebebbc58803647b65e1d5e88e02a5986e Mon Sep 17 00:00:00 2001
From: morrySnow <101034200+morrySnow@users.noreply.github.com>
Date: Wed, 20 Mar 2024 18:21:52 +0800
Subject: [PATCH] [opt](variable) user variable support expression rather than
 literal (#32492)

---
 .../doris/analysis/SetUserDefinedVar.java     | 36 ++++++++++++++++++-
 .../data/nereids_p0/test_user_var.out         | 13 ++++---
 .../suites/nereids_p0/test_user_var.groovy    | 12 ++++---
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/SetUserDefinedVar.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/SetUserDefinedVar.java
index 7b0c5d98aa..0b97c49e59 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/analysis/SetUserDefinedVar.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/SetUserDefinedVar.java
@@ -19,6 +19,12 @@ package org.apache.doris.analysis;
 
 import org.apache.doris.catalog.ScalarType;
 import org.apache.doris.common.AnalysisException;
+import org.apache.doris.common.Pair;
+import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.rewrite.ExprRewriter;
+
+import java.util.HashMap;
+import java.util.Map;
 
 public class SetUserDefinedVar extends SetVar {
     public SetUserDefinedVar(String variable, Expr value) {
@@ -28,12 +34,40 @@ public class SetUserDefinedVar extends SetVar {
     @Override
     public void analyze(Analyzer analyzer) throws AnalysisException  {
         Expr expression = getValue();
+        if (!(expression instanceof LiteralExpr)) {
+            boolean changed = true;
+            for (int i = 0; i < 1000 && changed; i++) {
+                Pair<Expr, Boolean> result = fold(analyzer, expression);
+                expression = result.first;
+                changed = result.second;
+            }
+        }
         if (expression instanceof NullLiteral) {
             setResult(NullLiteral.create(ScalarType.NULL));
         } else if (expression instanceof LiteralExpr) {
             setResult((LiteralExpr) expression);
         } else {
-            throw new AnalysisException("Unsupported to set the non-literal for user defined variables.");
+            throw new AnalysisException("Unsupported to set " + expression.toSql() + " user defined variables.");
         }
     }
+
+    public void registerExprId(Expr expr, Analyzer analyzer) {
+        if (expr.getId() == null) {
+            analyzer.registerExprId(expr);
+        }
+        for (Expr child : expr.getChildren()) {
+            registerExprId(child, analyzer);
+        }
+    }
+
+    private Pair<Expr, Boolean> fold(Analyzer analyzer, Expr expression) throws AnalysisException {
+        expression.analyze(analyzer);
+        registerExprId(expression, analyzer);
+        ExprRewriter rewriter = analyzer.getExprRewriter();
+        rewriter.reset();
+        Map<String, Expr> exprMap = new HashMap<>();
+        exprMap.put(expression.getId().toString(), expression);
+        rewriter.rewriteConstant(exprMap, analyzer, ConnectContext.get().getSessionVariable().toThrift());
+        return Pair.of(exprMap.get(expression.getId().toString()), rewriter.changed());
+    }
 }
diff --git a/regression-test/data/nereids_p0/test_user_var.out b/regression-test/data/nereids_p0/test_user_var.out
index 41aebfe022..5d6338dd7e 100644
--- a/regression-test/data/nereids_p0/test_user_var.out
+++ b/regression-test/data/nereids_p0/test_user_var.out
@@ -1,16 +1,19 @@
 -- This file is automatically generated. You should know what you did if you want to edit this
--- !select1 --
+-- !integer --
 1	0	-1
 
--- !select2 --
+-- !decimal --
 1.1	0.0	-1.1
 
--- !select3 --
+-- !string --
 H	
 
--- !select4 --
+-- !boolean --
 true	false
 
--- !select5 --
+-- !null_literal --
 \N	\N
 
+-- !function --
+4
+
diff --git a/regression-test/suites/nereids_p0/test_user_var.groovy b/regression-test/suites/nereids_p0/test_user_var.groovy
index f54e139632..d5eb164e3f 100644
--- a/regression-test/suites/nereids_p0/test_user_var.groovy
+++ b/regression-test/suites/nereids_p0/test_user_var.groovy
@@ -23,10 +23,12 @@ suite("test_user_var") {
     sql "SET @c1='H', @c2=''"
     sql "SET @d1=true, @d2=false"
     sql "SET @f1=null"
+    sql "set @func_1=(abs(1) + 1) * 2"
 
-    qt_select1 'select @a1, @a2, @a3;'
-    qt_select2 'select @b1, @b2, @b3;'
-    qt_select3 'select @c1, @c2;'
-    qt_select4 'select @d1, @d2;'
-    qt_select5 'select @f1, @f2;'
+    qt_integer 'select @a1, @a2, @a3;'
+    qt_decimal 'select @b1, @b2, @b3;'
+    qt_string 'select @c1, @c2;'
+    qt_boolean 'select @d1, @d2;'
+    qt_null_literal 'select @f1, @f2;'
+    qt_function 'select @func_1'
 }
\ No newline at end of file