StructInfoMap.java
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
package org.apache.doris.nereids.memo;
import org.apache.doris.catalog.TableIf;
import org.apache.doris.common.Pair;
import org.apache.doris.nereids.CascadesContext;
import org.apache.doris.nereids.rules.exploration.mv.StructInfo;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer;
import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation;
import org.apache.doris.nereids.trees.plans.logical.LogicalRelation;
import com.google.common.collect.Sets;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import javax.annotation.Nullable;
/**
* Representation for group in cascades optimizer.
*/
public class StructInfoMap {
public static final Logger LOG = LogManager.getLogger(StructInfoMap.class);
/**
* The map key is the relation id bit set to get corresponding plan accurately
*/
private final Map<BitSet, Pair<GroupExpression, List<BitSet>>> groupExpressionMap = new HashMap<>();
/**
* The map key is the relation id bit set to get corresponding plan accurately
*/
private final Map<BitSet, StructInfo> infoMap = new HashMap<>();
private long refreshVersion = 0;
/**
* get struct info according to table map
*
* @param tableMap the original table map
* @param group the group that the mv matched
* @return struct info or null if not found
*/
public @Nullable StructInfo getStructInfo(CascadesContext cascadesContext, BitSet tableMap, Group group,
Plan originPlan, boolean forceRefresh) {
StructInfo structInfo = infoMap.get(tableMap);
if (structInfo != null) {
return structInfo;
}
if (groupExpressionMap.isEmpty() || !groupExpressionMap.containsKey(tableMap)) {
refresh(group, cascadesContext, tableMap, new HashSet<>(),
forceRefresh);
group.getStructInfoMap().setRefreshVersion(cascadesContext.getMemo().getRefreshVersion());
}
if (groupExpressionMap.containsKey(tableMap)) {
Pair<GroupExpression, List<BitSet>> groupExpressionBitSetPair = getGroupExpressionWithChildren(tableMap);
structInfo = constructStructInfo(groupExpressionBitSetPair.first, groupExpressionBitSetPair.second,
originPlan, cascadesContext);
infoMap.put(tableMap, structInfo);
}
return structInfo;
}
public Set<BitSet> getTableMaps() {
return groupExpressionMap.keySet();
}
public Pair<GroupExpression, List<BitSet>> getGroupExpressionWithChildren(BitSet tableMap) {
return groupExpressionMap.get(tableMap);
}
public void setRefreshVersion(long refreshVersion) {
this.refreshVersion = refreshVersion;
}
private StructInfo constructStructInfo(GroupExpression groupExpression, List<BitSet> children,
Plan originPlan, CascadesContext cascadesContext) {
// this plan is not origin plan, should record origin plan in struct info
Plan plan = constructPlan(groupExpression, children);
return originPlan == null ? StructInfo.of(plan, cascadesContext)
: StructInfo.of(plan, originPlan, cascadesContext);
}
private Plan constructPlan(GroupExpression groupExpression, List<BitSet> children) {
List<Plan> childrenPlan = new ArrayList<>();
for (int i = 0; i < children.size(); i++) {
StructInfoMap structInfoMap = groupExpression.child(i).getStructInfoMap();
BitSet childMap = children.get(i);
Pair<GroupExpression, List<BitSet>> groupExpressionBitSetPair
= structInfoMap.getGroupExpressionWithChildren(childMap);
childrenPlan.add(
constructPlan(groupExpressionBitSetPair.first, groupExpressionBitSetPair.second));
}
// need to clear current group expression info by using withGroupExpression
// this plan would copy into memo, if with group expression, would cause err
return groupExpression.getPlan().withChildren(childrenPlan).withGroupExpression(Optional.empty());
}
/**
* refresh group expression map
*
* @param group the root group
* @param targetBitSet refreshed group expression table bitset must intersect with the targetBitSet
*
*/
public void refresh(Group group, CascadesContext cascadesContext,
BitSet targetBitSet,
Set<Integer> refreshedGroup,
boolean forceRefresh) {
StructInfoMap structInfoMap = group.getStructInfoMap();
refreshedGroup.add(group.getGroupId().asInt());
long memoVersion = cascadesContext.getMemo().getRefreshVersion();
if (!structInfoMap.getTableMaps().isEmpty() && memoVersion == structInfoMap.refreshVersion) {
return;
}
for (GroupExpression groupExpression : group.getLogicalExpressions()) {
List<Set<BitSet>> childrenTableMap = new LinkedList<>();
if (groupExpression.children().isEmpty()) {
BitSet leaf = constructLeaf(groupExpression, cascadesContext, forceRefresh);
if (leaf.isEmpty()) {
break;
}
groupExpressionMap.put(leaf, Pair.of(groupExpression, new LinkedList<>()));
continue;
}
for (Group child : groupExpression.children()) {
StructInfoMap childStructInfoMap = child.getStructInfoMap();
if (!refreshedGroup.contains(child.getGroupId().asInt())) {
childStructInfoMap.refresh(child, cascadesContext, targetBitSet, refreshedGroup, forceRefresh);
childStructInfoMap.setRefreshVersion(memoVersion);
}
Set<BitSet> filteredTableMaps = new HashSet<>();
for (BitSet tableMaps : child.getStructInfoMap().getTableMaps()) {
// filter the tableSet that used intersects with targetBitSet
if (!targetBitSet.isEmpty() && !tableMaps.intersects(targetBitSet)) {
continue;
}
filteredTableMaps.add(tableMaps);
}
if (!filteredTableMaps.isEmpty()) {
childrenTableMap.add(filteredTableMaps);
}
}
if (childrenTableMap.isEmpty()) {
continue;
}
// if groupExpression which has the same table set have refreshed, continue
BitSet eachGroupExpressionTableSet = new BitSet();
for (Set<BitSet> groupExpressionBitSet : childrenTableMap) {
for (BitSet bitSet : groupExpressionBitSet) {
eachGroupExpressionTableSet.or(bitSet);
}
}
if (groupExpressionMap.containsKey(eachGroupExpressionTableSet)) {
// for the group expressions of group, only need to refresh any of the group expression
// when they have the same group expression table set
continue;
}
// if cumulative child table map is different from current
// or current group expression map is empty, should update the groupExpressionMap currently
Collection<Pair<BitSet, List<BitSet>>> bitSetWithChildren = cartesianProduct(childrenTableMap);
for (Pair<BitSet, List<BitSet>> bitSetWithChild : bitSetWithChildren) {
groupExpressionMap.putIfAbsent(bitSetWithChild.first,
Pair.of(groupExpression, bitSetWithChild.second));
}
}
}
private BitSet constructLeaf(GroupExpression groupExpression, CascadesContext cascadesContext,
boolean forceRefresh) {
Plan plan = groupExpression.getPlan();
BitSet tableMap = new BitSet();
if (plan instanceof LogicalCatalogRelation) {
LogicalCatalogRelation logicalCatalogRelation = (LogicalCatalogRelation) plan;
TableIf table = logicalCatalogRelation.getTable();
// If disable materialized view nest rewrite, and mv already rewritten successfully once, doesn't construct
// table id map for nest mv rewrite
if (!forceRefresh && cascadesContext.getStatementContext()
.getMaterializationRewrittenSuccessSet().contains(table.getFullQualifiers())) {
return tableMap;
}
tableMap.set(logicalCatalogRelation.getRelationId().asInt());
}
// one row relation / CTE consumer
if (plan instanceof LogicalCTEConsumer || plan instanceof LogicalEmptyRelation
|| plan instanceof LogicalOneRowRelation) {
tableMap.set(((LogicalRelation) plan).getRelationId().asInt());
}
return tableMap;
}
private Collection<Pair<BitSet, List<BitSet>>> cartesianProduct(List<Set<BitSet>> childrenTableMap) {
Set<List<BitSet>> cartesianLists = Sets.cartesianProduct(childrenTableMap);
List<Pair<BitSet, List<BitSet>>> resultPairSet = new LinkedList<>();
for (List<BitSet> bitSetList : cartesianLists) {
BitSet bitSet = new BitSet();
for (BitSet b : bitSetList) {
bitSet.or(b);
}
resultPairSet.add(Pair.of(bitSet, bitSetList));
}
return resultPairSet;
}
@Override
public String toString() {
return "StructInfoMap{ groupExpressionMap = " + groupExpressionMap + ", infoMap = " + infoMap + '}';
}
}