xzj7019 commented on code in PR #28682: URL: https://github.com/apache/doris/pull/28682#discussion_r1434681278
########## 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: done ########## 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: done -- 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
