From 0ff191cdf33a2b0f37dea36a3134164a76dee922 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=B0=A2=E5=81=A5?= <jianxie0@gmail.com>
Date: Mon, 28 Aug 2023 19:03:43 +0800
Subject: [PATCH] [feature](Nereids) add expr depth limit and expr children
 limit in Nereids (#23569)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

#### `expr_depth_limit`

Default：3000

IsMutable：true

Limit on the depth of an expr tree.  Exceed this limit may cause long analysis time while holding db read lock.  Do not set this if you know what you are doing

#### `expr_children_limit`

Default：10000

IsMutable：true

Limit on the number of expr children of an expr tree.  Exceed this limit may cause long analysis time while holding database read lock.
---
 .../nereids/trees/expressions/Expression.java | 38 +++++++++++++++++++
 .../nereids/parser/NereidsParserTest.java     | 11 ++++++
 2 files changed, 49 insertions(+)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java
index dcb228bbf8..86122af5ad 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Expression.java
@@ -17,6 +17,7 @@
 
 package org.apache.doris.nereids.trees.expressions;
 
+import org.apache.doris.common.Config;
 import org.apache.doris.nereids.analyzer.Unbound;
 import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.trees.AbstractTreeNode;
@@ -38,6 +39,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 import org.apache.commons.lang3.StringUtils;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -51,12 +53,40 @@ public abstract class Expression extends AbstractTreeNode<Expression> implements
     // Mask this expression is generated by rule, should be removed.
     public boolean isGeneratedIsNotNull = false;
 
+    private final int depth;
+    private final int width;
+
     protected Expression(Expression... children) {
         super(children);
+        depth = Arrays.stream(children)
+                .mapToInt(e -> e.depth)
+                .max().orElse(0) + 1;
+        width = Arrays.stream(children)
+                .mapToInt(e -> e.width)
+                .sum() + (children.length == 0 ? 1 : 0);
+        checkLimit();
     }
 
     protected Expression(List<Expression> children) {
         super(Optional.empty(), children);
+        depth = children.stream()
+                .mapToInt(e -> e.depth)
+                .max().orElse(0) + 1;
+        width = children.stream()
+                .mapToInt(e -> e.width)
+                .sum() + (children.isEmpty() ? 1 : 0);
+        checkLimit();
+    }
+
+    private void checkLimit() {
+        if (depth > Config.expr_depth_limit) {
+            throw new AnalysisException(String.format("Exceeded the maximum depth of an "
+                    + "expression tree (%s).", Config.expr_depth_limit));
+        }
+        if (width > Config.expr_children_limit) {
+            throw new AnalysisException(String.format("Exceeded the maximum children of an "
+                    + "expression tree (%s).", Config.expr_children_limit));
+        }
     }
 
     public Alias alias(String alias) {
@@ -153,6 +183,14 @@ public abstract class Expression extends AbstractTreeNode<Expression> implements
         return children.get(index);
     }
 
+    public int getWidth() {
+        return width;
+    }
+
+    public int getDepth() {
+        return depth;
+    }
+
     @Override
     public Expression withChildren(List<Expression> children) {
         throw new RuntimeException();
diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java
index c7ecd44e92..b5777ecd85 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/parser/NereidsParserTest.java
@@ -346,4 +346,15 @@ public class NereidsParserTest extends ParserTestBase {
             Assertions.assertEquals(6, decimalV2Type.getScale());
         }
     }
+
+    @Test
+    void testParseExprDepthWidth() {
+        String sql = "SELECT 1+2 = 3 from t";
+        NereidsParser nereidsParser = new NereidsParser();
+        LogicalPlan logicalPlan = (LogicalPlan) nereidsParser.parseSingle(sql).child(0);
+        System.out.println(logicalPlan);
+        // alias (1 + 2 = 3)
+        Assertions.assertEquals(4, logicalPlan.getExpressions().get(0).getDepth());
+        Assertions.assertEquals(3, logicalPlan.getExpressions().get(0).getWidth());
+    }
 }