[refactor](fe) Refactor MV rewrite StructInfo lookup by relation ids

Author: yangtao555

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

@@ -274,8 +274,9 @@ public enum TableFrom {
     // Record used table and it's used partitions
     private final Multimap<List<String>, Pair<RelationId, Set<String>>> tableUsedPartitionNameMap =
             HashMultimap.create();
-    // Record query common table id to relation id mapping, this is used for mv rewrite
-    private final Multimap<Integer, Integer> commonTableIdToRelationIdToMap = HashMultimap.create();
+    // Record statement-scope table ids to relation ids for MV rewrite.
+    // One table id may map to multiple relation ids because of aliases and nested MV scan alternatives.
+    private final Multimap<Integer, Integer> tableIdToRelationIds = HashMultimap.create();
 
     // Record mtmv and valid partitions map because this is time-consuming behavior
     private final Map<BaseTableInfo, Collection<Partition>> mvCanRewritePartitionsMap = new HashMap<>();
@@ -304,8 +305,6 @@ public enum TableFrom {
     // mark is rewritten in RBO phase, if rewritten in RBO phase should set true
     private boolean preMvRewritten = false;
 
-    private final Set<List<String>> materializationRewrittenSuccessSet = new HashSet<>();
-
     private boolean isInsert = false;
     private Optional<Map<TableIf, Set<Expression>>> mvRefreshPredicates = Optional.empty();
 
@@ -1109,20 +1108,12 @@ public void setPreMvRewritten(boolean preMvRewritten) {
         this.preMvRewritten = preMvRewritten;
     }
 
-    public Set<List<String>> getMaterializationRewrittenSuccessSet() {
-        return materializationRewrittenSuccessSet;
-    }
-
-    public void addMaterializationRewrittenSuccess(List<String> materializationQualifier) {
-        this.materializationRewrittenSuccessSet.add(materializationQualifier);
-    }
-
     public Multimap<List<String>, Pair<RelationId, Set<String>>> getTableUsedPartitionNameMap() {
         return tableUsedPartitionNameMap;
     }
 
-    public Multimap<Integer, Integer> getCommonTableIdToRelationIdMap() {
-        return commonTableIdToRelationIdToMap;
+    public Multimap<Integer, Integer> getTableIdToRelationIds() {
+        return tableIdToRelationIds;
     }
 
     public Map<BaseTableInfo, Collection<Partition>> getMvCanRewritePartitionsMap() {

fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Group.java

@@ -83,7 +83,7 @@ public class Group {
     private List<PhysicalProperties> chosenEnforcerPropertiesList = new ArrayList<>();
     private List<Integer> chosenEnforcerIdList = new ArrayList<>();
 
-    private StructInfoMap structInfoMap = new StructInfoMap();
+    private final StructInfoMap structInfoMap;
 
     /**
      * Constructor for Group.
@@ -92,6 +92,7 @@ public class Group {
      */
     public Group(GroupId groupId, GroupExpression groupExpression, LogicalProperties logicalProperties) {
         this.groupId = groupId;
+        this.structInfoMap = new StructInfoMap(this);
         addGroupExpression(groupExpression);
         this.logicalProperties = logicalProperties;
         this.groupPlan = new GroupPlan(this);
@@ -104,6 +105,7 @@ public Group(GroupId groupId, GroupExpression groupExpression, LogicalProperties
      */
     public Group(GroupId groupId, LogicalProperties logicalProperties) {
         this.groupId = groupId;
+        this.structInfoMap = new StructInfoMap(this);
         this.logicalProperties = logicalProperties;
         this.groupPlan = new GroupPlan(this);
     }

fe/fe-core/src/main/java/org/apache/doris/nereids/memo/Memo.java

@@ -17,9 +17,9 @@
 
 package org.apache.doris.nereids.memo;
 
-import org.apache.doris.catalog.MTMV;
 import org.apache.doris.common.IdGenerator;
 import org.apache.doris.common.Pair;
+import org.apache.doris.nereids.StatementContext;
 import org.apache.doris.nereids.cost.Cost;
 import org.apache.doris.nereids.cost.CostCalculator;
 import org.apache.doris.nereids.metrics.EventChannel;
@@ -52,7 +52,6 @@
 import org.apache.logging.log4j.Logger;
 
 import java.util.ArrayList;
-import java.util.BitSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -61,7 +60,6 @@
 import java.util.Optional;
 import java.util.PriorityQueue;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
@@ -76,9 +74,7 @@ public class Memo {
             EventChannel.getDefaultChannel().addConsumers(new LogConsumer(GroupMergeEvent.class, EventChannel.LOG)));
     private static long stateId = 0;
     private final ConnectContext connectContext;
-    // The key is the query tableId, the value is the refresh version when last refresh, this is needed
-    // because struct info refresh base on target tableId.
-    private final Map<Integer, AtomicInteger> refreshVersion = new HashMap<>();
+    private final StatementContext statementContext;
     private final Map<Class<? extends AbstractMaterializedViewRule>, Set<Long>> materializationCheckSuccessMap =
             new LinkedHashMap<>();
     private final Map<Class<? extends AbstractMaterializedViewRule>, Set<Long>> materializationCheckFailMap =
@@ -93,11 +89,17 @@ public class Memo {
     public Memo() {
         this.root = null;
         this.connectContext = null;
+        this.statementContext = null;
     }
 
     public Memo(ConnectContext connectContext, Plan plan) {
+        this(connectContext == null ? null : connectContext.getStatementContext(), plan);
+    }
+
+    private Memo(StatementContext statementContext, Plan plan) {
+        this.statementContext = statementContext;
+        this.connectContext = statementContext == null ? null : statementContext.getConnectContext();
         this.root = init(plan);
-        this.connectContext = connectContext;
     }
 
     public static long getStateId() {
@@ -132,30 +134,6 @@ public int getGroupExpressionsSize() {
         return groupExpressions.size();
     }
 
-    /** get the refresh version map*/
-    public Map<Integer, AtomicInteger> getRefreshVersion() {
-        return refreshVersion;
-    }
-
-    /** return the incremented refresh version for the given commonTableId*/
-    public long incrementAndGetRefreshVersion(int commonTableId) {
-        return refreshVersion.compute(commonTableId, (k, v) -> {
-            if (v == null) {
-                return new AtomicInteger(1);
-            }
-            v.incrementAndGet();
-            return v;
-        }).get();
-    }
-
-    /** return the incremented refresh version for the given relationId set*/
-    public void incrementAndGetRefreshVersion(BitSet commonTableIdSet) {
-        for (int i = commonTableIdSet.nextSetBit(0); i >= 0;
-                i = commonTableIdSet.nextSetBit(i + 1)) {
-            incrementAndGetRefreshVersion(i);
-        }
-    }
-
     /**
      * Record materialization check result for performance
      */
@@ -379,6 +357,7 @@ public Plan copyOut(GroupExpression logicalExpression, boolean includeGroupExpre
      */
     private Group init(Plan plan) {
         Preconditions.checkArgument(!(plan instanceof GroupPlan), "Cannot init memo by a GroupPlan");
+        registerRelationIdentity(plan);
 
         // initialize children recursively
         List<Group> childrenGroups = new ArrayList<>(plan.arity());
@@ -481,15 +460,7 @@ private CopyInResult doCopyIn(Plan plan, @Nullable Group targetGroup, @Nullable
                     plan.getLogicalProperties(), targetGroup.getLogicalProperties());
             throw new IllegalStateException("Insert a plan into targetGroup but differ in logicalproperties");
         }
-        if (connectContext != null
-                && connectContext.getSessionVariable().isEnableMaterializedViewNestRewrite()
-                && plan instanceof LogicalCatalogRelation
-                && ((CatalogRelation) plan).getTable() instanceof MTMV
-                && !plan.getGroupExpression().isPresent()) {
-            TableId mvCommonTableId
-                    = this.connectContext.getStatementContext().getTableId(((CatalogRelation) plan).getTable());
-            incrementAndGetRefreshVersion(mvCommonTableId.asInt());
-        }
+        registerRelationIdentity(plan);
         Optional<GroupExpression> groupExpr = plan.getGroupExpression();
         if (groupExpr.isPresent()) {
             Preconditions.checkState(groupExpressions.containsKey(groupExpr.get()));
@@ -513,6 +484,24 @@ private CopyInResult doCopyIn(Plan plan, @Nullable Group targetGroup, @Nullable
         // TODO: need to derive logical property if generate new group. currently we not copy logical plan into
     }
 
+    private void registerRelationIdentity(Plan plan) {
+        if (statementContext == null) {
+            return;
+        }
+        if (plan instanceof LogicalCatalogRelation) {
+            // StructInfoMap searches query alternatives by relation id, but each MV context starts from the
+            // table ids in the MV definition. Register both original scans and nested MV scans copied into memo
+            // so those table ids can be expanded back to the currently available relation ids.
+            CatalogRelation catalogRelation = (CatalogRelation) plan;
+            TableId tableId = statementContext.getTableId(catalogRelation.getTable());
+            boolean relationIdentityChanged = statementContext.getTableIdToRelationIds()
+                    .put(tableId.asInt(), catalogRelation.getRelationId().asInt());
+            if (relationIdentityChanged) {
+                groups.values().forEach(group -> group.getStructInfoMap().clearCandidateCache());
+            }
+        }
+    }
+
     private List<Group> rewriteChildrenPlansToGroups(Plan plan, Group targetGroup) {
         List<Group> childrenGroups = Lists.newArrayList();
         for (int i = 0; i < plan.children().size(); i++) {