[Enhancement] shorten FE plan time in some corner case (backport #49137…

…) (#49222) Co-authored-by: before-Sunrise <[email protected]>
StarRocks · Aug 15, 2024 · 423a577 · 423a577
1 parent 58ef227
commit 423a577
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 25 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/Utils.java
@@ -149,38 +149,25 @@ private static void extractDisjunctiveImpl(ScalarOperator root, List<ScalarOpera
     }
 
     public static List<ColumnRefOperator> extractColumnRef(ScalarOperator root) {
-        if (null == root || !root.isVariable()) {
+        if (null == root) {
             return new LinkedList<>();
         }
 
-        LinkedList<ColumnRefOperator> list = new LinkedList<>();
-        if (OperatorType.VARIABLE.equals(root.getOpType())) {
-            list.add((ColumnRefOperator) root);
-            return list;
-        }
-
-        for (ScalarOperator child : root.getChildren()) {
-            list.addAll(extractColumnRef(child));
-        }
-
-        return list;
+        return root.getColumnRefs();
     }
 
     public static int countColumnRef(ScalarOperator root) {
-        return countColumnRef(root, 0);
-    }
-
-    private static int countColumnRef(ScalarOperator root, int count) {
-        if (null == root || !root.isVariable()) {
+        if (null == root) {
             return 0;
         }
 
         if (OperatorType.VARIABLE.equals(root.getOpType())) {
             return 1;
         }
 
+        int count = 0;
         for (ScalarOperator child : root.getChildren()) {
-            count += countColumnRef(child, count);
+            count += countColumnRef(child);
         }
 
         return count;

diff --git a/.../src/main/java/com/starrocks/sql/optimizer/operator/scalar/CompoundPredicateOperator.java b/.../src/main/java/com/starrocks/sql/optimizer/operator/scalar/CompoundPredicateOperator.java
@@ -38,8 +38,8 @@ public class CompoundPredicateOperator extends PredicateOperator {
     //      AND         AND
     //     /   \       /  \
     // subT1   a+1  And  subT5
-    //               /   \
-    //             subT3  a+1
+    //             /   \
+    //          subT3  a+1
     private int compoundTreeLeafNodeNumber;
     private Set<ScalarOperator> compoundTreeUniqueLeaves;
 

diff --git a/.../src/main/java/com/starrocks/sql/optimizer/rewrite/scalar/ExtractCommonPredicateRule.java b/.../src/main/java/com/starrocks/sql/optimizer/rewrite/scalar/ExtractCommonPredicateRule.java
@@ -50,7 +50,7 @@ public ScalarOperator visitCompoundPredicate(CompoundPredicateOperator predicate
         List<List<ScalarOperator>> orAndPredicates = Lists.newArrayList();
 
         for (ScalarOperator or : orLists) {
-            orAndPredicates.add(Lists.newArrayList(Utils.extractConjuncts(or)));
+            orAndPredicates.add(Utils.extractConjuncts(or));
         }
 
         // extract common predicate

diff --git a/...core/src/main/java/com/starrocks/sql/optimizer/rewrite/scalar/NormalizePredicateRule.java b/...core/src/main/java/com/starrocks/sql/optimizer/rewrite/scalar/NormalizePredicateRule.java
@@ -145,9 +145,9 @@ public ScalarOperator visitCompoundPredicate(CompoundPredicateOperator predicate
     @Nullable
     private Optional<ScalarOperator> getOptimizedCompoundTree(CompoundPredicateOperator parent) {
         // reset node first So we can apply NormalizePredicateRule to one tree many times
-        parent.setCompoundTreeUniqueLeaves(Sets.newLinkedHashSet());
+        parent.setCompoundTreeUniqueLeaves(null);
         parent.setCompoundTreeLeafNodeNumber(0);
-        Set<ScalarOperator> compoundTreeUniqueLeaves = parent.getCompoundTreeUniqueLeaves();
+        Set<ScalarOperator> compoundTreeUniqueLeaves = null;
 
         for (ScalarOperator child : parent.getChildren()) {
             if (child != null) {
@@ -157,12 +157,19 @@ private Optional<ScalarOperator> getOptimizedCompoundTree(CompoundPredicateOpera
                         (parent.isOr() && OperatorType.COMPOUND.equals(child.getOpType()) &&
                                 ((CompoundPredicateOperator) child).isOr())) {
                     CompoundPredicateOperator compoundChild = (CompoundPredicateOperator) (child);
-                    compoundTreeUniqueLeaves.addAll(compoundChild.getCompoundTreeUniqueLeaves());
+                    if (compoundTreeUniqueLeaves == null) {
+                        compoundTreeUniqueLeaves = compoundChild.getCompoundTreeUniqueLeaves();
+                    } else {
+                        compoundTreeUniqueLeaves.addAll(compoundChild.getCompoundTreeUniqueLeaves());
+                    }
                     parent.setCompoundTreeLeafNodeNumber(
                             compoundChild.getCompoundTreeLeafNodeNumber() + parent.getCompoundTreeLeafNodeNumber());
                 } else {
                     // child is leaf node in compound tree
                     // we cache CompoundPredicate's hash value to eliminate duplicate calculations
+                    if (compoundTreeUniqueLeaves == null) {
+                        compoundTreeUniqueLeaves = Sets.newLinkedHashSet();
+                    }
                     compoundTreeUniqueLeaves.add(new HashCachedScalarOperator(child));
                     parent.setCompoundTreeLeafNodeNumber(1 + parent.getCompoundTreeLeafNodeNumber());
                 }
@@ -176,6 +183,7 @@ private Optional<ScalarOperator> getOptimizedCompoundTree(CompoundPredicateOpera
                 }
             }
         }
+        parent.setCompoundTreeUniqueLeaves(compoundTreeUniqueLeaves);
 
         // this tree can be optimized
         if (compoundTreeUniqueLeaves.size() != parent.getCompoundTreeLeafNodeNumber()) {
@@ -247,7 +255,12 @@ public ScalarOperator visitInPredicate(InPredicateOperator predicate, ScalarOper
             result.add(newOp);
         });
 
-        return isIn ? Utils.compoundOr(result) : Utils.compoundAnd(result);
+        ScalarOperator res = isIn ? Utils.compoundOr(result) : Utils.compoundAnd(result);
+        if (res instanceof CompoundPredicateOperator) {
+            return visitCompoundPredicate((CompoundPredicateOperator) res, context);
+        } else {
+            return res;
+        }
     }
 
     // rewrite collection element to subfiled

diff --git a/.../src/test/java/com/starrocks/sql/optimizer/rewrite/scalar/NormalizePredicateRuleTest.java b/.../src/test/java/com/starrocks/sql/optimizer/rewrite/scalar/NormalizePredicateRuleTest.java
@@ -118,4 +118,28 @@ public void testCompound() {
 
         assertTrue(result.isConstantTrue());
     }
+
+    @Test
+    public void testCompound1() {
+        NormalizePredicateRule rule = new NormalizePredicateRule();
+        ScalarOperatorRewriteContext context = new ScalarOperatorRewriteContext();
+
+        InPredicateOperator inOp = new InPredicateOperator(
+                true,
+                ConstantOperator.createInt(1),
+                new ColumnRefOperator(0, Type.INT, "col1", true),
+                new ColumnRefOperator(0, Type.INT, "col1", true)
+        );
+
+        CompoundPredicateOperator compoundPredicateOperator =
+                new CompoundPredicateOperator(CompoundPredicateOperator.CompoundType.AND, inOp,
+                        new BinaryPredicateOperator(BinaryType.GE,
+                                ConstantOperator.createInt(1),
+                                new ColumnRefOperator(1, Type.INT, "test1", true))
+                );
+
+        ScalarOperatorRewriter operatorRewriter = new ScalarOperatorRewriter();
+        ScalarOperator res =
+                operatorRewriter.rewrite(compoundPredicateOperator, Lists.newArrayList(new NormalizePredicateRule()));
+    }
 }