[improvement](mtmv) Support to get tables in materialized view when collecting table in plan

fe/fe-core/src/main/java/org/apache/doris/mtmv/MTMVPlanUtil.java

@@ -95,17 +95,17 @@ private static Set<BaseTableInfo> getBaseTables(Plan plan) {
  TableCollectorContext collectorContext =
  new TableCollector.TableCollectorContext(
  com.google.common.collect.Sets
- .newHashSet(TableType.values()));
+ .newHashSet(TableType.values()), true);
  plan.accept(TableCollector.INSTANCE, collectorContext);
-List<TableIf> collectedTables = collectorContext.getCollectedTables();
+Set<TableIf> collectedTables = collectorContext.getCollectedTables();
  return transferTableIfToInfo(collectedTables);
  }
 
  private static Set<BaseTableInfo> getBaseViews(Plan plan) {
  return Sets.newHashSet();
  }
 
- private static Set<BaseTableInfo> transferTableIfToInfo(List<TableIf> tables) {
+ private static Set<BaseTableInfo> transferTableIfToInfo(Set<TableIf> tables) {
  Set<BaseTableInfo> result = com.google.common.collect.Sets.newHashSet();
  for (TableIf table : tables) {
  result.add(new BaseTableInfo(table));

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/exploration/mv/InitMaterializationContextHook.java

@@ -66,9 +66,9 @@ public void initMaterializationContext(CascadesContext cascadesContext) {
  return;
  }
  Plan rewritePlan = cascadesContext.getRewritePlan();
- TableCollectorContext collectorContext = new TableCollectorContext(Sets.newHashSet());
+ TableCollectorContext collectorContext = new TableCollectorContext(Sets.newHashSet(), true);
  rewritePlan.accept(TableCollector.INSTANCE, collectorContext);
-List<TableIf> collectedTables = collectorContext.getCollectedTables();
+Set<TableIf> collectedTables = collectorContext.getCollectedTables();
  if (collectedTables.isEmpty()) {
  return;
  }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateMTMVInfo.java

@@ -320,9 +320,9 @@ private PartitionDesc generatePartitionDesc(MTMVRelatedTableIf relatedTable, Con
 
  private void analyzeBaseTables(Plan plan) {
  TableCollectorContext collectorContext =
- new TableCollector.TableCollectorContext(Sets.newHashSet(TableType.MATERIALIZED_VIEW));
+ new TableCollector.TableCollectorContext(Sets.newHashSet(TableType.MATERIALIZED_VIEW), true);
  plan.accept(TableCollector.INSTANCE, collectorContext);
-List<TableIf> collectedTables = collectorContext.getCollectedTables();
+Set<TableIf> collectedTables = collectorContext.getCollectedTables();
  if (!CollectionUtils.isEmpty(collectedTables)) {
  throw new AnalysisException("can not contain MATERIALIZED_VIEW");
  }

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

@@ -23,6 +23,7 @@
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
 import org.apache.doris.nereids.trees.plans.RelationId;
+import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
 import org.apache.doris.nereids.util.Utils;
 
 import java.util.List;
@@ -68,4 +69,9 @@ public Plan withGroupExprLogicalPropChildren(Optional<GroupExpression> groupExpr
  Optional<LogicalProperties> logicalProperties, List<Plan> children) {
  return new LogicalTestScan(relationId, table, qualifier, groupExpression, logicalProperties);
  }
+
+ @Override
+ public <R, C> R accept(PlanVisitor<R, C> visitor, C context) {
+ return visitor.visitLogicalTestScan(this, context);
+ }
 }

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

@@ -20,6 +20,7 @@
 import org.apache.doris.nereids.analyzer.UnboundOneRowRelation;
 import org.apache.doris.nereids.analyzer.UnboundRelation;
 import org.apache.doris.nereids.analyzer.UnboundTVFRelation;
+import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalDeferMaterializeOlapScan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalEsScan;
@@ -31,6 +32,8 @@
 import org.apache.doris.nereids.trees.plans.logical.LogicalRelation;
 import org.apache.doris.nereids.trees.plans.logical.LogicalSchemaScan;
 import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation;
+import org.apache.doris.nereids.trees.plans.logical.LogicalTestScan;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalCatalogRelation;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalDeferMaterializeOlapScan;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalEmptyRelation;
 import org.apache.doris.nereids.trees.plans.physical.PhysicalEsScan;
@@ -54,8 +57,16 @@ public interface RelationVisitor<R, C> {
 
  R visitLogicalRelation(LogicalRelation logicalRelation, C context);
 
+ default R visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, C context) {
+ return visitLogicalRelation(catalogRelation, context);
+ }
+
  R visitPhysicalRelation(PhysicalRelation physicalRelation, C context);
 
+ default R visitPhysicalCatalogRelation(PhysicalCatalogRelation catalogRelation, C context) {
+ return visitPhysicalRelation(catalogRelation, context);
+ }
+
  // *******************************
  // unbound relations
  // *******************************
@@ -81,42 +92,46 @@ default R visitLogicalEmptyRelation(LogicalEmptyRelation emptyRelation, C contex
  }
 
  default R visitLogicalEsScan(LogicalEsScan esScan, C context) {
- return visitLogicalRelation(esScan, context);
+ return visitLogicalCatalogRelation(esScan, context);
  }
 
  default R visitLogicalFileScan(LogicalFileScan fileScan, C context) {
- return visitLogicalRelation(fileScan, context);
+ return visitLogicalCatalogRelation(fileScan, context);
  }
 
  default R visitLogicalJdbcScan(LogicalJdbcScan jdbcScan, C context) {
- return visitLogicalRelation(jdbcScan, context);
+ return visitLogicalCatalogRelation(jdbcScan, context);
  }
 
  default R visitLogicalOdbcScan(LogicalOdbcScan odbcScan, C context) {
- return visitLogicalRelation(odbcScan, context);
+ return visitLogicalCatalogRelation(odbcScan, context);
  }
 
  default R visitLogicalOlapScan(LogicalOlapScan olapScan, C context) {
- return visitLogicalRelation(olapScan, context);
+ return visitLogicalCatalogRelation(olapScan, context);
  }
 
  default R visitLogicalDeferMaterializeOlapScan(
  LogicalDeferMaterializeOlapScan deferMaterializeOlapScan, C context) {
- return visitLogicalRelation(deferMaterializeOlapScan, context);
+ return visitLogicalCatalogRelation(deferMaterializeOlapScan, context);
  }
 
  default R visitLogicalOneRowRelation(LogicalOneRowRelation oneRowRelation, C context) {
  return visitLogicalRelation(oneRowRelation, context);
  }
 
  default R visitLogicalSchemaScan(LogicalSchemaScan schemaScan, C context) {
- return visitLogicalRelation(schemaScan, context);
+ return visitLogicalCatalogRelation(schemaScan, context);
  }
 
  default R visitLogicalTVFRelation(LogicalTVFRelation tvfRelation, C context) {
  return visitLogicalRelation(tvfRelation, context);
  }
 
+ default R visitLogicalTestScan(LogicalTestScan testScan, C context) {
+ return visitLogicalCatalogRelation(testScan, context);
+ }
+
  // *******************************
  // physical relations
  // *******************************
@@ -126,36 +141,36 @@ default R visitPhysicalEmptyRelation(PhysicalEmptyRelation emptyRelation, C cont
  }
 
  default R visitPhysicalEsScan(PhysicalEsScan esScan, C context) {
- return visitPhysicalRelation(esScan, context);
+ return visitPhysicalCatalogRelation(esScan, context);
  }
 
  default R visitPhysicalFileScan(PhysicalFileScan fileScan, C context) {
- return visitPhysicalRelation(fileScan, context);
+ return visitPhysicalCatalogRelation(fileScan, context);
  }
 
  default R visitPhysicalJdbcScan(PhysicalJdbcScan jdbcScan, C context) {
- return visitPhysicalRelation(jdbcScan, context);
+ return visitPhysicalCatalogRelation(jdbcScan, context);
  }
 
  default R visitPhysicalOdbcScan(PhysicalOdbcScan odbcScan, C context) {
- return visitPhysicalRelation(odbcScan, context);
+ return visitPhysicalCatalogRelation(odbcScan, context);
  }
 
  default R visitPhysicalOlapScan(PhysicalOlapScan olapScan, C context) {
- return visitPhysicalRelation(olapScan, context);
+ return visitPhysicalCatalogRelation(olapScan, context);
  }
 
  default R visitPhysicalDeferMaterializeOlapScan(
  PhysicalDeferMaterializeOlapScan deferMaterializeOlapScan, C context) {
- return visitPhysicalRelation(deferMaterializeOlapScan, context);
+ return visitPhysicalCatalogRelation(deferMaterializeOlapScan, context);
  }
 
  default R visitPhysicalOneRowRelation(PhysicalOneRowRelation oneRowRelation, C context) {
  return visitPhysicalRelation(oneRowRelation, context);
  }
 
  default R visitPhysicalSchemaScan(PhysicalSchemaScan schemaScan, C context) {
- return visitPhysicalRelation(schemaScan, context);
+ return visitPhysicalCatalogRelation(schemaScan, context);
  }
 
  default R visitPhysicalTVFRelation(PhysicalTVFRelation tvfRelation, C context) {

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

@@ -17,53 +17,100 @@
 
 package org.apache.doris.nereids.trees.plans.visitor;
 
+import org.apache.doris.catalog.MTMV;
 import org.apache.doris.catalog.TableIf;
 import org.apache.doris.catalog.TableIf.TableType;
+import org.apache.doris.common.AnalysisException;
+import org.apache.doris.mtmv.MTMVCache;
+import org.apache.doris.mtmv.MTMVPlanUtil;
 import org.apache.doris.nereids.trees.plans.Plan;
-import org.apache.doris.nereids.trees.plans.algebra.CatalogRelation;
+import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
+import org.apache.doris.nereids.trees.plans.physical.PhysicalCatalogRelation;
 import org.apache.doris.nereids.trees.plans.visitor.TableCollector.TableCollectorContext;
 
-import java.util.ArrayList;
-import java.util.List;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.util.HashSet;
 import java.util.Set;
 
 /**
  * Collect the table in plan
  * Note: will not get table if table is eliminated by EmptyRelation in rewrite.
+ * View expand is in RBO, if call this method with the plan after RBO, this will get base tables in view, or will not.
+ * Materialized view is extended or not can be controlled by the field expand
  */
-public class TableCollector extends DefaultPlanVisitor<Void, TableCollectorContext> {
+public class TableCollector extends DefaultPlanVisitor<Plan, TableCollectorContext> {
 
  public static final TableCollector INSTANCE = new TableCollector();
+ private static final Logger LOG = LogManager.getLogger(TableCollector.class);
 
  @Override
- public Void visit(Plan plan, TableCollectorContext context) {
- if (plan instanceof CatalogRelation) {
- TableIf table = ((CatalogRelation) plan).getTable();
- if (context.getTargetTableTypes().isEmpty() || context.getTargetTableTypes().contains(table.getType())) {
- context.getCollectedTables().add(table);
- }
+ public Plan visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, TableCollectorContext context) {
+ TableIf table = catalogRelation.getTable();
+ if (context.getTargetTableTypes().isEmpty() || context.getTargetTableTypes().contains(table.getType())) {
+ context.getCollectedTables().add(table);
+ }
+ if (table instanceof MTMV) {
+ expandMvAndCollect((MTMV) table, context);
+ }
+ return catalogRelation;
+ }
+
+ @Override
+ public Plan visitPhysicalCatalogRelation(PhysicalCatalogRelation catalogRelation, TableCollectorContext context) {
+ TableIf table = catalogRelation.getTable();
+ if (context.getTargetTableTypes().isEmpty() || context.getTargetTableTypes().contains(table.getType())) {
+ context.getCollectedTables().add(table);
+ }
+ if (table instanceof MTMV) {
+ expandMvAndCollect((MTMV) table, context);
+ }
+ return catalogRelation;
+ }
+
+ private void expandMvAndCollect(MTMV mtmv, TableCollectorContext context) {
+ if (!context.isExpand()) {
+ return;
+ }
+ try {
+ MTMVCache expandedMv = MTMVCache.from(mtmv, MTMVPlanUtil.createMTMVContext(mtmv));
+ expandedMv.getLogicalPlan().accept(this, context);
+ } catch (AnalysisException e) {
+ LOG.error(String.format(
+ "table collector expand fail, mtmv name is %s, targetTableTypes is %s",
+ mtmv.getName(), context.targetTableTypes), e);
+ throw new org.apache.doris.nereids.exceptions.AnalysisException(
+ String.format("expand mv and collect table fail, mv name is %s, mv sql is %s",
+ mtmv.getName(), mtmv.getQuerySql()), e);
  }
- return super.visit(plan, context);
  }
 
  /**
  * The context for table collecting, it contains the target collect table types
  * and the result of collect.
  */
  public static final class TableCollectorContext {
- private final List<TableIf> collectedTables = new ArrayList<>();
+ private final Set<TableIf> collectedTables = new HashSet<>();
  private final Set<TableType> targetTableTypes;
+ // if expand the mv or not
+ private final boolean expand;
 
- public TableCollectorContext(Set<TableType> targetTableTypes) {
+ public TableCollectorContext(Set<TableType> targetTableTypes, boolean expand) {
  this.targetTableTypes = targetTableTypes;
+ this.expand = expand;
  }
 
- public List<TableIf> getCollectedTables() {
+ public Set<TableIf> getCollectedTables() {
  return collectedTables;
  }
 
  public Set<TableType> getTargetTableTypes() {
  return targetTableTypes;
  }
+
+ public boolean isExpand() {
+ return expand;
+ }
  }
 }