morrySnow commented on code in PR #28682: URL: https://github.com/apache/doris/pull/28682#discussion_r1434049416
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject Review Comment: nit: `instanceOf` is including `!= null`, so we could remove `!= null` condition ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalAggregate aggrRoot = context.tableToAggrRootMap.get(table); + if (aggrRoot == null) { + return false; + } else if (!checkAggrKeyOnUkOrPk(aggrRoot, table)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkNoFilterOnPullUpTable(List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) { + for (LogicalCatalogRelation table : pullUpTableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else { + List<LogicalFilter> filterList = new ArrayList<>(); + filterList.addAll((Collection<? extends LogicalFilter>) + joinRoot.collect(LogicalFilter.class::isInstance)); + for (LogicalFilter filter : filterList) { + if (filter.child().equals(context.pullUpTable)) { + return false; + } + } + } + } + return true; + } + + private boolean checkAggrKeyOnUkOrPk(LogicalAggregate aggregate, LogicalCatalogRelation table) { + List<Expression> groupByKeys = aggregate.getGroupByExpressions(); + boolean isAllSlotReference = groupByKeys.stream().allMatch(e -> e instanceof SlotReference); + if (!isAllSlotReference) { + return false; + } else { + Set<String> ukInfo = getUkInfoFromConstraint(table); + Set<String> pkInfo = getPkInfoFromConstraint(table); + if (ukInfo == null || pkInfo == null || ukInfo.size() != 1 || pkInfo.size() != 1) { Review Comment: why size == 1? ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalAggregate aggrRoot = context.tableToAggrRootMap.get(table); + if (aggrRoot == null) { + return false; + } else if (!checkAggrKeyOnUkOrPk(aggrRoot, table)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkNoFilterOnPullUpTable(List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) { + for (LogicalCatalogRelation table : pullUpTableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else { + List<LogicalFilter> filterList = new ArrayList<>(); + filterList.addAll((Collection<? extends LogicalFilter>) + joinRoot.collect(LogicalFilter.class::isInstance)); + for (LogicalFilter filter : filterList) { + if (filter.child().equals(context.pullUpTable)) { + return false; + } + } + } + } + return true; + } + + private boolean checkAggrKeyOnUkOrPk(LogicalAggregate aggregate, LogicalCatalogRelation table) { + List<Expression> groupByKeys = aggregate.getGroupByExpressions(); + boolean isAllSlotReference = groupByKeys.stream().allMatch(e -> e instanceof SlotReference); + if (!isAllSlotReference) { + return false; + } else { + Set<String> ukInfo = getUkInfoFromConstraint(table); + Set<String> pkInfo = getPkInfoFromConstraint(table); + if (ukInfo == null || pkInfo == null || ukInfo.size() != 1 || pkInfo.size() != 1) { + return false; + } else { + String ukName = ukInfo.iterator().next(); + String pkName = pkInfo.iterator().next(); + for (Object expr : aggregate.getGroupByExpressions()) { + SlotReference slot = (SlotReference) expr; + if (table.getOutputExprIds().contains(slot.getExprId()) + && (slot.getName().equals(ukName) || slot.getName().equals(pkName))) { + return true; + } + } + return false; + } + } + } + + private boolean checkJoinConditionOnPk(LogicalJoin joinRoot, LogicalCatalogRelation table, PullUpContext context) { + Set<String> pkInfos = getPkInfoFromConstraint(table); + if (pkInfos == null || pkInfos.size() != 1) { + return false; + } + String pkSlot = pkInfos.iterator().next(); + List<LogicalJoin> joinList = new ArrayList<>(); + joinList.addAll((Collection<? extends LogicalJoin>) joinRoot.collect(LogicalJoin.class::isInstance)); + boolean found = false; + for (LogicalJoin join : joinList) { + List<Expression> conditions = join.getHashJoinConjuncts(); + List<LogicalCatalogRelation> basicTableList = new ArrayList<>(); + basicTableList.addAll((Collection<? extends LogicalCatalogRelation>) join + .collect(LogicalCatalogRelation.class::isInstance)); + for (Expression equalTo : conditions) { + if (equalTo instanceof EqualTo Review Comment: EqualTo -> EqualPredicate? EqualPredicate include EqualTo and NullSafeEqualTo ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalAggregate aggrRoot = context.tableToAggrRootMap.get(table); + if (aggrRoot == null) { + return false; + } else if (!checkAggrKeyOnUkOrPk(aggrRoot, table)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkNoFilterOnPullUpTable(List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) { + for (LogicalCatalogRelation table : pullUpTableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else { + List<LogicalFilter> filterList = new ArrayList<>(); + filterList.addAll((Collection<? extends LogicalFilter>) + joinRoot.collect(LogicalFilter.class::isInstance)); + for (LogicalFilter filter : filterList) { + if (filter.child().equals(context.pullUpTable)) { + return false; + } + } + } + } + return true; + } + + private boolean checkAggrKeyOnUkOrPk(LogicalAggregate aggregate, LogicalCatalogRelation table) { + List<Expression> groupByKeys = aggregate.getGroupByExpressions(); + boolean isAllSlotReference = groupByKeys.stream().allMatch(e -> e instanceof SlotReference); + if (!isAllSlotReference) { + return false; + } else { + Set<String> ukInfo = getUkInfoFromConstraint(table); + Set<String> pkInfo = getPkInfoFromConstraint(table); + if (ukInfo == null || pkInfo == null || ukInfo.size() != 1 || pkInfo.size() != 1) { + return false; + } else { + String ukName = ukInfo.iterator().next(); + String pkName = pkInfo.iterator().next(); + for (Object expr : aggregate.getGroupByExpressions()) { + SlotReference slot = (SlotReference) expr; + if (table.getOutputExprIds().contains(slot.getExprId()) + && (slot.getName().equals(ukName) || slot.getName().equals(pkName))) { + return true; + } + } + return false; + } + } + } + + private boolean checkJoinConditionOnPk(LogicalJoin joinRoot, LogicalCatalogRelation table, PullUpContext context) { + Set<String> pkInfos = getPkInfoFromConstraint(table); + if (pkInfos == null || pkInfos.size() != 1) { + return false; + } + String pkSlot = pkInfos.iterator().next(); + List<LogicalJoin> joinList = new ArrayList<>(); + joinList.addAll((Collection<? extends LogicalJoin>) joinRoot.collect(LogicalJoin.class::isInstance)); + boolean found = false; + for (LogicalJoin join : joinList) { + List<Expression> conditions = join.getHashJoinConjuncts(); + List<LogicalCatalogRelation> basicTableList = new ArrayList<>(); + basicTableList.addAll((Collection<? extends LogicalCatalogRelation>) join + .collect(LogicalCatalogRelation.class::isInstance)); + for (Expression equalTo : conditions) { + if (equalTo instanceof EqualTo + && ((EqualTo) equalTo).left() instanceof SlotReference + && ((EqualTo) equalTo).right() instanceof SlotReference) { + SlotReference leftSlot = (SlotReference) ((EqualTo) equalTo).left(); + SlotReference rightSlot = (SlotReference) ((EqualTo) equalTo).right(); + if (table.getOutputExprIds().contains(leftSlot.getExprId()) + && pkSlot.equals(leftSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation rightTable = findTableFromSlot(rightSlot, basicTableList); + if (rightTable != null && getFkInfoFromConstraint(rightTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(rightTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(rightSlot.getName())) { + found = true; + context.replaceColumns.add(rightSlot); + context.pullUpTableToPkSlotMap.put(table, leftSlot); + break; + } + } + } + } + } else if (table.getOutputExprIds().contains(rightSlot.getExprId()) + && pkSlot.equals(rightSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation leftTable = findTableFromSlot(leftSlot, basicTableList); + if (leftTable != null && getFkInfoFromConstraint(leftTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(leftTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(leftSlot.getName())) { + found = true; + context.replaceColumns.add(leftSlot); + context.pullUpTableToPkSlotMap.put(table, rightSlot); + break; + } + } + } + } + } + if (found) { + break; + } + } + } + if (found) { + break; + } + } + return found; + } + + private LogicalCatalogRelation findTableFromSlot(SlotReference targetSlot, + List<LogicalCatalogRelation> tableList) { + for (LogicalCatalogRelation table : tableList) { + if (table.getOutputExprIds().contains(targetSlot.getExprId())) { + return table; + } + } + return null; + } + + private ForeignKeyConstraint getFkInfoFromConstraint(LogicalCatalogRelation table) { + table.getTable().readLock(); Review Comment: do not need lock, since we have locked table before do optimize ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { Review Comment: nit: this function is very similar to `checkJoinCondition`. maybe we could reuse some code ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; Review Comment: nit: could be static ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalAggregate aggrRoot = context.tableToAggrRootMap.get(table); + if (aggrRoot == null) { + return false; + } else if (!checkAggrKeyOnUkOrPk(aggrRoot, table)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkNoFilterOnPullUpTable(List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) { + for (LogicalCatalogRelation table : pullUpTableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else { + List<LogicalFilter> filterList = new ArrayList<>(); + filterList.addAll((Collection<? extends LogicalFilter>) + joinRoot.collect(LogicalFilter.class::isInstance)); + for (LogicalFilter filter : filterList) { + if (filter.child().equals(context.pullUpTable)) { + return false; + } + } + } + } + return true; + } + + private boolean checkAggrKeyOnUkOrPk(LogicalAggregate aggregate, LogicalCatalogRelation table) { + List<Expression> groupByKeys = aggregate.getGroupByExpressions(); + boolean isAllSlotReference = groupByKeys.stream().allMatch(e -> e instanceof SlotReference); + if (!isAllSlotReference) { + return false; + } else { + Set<String> ukInfo = getUkInfoFromConstraint(table); + Set<String> pkInfo = getPkInfoFromConstraint(table); + if (ukInfo == null || pkInfo == null || ukInfo.size() != 1 || pkInfo.size() != 1) { + return false; + } else { + String ukName = ukInfo.iterator().next(); + String pkName = pkInfo.iterator().next(); + for (Object expr : aggregate.getGroupByExpressions()) { + SlotReference slot = (SlotReference) expr; + if (table.getOutputExprIds().contains(slot.getExprId()) + && (slot.getName().equals(ukName) || slot.getName().equals(pkName))) { + return true; + } + } + return false; + } + } + } + + private boolean checkJoinConditionOnPk(LogicalJoin joinRoot, LogicalCatalogRelation table, PullUpContext context) { + Set<String> pkInfos = getPkInfoFromConstraint(table); + if (pkInfos == null || pkInfos.size() != 1) { + return false; + } + String pkSlot = pkInfos.iterator().next(); + List<LogicalJoin> joinList = new ArrayList<>(); + joinList.addAll((Collection<? extends LogicalJoin>) joinRoot.collect(LogicalJoin.class::isInstance)); + boolean found = false; + for (LogicalJoin join : joinList) { + List<Expression> conditions = join.getHashJoinConjuncts(); + List<LogicalCatalogRelation> basicTableList = new ArrayList<>(); + basicTableList.addAll((Collection<? extends LogicalCatalogRelation>) join + .collect(LogicalCatalogRelation.class::isInstance)); + for (Expression equalTo : conditions) { + if (equalTo instanceof EqualTo + && ((EqualTo) equalTo).left() instanceof SlotReference + && ((EqualTo) equalTo).right() instanceof SlotReference) { + SlotReference leftSlot = (SlotReference) ((EqualTo) equalTo).left(); + SlotReference rightSlot = (SlotReference) ((EqualTo) equalTo).right(); + if (table.getOutputExprIds().contains(leftSlot.getExprId()) + && pkSlot.equals(leftSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation rightTable = findTableFromSlot(rightSlot, basicTableList); + if (rightTable != null && getFkInfoFromConstraint(rightTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(rightTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(rightSlot.getName())) { + found = true; + context.replaceColumns.add(rightSlot); + context.pullUpTableToPkSlotMap.put(table, leftSlot); + break; + } + } + } + } + } else if (table.getOutputExprIds().contains(rightSlot.getExprId()) + && pkSlot.equals(rightSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation leftTable = findTableFromSlot(leftSlot, basicTableList); + if (leftTable != null && getFkInfoFromConstraint(leftTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(leftTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(leftSlot.getName())) { + found = true; + context.replaceColumns.add(leftSlot); + context.pullUpTableToPkSlotMap.put(table, rightSlot); + break; + } + } + } + } + } + if (found) { + break; + } + } + } + if (found) { + break; + } + } + return found; + } + + private LogicalCatalogRelation findTableFromSlot(SlotReference targetSlot, + List<LogicalCatalogRelation> tableList) { + for (LogicalCatalogRelation table : tableList) { + if (table.getOutputExprIds().contains(targetSlot.getExprId())) { + return table; + } + } + return null; + } + + private ForeignKeyConstraint getFkInfoFromConstraint(LogicalCatalogRelation table) { + table.getTable().readLock(); + try { + for (Map.Entry<String, Constraint> constraintMap : table.getTable().getConstraintsMap().entrySet()) { + Constraint constraint = constraintMap.getValue(); + if (constraint instanceof ForeignKeyConstraint) { + return (ForeignKeyConstraint) constraint; + } + } + return null; + } finally { + table.getTable().readUnlock(); + } + } + + private Set<String> getPkInfoFromConstraint(LogicalCatalogRelation table) { + table.getTable().readLock(); + try { + for (Map.Entry<String, Constraint> constraintMap : table.getTable().getConstraintsMap().entrySet()) { + Constraint constraint = constraintMap.getValue(); + if (constraint instanceof PrimaryKeyConstraint) { + return ((PrimaryKeyConstraint) constraint).getPrimaryKeyNames(); + } + } + return null; + } finally { + table.getTable().readUnlock(); + } + } + + private Set<String> getUkInfoFromConstraint(LogicalCatalogRelation table) { + table.getTable().readLock(); + try { + for (Map.Entry<String, Constraint> constraintMap : table.getTable().getConstraintsMap().entrySet()) { + Constraint constraint = constraintMap.getValue(); + if (constraint instanceof UniqueConstraint) { + return ((UniqueConstraint) constraint).getUniqueColumnNames(); + } + } + return null; + } finally { + table.getTable().readUnlock(); + } + } + + private boolean checkJoinRoot(LogicalJoin joinRoot) { + List<LogicalPlan> joinChildrenPlans = Lists.newArrayList(); + joinChildrenPlans.addAll((Collection<? extends LogicalPlan>) joinRoot + .collect(LogicalPlan.class::isInstance)); + boolean planTypeMatch = joinChildrenPlans.stream() + .allMatch(p -> SUPPORTED_PLAN_TYPE.stream().anyMatch(c -> c.isInstance(p))); + if (!planTypeMatch) { + return false; + } + + List<LogicalJoin> allJoinNodes = Lists.newArrayList(); + allJoinNodes.addAll((Collection<? extends LogicalJoin>) joinRoot.collect(LogicalJoin.class::isInstance)); + boolean joinTypeMatch = allJoinNodes.stream().allMatch(e -> e.getJoinType() == JoinType.INNER_JOIN); + boolean joinConditionMatch = allJoinNodes.stream() + .allMatch(e -> !e.getHashJoinConjuncts().isEmpty() && e.getOtherJoinConjuncts().isEmpty()); + if (!joinTypeMatch || !joinConditionMatch) { + return false; + } + + return true; + } + + private boolean checkAggrRoot(LogicalAggregate aggrRoot) { + for (Object expr : aggrRoot.getGroupByExpressions()) { + if (!(expr instanceof NamedExpression)) { + return false; + } + } + return true; + } + + private List<LogicalCatalogRelation> getTableListUnderJoin(LogicalJoin joinRoot) { + List<LogicalCatalogRelation> tableLists = new ArrayList<>(); + tableLists.addAll((Collection<? extends LogicalCatalogRelation>) joinRoot + .collect(LogicalCatalogRelation.class::isInstance)); + return tableLists; + } + + private String makeQualifiedName(LogicalCatalogRelation table) { + String dbName = table.getTable().getDatabase().getFullName(); + String tableName = table.getTable().getName(); + return dbName + ":" + tableName; Review Comment: why use `:`? ########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PullUpJoinFromUnionAll.java: ########## @@ -0,0 +1,688 @@ +// 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.rules.rewrite; + +import org.apache.doris.catalog.constraint.Constraint; +import org.apache.doris.catalog.constraint.ForeignKeyConstraint; +import org.apache.doris.catalog.constraint.PrimaryKeyConstraint; +import org.apache.doris.catalog.constraint.UniqueConstraint; +import org.apache.doris.nereids.jobs.JobContext; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.ExprId; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.plans.JoinHint; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter; +import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * Pull up join from union all rule. + */ +public class PullUpJoinFromUnionAll extends OneRewriteRuleFactory { + private static final Set<Class<? extends LogicalPlan>> SUPPORTED_PLAN_TYPE = ImmutableSet.of( + LogicalFilter.class, + LogicalJoin.class, + LogicalProject.class, + LogicalCatalogRelation.class + ); + + private static class PullUpContext { + public final String unifiedOutputAlias = "PULL_UP_UNIFIED_OUTPUT_ALIAS"; + public final Map<String, List<LogicalCatalogRelation>> pullUpCandidatesMaps = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalJoin> tableToJoinRootMap = Maps.newHashMap(); + public final Map<LogicalCatalogRelation, LogicalAggregate> tableToAggrRootMap = Maps.newHashMap(); + public final Map<NamedExpression, SlotReference> origChild0ToNewUnionOutputMap = Maps.newHashMap(); + public final List<LogicalAggregate> aggrChildList = Lists.newArrayList(); + public final List<LogicalJoin> joinChildList = Lists.newArrayList(); + public final List<SlotReference> replaceColumns = Lists.newArrayList(); + public final Map<LogicalCatalogRelation, Slot> pullUpTableToPkSlotMap = Maps.newHashMap(); + public int replacedColumnIndex = -1; + public LogicalCatalogRelation pullUpTable; + + // the slot will replace the original pk in group by and select list + public SlotReference replaceColumn; + public boolean needAddReplaceColumn = false; + + public PullUpContext() {} + + public void setReplacedColumn(SlotReference slot) { + this.replaceColumn = slot; + } + + public void setPullUpTable(LogicalCatalogRelation table) { + this.pullUpTable = table; + } + + public void setNeedAddReplaceColumn(boolean needAdd) { + this.needAddReplaceColumn = needAdd; + } + } + + @Override + public Rule build() { + return logicalUnion() + .when(union -> union.getQualifier() != Qualifier.DISTINCT) + .then(union -> { + PullUpContext context = new PullUpContext(); + if (!checkUnionPattern(union, context) + || !checkJoinCondition(context) + || !checkGroupByKeys(context)) { + return null; + } + // only support single table pull up currently + if (context.pullUpCandidatesMaps.entrySet().size() != 1) { + return null; + } + + List<LogicalCatalogRelation> pullUpTableList = context.pullUpCandidatesMaps + .entrySet().iterator().next().getValue(); + if (pullUpTableList.size() != union.children().size() + || context.replaceColumns.size() != union.children().size() + || !checkNoFilterOnPullUpTable(pullUpTableList, context)) { + return null; + } + // make new union node + LogicalUnion newUnionNode = makeNewUnionNode(union, pullUpTableList, context); + // make new join node + LogicalJoin newJoin = makeNewJoin(newUnionNode, pullUpTableList.get(0), context); + // add project on pull up table with origin union output + List<NamedExpression> newProjectOutputs = makeNewProjectOutputs(union, newJoin, context); + + return new LogicalProject(newProjectOutputs, newJoin); + }).toRule(RuleType.PULL_UP_JOIN_FROM_UNIONALL); + } + + private boolean checkUnionPattern(LogicalUnion union, PullUpContext context) { + int tableListNumber = -1; + for (Plan child : union.children()) { + if (!(child instanceof LogicalProject + && child.child(0) != null + && child.child(0) instanceof LogicalAggregate + && child.child(0).child(0) != null + && child.child(0).child(0) instanceof LogicalProject + && child.child(0).child(0).child(0) != null + && child.child(0).child(0).child(0) instanceof LogicalJoin)) { + return false; + } + LogicalAggregate aggrRoot = (LogicalAggregate) child.child(0); + if (!checkAggrRoot(aggrRoot)) { + return false; + } + context.aggrChildList.add(aggrRoot); + LogicalJoin joinRoot = (LogicalJoin) aggrRoot.child().child(0); + // check join under union is spj + if (!checkJoinRoot(joinRoot)) { + return false; + } + context.joinChildList.add(joinRoot); + + List<LogicalCatalogRelation> tableList = getTableListUnderJoin(joinRoot); + // add into table -> joinRoot map + for (LogicalCatalogRelation table : tableList) { + context.tableToJoinRootMap.put(table, joinRoot); + context.tableToAggrRootMap.put(table, aggrRoot); + } + if (tableListNumber == -1) { + tableListNumber = tableList.size(); + } else { + // check all union children have the same number of tables + if (tableListNumber != tableList.size()) { + return false; + } + } + + for (LogicalCatalogRelation table : tableList) { + // key: qualified table name + // value: table list in all union children + String qName = makeQualifiedName(table); + if (context.pullUpCandidatesMaps.get(qName) == null) { + List<LogicalCatalogRelation> newList = new ArrayList<>(); + newList.add(table); + context.pullUpCandidatesMaps.put(qName, newList); + } else { + context.pullUpCandidatesMaps.get(qName).add(table); + } + } + } + int expectedNumber = union.children().size(); + List<String> toBeRemoved = new ArrayList<>(); + // check the pull up table candidate exists in all union children + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + if (e.getValue().size() != expectedNumber) { + toBeRemoved.add(e.getKey()); + } + } + for (String key : toBeRemoved) { + context.pullUpCandidatesMaps.remove(key); + } + return !context.pullUpCandidatesMaps.isEmpty(); + } + + private boolean checkJoinCondition(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else if (!checkJoinConditionOnPk(joinRoot, table, context)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkGroupByKeys(PullUpContext context) { + List<String> toBeRemoved = new ArrayList<>(); + for (Map.Entry<String, List<LogicalCatalogRelation>> e : context.pullUpCandidatesMaps.entrySet()) { + List<LogicalCatalogRelation> tableList = e.getValue(); + boolean allFound = true; + for (LogicalCatalogRelation table : tableList) { + LogicalAggregate aggrRoot = context.tableToAggrRootMap.get(table); + if (aggrRoot == null) { + return false; + } else if (!checkAggrKeyOnUkOrPk(aggrRoot, table)) { + allFound = false; + break; + } + } + if (!allFound) { + toBeRemoved.add(e.getKey()); + } + } + for (String table : toBeRemoved) { + context.pullUpCandidatesMaps.remove(table); + } + + if (context.pullUpCandidatesMaps.isEmpty()) { + return false; + } + return true; + } + + private boolean checkNoFilterOnPullUpTable(List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) { + for (LogicalCatalogRelation table : pullUpTableList) { + LogicalJoin joinRoot = context.tableToJoinRootMap.get(table); + if (joinRoot == null) { + return false; + } else { + List<LogicalFilter> filterList = new ArrayList<>(); + filterList.addAll((Collection<? extends LogicalFilter>) + joinRoot.collect(LogicalFilter.class::isInstance)); + for (LogicalFilter filter : filterList) { + if (filter.child().equals(context.pullUpTable)) { + return false; + } + } + } + } + return true; + } + + private boolean checkAggrKeyOnUkOrPk(LogicalAggregate aggregate, LogicalCatalogRelation table) { + List<Expression> groupByKeys = aggregate.getGroupByExpressions(); + boolean isAllSlotReference = groupByKeys.stream().allMatch(e -> e instanceof SlotReference); + if (!isAllSlotReference) { + return false; + } else { + Set<String> ukInfo = getUkInfoFromConstraint(table); + Set<String> pkInfo = getPkInfoFromConstraint(table); + if (ukInfo == null || pkInfo == null || ukInfo.size() != 1 || pkInfo.size() != 1) { + return false; + } else { + String ukName = ukInfo.iterator().next(); + String pkName = pkInfo.iterator().next(); + for (Object expr : aggregate.getGroupByExpressions()) { + SlotReference slot = (SlotReference) expr; + if (table.getOutputExprIds().contains(slot.getExprId()) + && (slot.getName().equals(ukName) || slot.getName().equals(pkName))) { + return true; + } + } + return false; + } + } + } + + private boolean checkJoinConditionOnPk(LogicalJoin joinRoot, LogicalCatalogRelation table, PullUpContext context) { + Set<String> pkInfos = getPkInfoFromConstraint(table); + if (pkInfos == null || pkInfos.size() != 1) { + return false; + } + String pkSlot = pkInfos.iterator().next(); + List<LogicalJoin> joinList = new ArrayList<>(); + joinList.addAll((Collection<? extends LogicalJoin>) joinRoot.collect(LogicalJoin.class::isInstance)); + boolean found = false; + for (LogicalJoin join : joinList) { + List<Expression> conditions = join.getHashJoinConjuncts(); + List<LogicalCatalogRelation> basicTableList = new ArrayList<>(); + basicTableList.addAll((Collection<? extends LogicalCatalogRelation>) join + .collect(LogicalCatalogRelation.class::isInstance)); + for (Expression equalTo : conditions) { + if (equalTo instanceof EqualTo + && ((EqualTo) equalTo).left() instanceof SlotReference + && ((EqualTo) equalTo).right() instanceof SlotReference) { + SlotReference leftSlot = (SlotReference) ((EqualTo) equalTo).left(); + SlotReference rightSlot = (SlotReference) ((EqualTo) equalTo).right(); + if (table.getOutputExprIds().contains(leftSlot.getExprId()) + && pkSlot.equals(leftSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation rightTable = findTableFromSlot(rightSlot, basicTableList); + if (rightTable != null && getFkInfoFromConstraint(rightTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(rightTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(rightSlot.getName())) { + found = true; + context.replaceColumns.add(rightSlot); + context.pullUpTableToPkSlotMap.put(table, leftSlot); + break; + } + } + } + } + } else if (table.getOutputExprIds().contains(rightSlot.getExprId()) + && pkSlot.equals(rightSlot.getName())) { + // pk-fk join condition, check other side's join key is on fk + LogicalCatalogRelation leftTable = findTableFromSlot(leftSlot, basicTableList); + if (leftTable != null && getFkInfoFromConstraint(leftTable) != null) { + ForeignKeyConstraint fkInfo = getFkInfoFromConstraint(leftTable); + if (fkInfo.getReferencedTable().getId() == table.getTable().getId()) { + for (Map.Entry<String, String> entry : fkInfo.getForeignToReference().entrySet()) { + if (entry.getValue().equals(pkSlot) && entry.getKey().equals(leftSlot.getName())) { + found = true; + context.replaceColumns.add(leftSlot); + context.pullUpTableToPkSlotMap.put(table, rightSlot); + break; + } + } + } + } + } + if (found) { + break; + } + } + } + if (found) { + break; + } + } + return found; + } + + private LogicalCatalogRelation findTableFromSlot(SlotReference targetSlot, Review Comment: we could use SlotReference#table attr for check more accurately -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org For additional commands, e-mail: commits-h...@doris.apache.org
