From 8342eb0b02c18f0d5703af1f01eef7dcd6495e16 Mon Sep 17 00:00:00 2001
From: emmymiao87 <522274284@qq.com>
Date: Wed, 18 Dec 2019 14:50:32 +0800
Subject: [PATCH] Only UDA function could be supported in correlated subquery

Those query of issue could not be supported. #2483 #2493
Those query is forbidden:
query1: select * from t1 where k1=(select k1 from t2 where t1.k2=t2.k2);
query2: select * from t1 where k1=(select distinct k1 from t2 where t1.k2=t2.k2);
Only sum, max, min, avg and count function could appear on select clause for correlated subquery. #2420
Those query is legal:
query1: select * from t1 where k1=(select avg(k1) from t2 where t1.k2=t2.k2);
---
 .../java/org/apache/doris/analysis/Expr.java   | 18 ++++++++++++++++++
 .../apache/doris/analysis/StmtRewriter.java    | 15 +++++++++++----
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/fe/src/main/java/org/apache/doris/analysis/Expr.java b/fe/src/main/java/org/apache/doris/analysis/Expr.java
index 0866fe555a..8be5f09bbb 100644
--- a/fe/src/main/java/org/apache/doris/analysis/Expr.java
+++ b/fe/src/main/java/org/apache/doris/analysis/Expr.java
@@ -123,6 +123,24 @@ abstract public class Expr extends TreeNode<Expr> implements ParseNode, Cloneabl
                 }
             };
     */
+    // Returns true if an Expr is a builtin aggregate function.
+    public final static com.google.common.base.Predicate<Expr> IS_BUILTIN_AGG_FN =
+            new com.google.common.base.Predicate<Expr>() {
+                @Override
+                public boolean apply(Expr arg) {
+                    if (arg instanceof FunctionCallExpr) {
+                        String fnName = ((FunctionCallExpr)arg).getFnName().getFunction();
+                        return  (fnName.equalsIgnoreCase("sum")
+                                || fnName.equalsIgnoreCase("max")
+                                || fnName.equalsIgnoreCase("min")
+                                || fnName.equalsIgnoreCase("avg")
+                                || fnName.equalsIgnoreCase("count"));
+                    } else {
+                        return false;
+                    }
+                }
+            };
+
 
     public final static com.google.common.base.Predicate<Expr> IS_TRUE_LITERAL =
             new com.google.common.base.Predicate<Expr>() {
diff --git a/fe/src/main/java/org/apache/doris/analysis/StmtRewriter.java b/fe/src/main/java/org/apache/doris/analysis/StmtRewriter.java
index 34101df358..aa995269ae 100644
--- a/fe/src/main/java/org/apache/doris/analysis/StmtRewriter.java
+++ b/fe/src/main/java/org/apache/doris/analysis/StmtRewriter.java
@@ -744,10 +744,17 @@ public class StmtRewriter {
         Preconditions.checkNotNull(stmt);
         // Grouping and/or aggregation (including analytic functions) is only
         // allowed on correlated EXISTS subqueries
-        if ((expr instanceof BinaryPredicate
-                && (stmt.hasAggInfo() || stmt.hasAnalyticInfo()))
-                || (expr instanceof InPredicate
-                && (stmt.hasAggInfo() || stmt.hasAnalyticInfo()))) {
+        if (expr instanceof BinaryPredicate) {
+            if (stmt.getSelectList().getItems().size() != 1) {
+                throw new AnalysisException("The subquery only support one item in select clause");
+            }
+            SelectListItem item = stmt.getSelectList().getItems().get(0);
+            if (!item.getExpr().contains(Expr.IS_BUILTIN_AGG_FN)) {
+                throw new AnalysisException("The select item of correlated subquery should only be sum, min, max, avg"
+                                                    + " and count. Current subquery:" + stmt.toSql());
+            }
+        }
+        if (expr instanceof InPredicate && (stmt.hasAggInfo() || stmt.hasAnalyticInfo())) {
             LOG.warn("canRewriteCorrelatedSubquery fail, expr={} subquery={}", expr.toSql(), stmt.toSql());
             throw new AnalysisException("Unsupported correlated subquery with grouping "
                     + "and/or aggregation: "