[opt](nereids)new way to set pre-agg status (#34738)

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -36,6 +36,7 @@
 import org.apache.doris.nereids.rules.rewrite.AddDefaultLimit;
 import org.apache.doris.nereids.rules.rewrite.AdjustConjunctsReturnType;
 import org.apache.doris.nereids.rules.rewrite.AdjustNullable;
+import org.apache.doris.nereids.rules.rewrite.AdjustPreAggStatus;
 import org.apache.doris.nereids.rules.rewrite.AggScalarSubQueryToWindowFunction;
 import org.apache.doris.nereids.rules.rewrite.BuildAggForUnion;
 import org.apache.doris.nereids.rules.rewrite.CTEInline;
@@ -391,6 +392,9 @@ public class Rewriter extends AbstractBatchJobExecutor {
                     bottomUp(RuleSet.PUSH_DOWN_FILTERS),
                     custom(RuleType.ELIMINATE_UNNECESSARY_PROJECT, EliminateUnnecessaryProject::new)
             ),
+            topic("adjust preagg status",
+                    topDown(new AdjustPreAggStatus())
+            ),
             topic("topn optimize",
                     topDown(new DeferMaterializeTopNResult())
             ),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java

@@ -242,6 +242,21 @@ public enum RuleType {
     MATERIALIZED_INDEX_PROJECT_SCAN(RuleTypeClass.REWRITE),
     MATERIALIZED_INDEX_PROJECT_FILTER_SCAN(RuleTypeClass.REWRITE),
     MATERIALIZED_INDEX_FILTER_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_PROJECT_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_FILTER_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_REPEAT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_REPEAT_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_REPEAT_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_REPEAT_PROJECT_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_AGG_REPEAT_FILTER_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_PROJECT_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_PROJECT_FILTER_SCAN(RuleTypeClass.REWRITE),
+    PREAGG_STATUS_FILTER_PROJECT_SCAN(RuleTypeClass.REWRITE),
     REDUCE_AGGREGATE_CHILD_OUTPUT_ROWS(RuleTypeClass.REWRITE),
 
     OLAP_SCAN_PARTITION_PRUNE(RuleTypeClass.REWRITE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindRelation.java

@@ -206,7 +206,7 @@ private LogicalPlan makeOlapScan(TableIf table, UnboundRelation unboundRelation,
                 }
                 PreAggStatus preAggStatus
                         = olapTable.getIndexMetaByIndexId(indexId).getKeysType().equals(KeysType.DUP_KEYS)
-                        ? PreAggStatus.on()
+                        ? PreAggStatus.unset()
                         : PreAggStatus.off("For direct index scan.");
 
                 scan = new LogicalOlapScan(unboundRelation.getRelationId(),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/mv/AbstractSelectMaterializedIndexRule.java

@@ -88,17 +88,7 @@ protected boolean shouldSelectIndexWithAgg(LogicalOlapScan scan) {
             case AGG_KEYS:
             case UNIQUE_KEYS:
             case DUP_KEYS:
-                // SelectMaterializedIndexWithAggregate(R1) run before SelectMaterializedIndexWithoutAggregate(R2)
-                // if R1 selects baseIndex and preAggStatus is off
-                // we should give a chance to R2 to check if some prefix-index can be selected
-                // so if R1 selects baseIndex and preAggStatus is off, we keep scan's index unselected in order to
-                // let R2 to get a chance to do its work
-                // at last, after R1, the scan may be the 4 status
-                // 1. preAggStatus is ON and baseIndex is selected, it means select baseIndex is correct.
-                // 2. preAggStatus is ON and some other Index is selected, this is correct, too.
-                // 3. preAggStatus is OFF, no index is selected, it means R2 could get a chance to run
-                // so we check the preAggStatus and if some index is selected to make sure R1 can be run only once
-                return scan.getPreAggStatus().isOn() && !scan.isIndexSelected();
+                return !scan.isIndexSelected();
             default:
                 return false;
         }

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/mv/SelectMaterializedIndexWithAggregate.java

@@ -211,7 +211,7 @@ public List<Rule> buildRules() {
                                     result.exprRewriteMap.projectExprMap);
                             LogicalProject<LogicalOlapScan> newProject = new LogicalProject<>(
                                     generateNewOutputsWithMvOutputs(mvPlan, newProjectList),
-                                    scan.withMaterializedIndexSelected(result.preAggStatus, result.indexId));
+                                    scan.withMaterializedIndexSelected(result.indexId));
                             return new LogicalProject<>(generateProjectsAlias(agg.getOutputs(), slotContext),
                                     new ReplaceExpressions(slotContext)
                                             .replace(
@@ -259,9 +259,6 @@ public List<Rule> buildRules() {
                                                 filter.getExpressions(), project.getExpressions()
                                         ))
                                 );
-                                if (mvPlanWithoutAgg.getSelectedIndexId() == result.indexId) {
-                                    mvPlanWithoutAgg = mvPlanWithoutAgg.withPreAggStatus(result.preAggStatus);
-                                }
                                 SlotContext slotContextWithoutAgg = generateBaseScanExprToMvExpr(mvPlanWithoutAgg);
 
                                 return agg.withChildren(new LogicalProject(
@@ -535,7 +532,7 @@ public List<Rule> buildRules() {
                                     result.exprRewriteMap.projectExprMap);
                             LogicalProject<Plan> newProject = new LogicalProject<>(
                                     generateNewOutputsWithMvOutputs(mvPlan, newProjectList),
-                                    scan.withMaterializedIndexSelected(result.preAggStatus, result.indexId));
+                                    scan.withMaterializedIndexSelected(result.indexId));
 
                             return new LogicalProject<>(generateProjectsAlias(agg.getOutputs(), slotContext),
                                     new ReplaceExpressions(slotContext).replace(new LogicalAggregate<>(
@@ -552,16 +549,7 @@ public List<Rule> buildRules() {
     }
 
     private static LogicalOlapScan createLogicalOlapScan(LogicalOlapScan scan, SelectResult result) {
-        LogicalOlapScan mvPlan;
-        if (result.preAggStatus.isOff()) {
-            // we only set preAggStatus and make index unselected to let SelectMaterializedIndexWithoutAggregate
-            // have a chance to run and select proper index
-            mvPlan = scan.withPreAggStatus(result.preAggStatus);
-        } else {
-            mvPlan =
-                    scan.withMaterializedIndexSelected(result.preAggStatus, result.indexId);
-        }
-        return mvPlan;
+        return scan.withMaterializedIndexSelected(result.indexId);
     }
 
     ///////////////////////////////////////////////////////////////////////////

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/mv/SelectMaterializedIndexWithoutAggregate.java

@@ -27,7 +27,6 @@
 import org.apache.doris.nereids.rules.rewrite.RewriteRuleFactory;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.Slot;
-import org.apache.doris.nereids.trees.plans.PreAggStatus;
 import org.apache.doris.nereids.trees.plans.logical.LogicalFilter;
 import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
@@ -185,7 +184,7 @@ public static LogicalOlapScan select(
                 break;
             case DUP_KEYS:
                 if (table.getIndexIdToMeta().size() == 1) {
-                    return scan.withMaterializedIndexSelected(PreAggStatus.on(), baseIndexId);
+                    return scan.withMaterializedIndexSelected(baseIndexId);
                 }
                 break;
             default:
@@ -210,19 +209,10 @@ public static LogicalOlapScan select(
             // this is fail-safe for select mv
             // select baseIndex if bestIndex's slots' data types are different from baseIndex
             bestIndex = isSameDataType(scan, bestIndex, requiredSlots.get()) ? bestIndex : baseIndexId;
-            return scan.withMaterializedIndexSelected(PreAggStatus.on(), bestIndex);
+            return scan.withMaterializedIndexSelected(bestIndex);
         } else {
-            final PreAggStatus preAggStatus;
-            if (preAggEnabledByHint(scan)) {
-                // PreAggStatus could be enabled by pre-aggregation hint for agg-keys and unique-keys.
-                preAggStatus = PreAggStatus.on();
-            } else {
-                // if PreAggStatus is OFF, we use the message from SelectMaterializedIndexWithAggregate
-                preAggStatus = scan.getPreAggStatus().isOff() ? scan.getPreAggStatus()
-                        : PreAggStatus.off("No aggregate on scan.");
-            }
             if (table.getIndexIdToMeta().size() == 1) {
-                return scan.withMaterializedIndexSelected(preAggStatus, baseIndexId);
+                return scan.withMaterializedIndexSelected(baseIndexId);
             }
             int baseIndexKeySize = table.getKeyColumnsByIndexId(table.getBaseIndexId()).size();
             // No aggregate on scan.
@@ -235,13 +225,13 @@ public static LogicalOlapScan select(
 
             if (candidates.size() == 1) {
                 // `candidates` only have base index.
-                return scan.withMaterializedIndexSelected(preAggStatus, baseIndexId);
+                return scan.withMaterializedIndexSelected(baseIndexId);
             } else {
                 long bestIndex = selectBestIndex(candidates, scan, predicatesSupplier.get(), requiredExpr.get());
                 // this is fail-safe for select mv
                 // select baseIndex if bestIndex's slots' data types are different from baseIndex
                 bestIndex = isSameDataType(scan, bestIndex, requiredSlots.get()) ? bestIndex : baseIndexId;
-                return scan.withMaterializedIndexSelected(preAggStatus, bestIndex);
+                return scan.withMaterializedIndexSelected(bestIndex);
             }
         }
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PreAggStatus.java

@@ -26,10 +26,11 @@
 public class PreAggStatus {
 
     private enum Status {
-        ON, OFF
+        ON, OFF, UNSET
     }
 
     private static final PreAggStatus PRE_AGG_ON = new PreAggStatus(Status.ON, "");
+    private static final PreAggStatus PRE_AGG_UNSET = new PreAggStatus(Status.UNSET, "");
     private final Status status;
     private final String offReason;
 
@@ -46,6 +47,10 @@ public boolean isOff() {
         return status == Status.OFF;
     }
 
+    public boolean isUnset() {
+        return status == Status.UNSET;
+    }
+
     public String getOffReason() {
         return offReason;
     }
@@ -58,6 +63,10 @@ public PreAggStatus offOrElse(Supplier<PreAggStatus> supplier) {
         }
     }
 
+    public static PreAggStatus unset() {
+        return PRE_AGG_UNSET;
+    }
+
     public static PreAggStatus on() {
         return PRE_AGG_ON;
     }
@@ -70,8 +79,10 @@ public static PreAggStatus off(String reason) {
     public String toString() {
         if (status == Status.ON) {
             return "ON";
-        } else {
+        } else if (status == Status.OFF) {
             return "OFF, " + offReason;
+        } else {
+            return "UNSET";
         }
     }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java

@@ -126,15 +126,15 @@ public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier) {
         this(id, table, qualifier, Optional.empty(), Optional.empty(),
                 table.getPartitionIds(), false,
                 ImmutableList.of(),
-                -1, false, PreAggStatus.on(), ImmutableList.of(), ImmutableList.of(),
+                -1, false, PreAggStatus.unset(), ImmutableList.of(), ImmutableList.of(),
                 Maps.newHashMap(), Optional.empty(), false, false);
     }
 
     public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier, List<Long> tabletIds,
             List<String> hints, Optional<TableSample> tableSample) {
         this(id, table, qualifier, Optional.empty(), Optional.empty(),
                 table.getPartitionIds(), false, tabletIds,
-                -1, false, PreAggStatus.on(), ImmutableList.of(), hints, Maps.newHashMap(),
+                -1, false, PreAggStatus.unset(), ImmutableList.of(), hints, Maps.newHashMap(),
                 tableSample, false, false);
     }
 
@@ -143,7 +143,7 @@ public LogicalOlapScan(RelationId id, OlapTable table, List<String> qualifier, L
         this(id, table, qualifier, Optional.empty(), Optional.empty(),
                 // must use specifiedPartitions here for prune partition by sql like 'select * from t partition p1'
                 specifiedPartitions, false, tabletIds,
-                -1, false, PreAggStatus.on(), specifiedPartitions, hints, Maps.newHashMap(),
+                -1, false, PreAggStatus.unset(), specifiedPartitions, hints, Maps.newHashMap(),
                 tableSample, false, false);
     }
 
@@ -275,11 +275,11 @@ public LogicalOlapScan withSelectedPartitionIds(List<Long> selectedPartitionIds)
                 hints, cacheSlotWithSlotName, tableSample, directMvScan, projectPulledUp);
     }
 
-    public LogicalOlapScan withMaterializedIndexSelected(PreAggStatus preAgg, long indexId) {
+    public LogicalOlapScan withMaterializedIndexSelected(long indexId) {
         return new LogicalOlapScan(relationId, (Table) table, qualifier,
                 Optional.empty(), Optional.of(getLogicalProperties()),
                 selectedPartitionIds, partitionPruned, selectedTabletIds,
-                indexId, true, preAgg, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName,
+                indexId, true, PreAggStatus.unset(), manuallySpecifiedPartitions, hints, cacheSlotWithSlotName,
                 tableSample, directMvScan, projectPulledUp);
     }
 
@@ -432,6 +432,10 @@ public boolean isDirectMvScan() {
         return directMvScan;
     }
 
+    public boolean isPreAggStatusUnSet() {
+        return preAggStatus.isUnset();
+    }
+
     private List<SlotReference> createSlotsVectorized(List<Column> columns) {
         List<String> qualified = qualified();
         Object[] slots = new Object[columns.size()];