XieJiann commented on code in PR #28682:
URL: https://github.com/apache/doris/pull/28682#discussion_r1434651090

##########
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;

Review Comment:
   map agg.getGroupExpressions to Slot and use 
`getLogicalProperty.getFunctionalDependence.isUnique()`



-- 
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

Reply via email to