From 0e56290e4ecb62f877432700f80d9f4995ced316 Mon Sep 17 00:00:00 2001
From: minghong <englefly@gmail.com>
Date: Thu, 26 Sep 2024 15:56:38 +0800
Subject: [PATCH] fe

---
 .../doris/nereids/stats/StatsCalculator.java  | 43 ++++++++++++-------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/StatsCalculator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/StatsCalculator.java
index bfff21ab3a500b..1695d03359346b 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/StatsCalculator.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/stats/StatsCalculator.java
@@ -397,9 +397,28 @@ private void checkIfUnknownStatsUsedAsKey(StatisticsBuilder builder) {
         }
     }
 
+    private double getOlapTableRowCount(OlapScan olapScan) {
+        OlapTable olapTable = olapScan.getTable();
+        AnalysisManager analysisManager = Env.getCurrentEnv().getAnalysisManager();
+        TableStatsMeta tableMeta = analysisManager.findTableStatsStatus(olapScan.getTable().getId());
+        double rowCount = -1;
+        if (tableMeta != null && tableMeta.userInjected) {
+            rowCount = tableMeta.getRowCount(olapScan.getSelectedIndexId());
+        } else {
+            rowCount = olapTable.getRowCountForIndex(olapScan.getSelectedIndexId(), true);
+            if (rowCount == -1) {
+                if (tableMeta != null) {
+                    rowCount = tableMeta.getRowCount(olapScan.getSelectedIndexId());
+                }
+            }
+        }
+        return rowCount;
+    }
+
     private Statistics computeOlapScan(OlapScan olapScan) {
         OlapTable olapTable = olapScan.getTable();
-        double tableRowCount = olapTable.getRowCountForIndex(olapScan.getSelectedIndexId());
+        double tableRowCount = getOlapTableRowCount(olapScan);
+        tableRowCount = Math.max(1, tableRowCount);
         if (tableRowCount <= 0) {
             AnalysisManager analysisManager = Env.getCurrentEnv().getAnalysisManager();
             TableStatsMeta tableMeta = analysisManager.findTableStatsStatus(olapScan.getTable().getId());
@@ -471,11 +490,13 @@ private Statistics computeOlapScan(OlapScan olapScan) {
                 builder.putColumnStatistics(slot, ColumnStatistic.UNKNOWN);
             }
         }
-
+        boolean useTableLevelStats = true;
         if (olapScan.getSelectedPartitionIds().size() < olapScan.getTable().getPartitionNum()) {
             // partition pruned
+            // try to use selected partition stats, if failed, fall back to table stats
             double selectedPartitionsRowCount = getSelectedPartitionRowCount(olapScan);
-            if (selectedPartitionsRowCount > 0) {
+            if (selectedPartitionsRowCount >= 0) {
+                useTableLevelStats = false;
                 List<String> selectedPartitionNames = new ArrayList<>(olapScan.getSelectedPartitionIds().size());
                 olapScan.getSelectedPartitionIds().forEach(id -> {
                     selectedPartitionNames.add(olapScan.getTable().getPartition(id).getName());
@@ -489,19 +510,11 @@ private Statistics computeOlapScan(OlapScan olapScan) {
                 }
                 checkIfUnknownStatsUsedAsKey(builder);
                 builder.setRowCount(selectedPartitionsRowCount);
-            } else {
-                // if partition row count is invalid (-1), fallback to table stats
-                for (SlotReference slot : visibleOutputSlots) {
-                    ColumnStatistic cache = getColumnStatsFromTableCache((CatalogRelation) olapScan, slot);
-                    ColumnStatisticBuilder colStatsBuilder = new ColumnStatisticBuilder(cache);
-                    colStatsBuilder.setCount(tableRowCount);
-                    adjustColStats(olapScan, slot, colStatsBuilder);
-                    builder.putColumnStatistics(slot, colStatsBuilder.build());
-                }
-                checkIfUnknownStatsUsedAsKey(builder);
-                builder.setRowCount(tableRowCount);
             }
-        } else {
+        }
+        // 1. no partition is pruned, or
+        // 2. fall back to table stats
+        if (useTableLevelStats) {
             // get table level stats
             for (SlotReference slot : visibleOutputSlots) {
                 ColumnStatistic cache = getColumnStatsFromTableCache((CatalogRelation) olapScan, slot);