[BugFix] Operator's statistics not be estimated during JoinReorder of…

… table-pruning phase (#51304) Signed-off-by: satanson <[email protected]>
StarRocks · Sep 24, 2024 · 61ebfc7 · 61ebfc7
1 parent 9222444
commit 61ebfc7
Show file tree

Hide file tree

Showing 3 changed files with 4,234 additions and 0 deletions.
diff --git a/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/join/ReorderJoinRule.java b/fe/fe-core/src/main/java/com/starrocks/sql/optimizer/rule/join/ReorderJoinRule.java
@@ -14,6 +14,7 @@
 
 package com.starrocks.sql.optimizer.rule.join;
 
+import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -336,6 +337,14 @@ private OptExpression deriveNewOptExpression(OptExpression optExpression,
             LogicalProperty newProperty = new LogicalProperty(optExpression.getLogicalProperty());
             newProperty.setOutputColumns(newCols);
 
+            if (!Optional.ofNullable(optExpression.getStatistics()).isPresent()) {
+                ExpressionContext expressionContext = new ExpressionContext(optExpression);
+                StatisticsCalculator statisticsCalculator = new StatisticsCalculator(
+                        expressionContext, optimizerContext.getColumnRefFactory(), optimizerContext);
+                statisticsCalculator.estimatorStats();
+                optExpression.setStatistics(expressionContext.getStatistics());
+            }
+            Preconditions.checkState(optExpression.getStatistics() != null);
             Statistics newStats = Statistics.buildFrom(optExpression.getStatistics()).build();
             Iterator<Map.Entry<ColumnRefOperator, ColumnStatistic>>
                     iterator = newStats.getColumnStatistics().entrySet().iterator();

diff --git a/fe/fe-core/src/test/java/com/starrocks/sql/plan/ReplayFromDumpTest.java b/fe/fe-core/src/test/java/com/starrocks/sql/plan/ReplayFromDumpTest.java
@@ -33,6 +33,8 @@
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.stream.Stream;
+
 public class ReplayFromDumpTest extends ReplayFromDumpTestBase {
     @Test
     public void testForceRuleBasedRewrite() throws Exception {
@@ -974,4 +976,21 @@ public void testQueryCacheMisuseExogenousRuntimeFilter() throws Exception {
             connectContext.getSessionVariable().replayFromJson(savedSv);
         }
     }
+
+    @Test
+    public void testPruneTableNPE() throws Exception {
+        String savedSv = connectContext.getSessionVariable().getJsonString();
+        try {
+            connectContext.getSessionVariable().setEnableCboTablePrune(true);
+            connectContext.getSessionVariable().setEnableRboTablePrune(true);
+            Pair<QueryDumpInfo, String> replayPair =
+                    getPlanFragment(getDumpInfoFromFile("query_dump/prune_table_npe"),
+                            null, TExplainLevel.NORMAL);
+            long numHashJoins = Stream.of(replayPair.second.split("\n"))
+                    .filter(ln -> ln.contains("HASH JOIN")).count();
+            Assert.assertEquals(numHashJoins, 2);
+        } finally {
+            connectContext.getSessionVariable().replayFromJson(savedSv);
+        }
+    }
 }