// 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.ForeignKeyConstraint;
import org.apache.doris.catalog.constraint.PrimaryKeyConstraint;
import org.apache.doris.catalog.constraint.UniqueConstraint;
import org.apache.doris.nereids.hint.DistributeHint;
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.DistributeType;
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 static 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) instanceof LogicalAggregate
                    && child.child(0).child(0) instanceof LogicalProject
                    && 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) {
        Set<ForeignKeyConstraint> foreignKeyConstraints = table.getTable().getForeignKeyConstraints();
        if (foreignKeyConstraints.isEmpty()) {
            return null;
        }
        return foreignKeyConstraints.stream().iterator().next();
    }

    private Set<String> getPkInfoFromConstraint(LogicalCatalogRelation table) {
        Set<PrimaryKeyConstraint> primaryKeyConstraints = table.getTable().getPrimaryKeyConstraints();
        if (primaryKeyConstraints.isEmpty()) {
            return null;
        }
        return primaryKeyConstraints.stream().iterator().next().getPrimaryKeyNames();
    }

    private Set<String> getUkInfoFromConstraint(LogicalCatalogRelation table) {
        Set<UniqueConstraint> uniqueConstraints = table.getTable().getUniqueConstraints();
        if (uniqueConstraints.isEmpty()) {
            return null;
        }
        return uniqueConstraints.stream().iterator().next().getUniqueColumnNames();
    }

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

    private Plan doPullUpJoinFromUnionAll(Plan unionChildPlan, PullUpContext context) {
        return PullUpRewriter.INSTANCE.rewrite(unionChildPlan, context);
    }

    private List<Plan> doWrapReplaceColumnForUnionChildren(List<Plan> unionChildren, PullUpContext context) {
        List<Plan> newUnionChildren = new ArrayList<>();
        for (int i = 0; i < unionChildren.size(); i++) {
            // has been checked before
            LogicalProject oldProject = (LogicalProject) unionChildren.get(i);
            List<NamedExpression> newNamedExpressionList = new ArrayList<>();
            for (int j = 0; j < oldProject.getProjects().size(); j++) {
                Object child = oldProject.getProjects().get(j);
                if (context.replaceColumns.contains(child)) {
                    Alias newExpr = new Alias((Expression) child, context.unifiedOutputAlias);
                    newNamedExpressionList.add(newExpr);
                    context.replacedColumnIndex = j;
                } else {
                    newNamedExpressionList.add((NamedExpression) child);
                }
            }
            LogicalProject newProject = new LogicalProject(newNamedExpressionList, (LogicalPlan) oldProject.child());
            newUnionChildren.add(newProject);
        }
        return newUnionChildren;
    }

    private List<NamedExpression> makeNewProjectOutputs(LogicalUnion origUnion,
            LogicalJoin newJoin, PullUpContext context) {
        List<NamedExpression> newProjectOutputs = new ArrayList<>();
        List<Slot> origUnionSlots = origUnion.getOutput();
        List<NamedExpression> origUnionChildOutput = ((LogicalProject) origUnion.child(0)).getOutputs();
        for (int i = 0; i < origUnionChildOutput.size(); i++) {
            NamedExpression unionOutputExpr = origUnionChildOutput.get(i);
            if (unionOutputExpr instanceof Alias) {
                if (!(unionOutputExpr.child(0) instanceof Literal)) {
                    unionOutputExpr = (Slot) unionOutputExpr.child(0);
                }
            }
            boolean found = false;
            Slot matchedJoinSlot = null;
            for (Slot joinOutput : newJoin.getOutput()) {
                Slot slot = joinOutput;
                if (context.origChild0ToNewUnionOutputMap.get(slot) != null) {
                    slot = context.origChild0ToNewUnionOutputMap.get(slot);
                }
                if (slot.equals(unionOutputExpr) || slot.getExprId() == unionOutputExpr.getExprId()) {
                    matchedJoinSlot = joinOutput;
                    found = true;
                    break;
                }
            }
            if (found) {
                ExprId exprId = origUnionSlots.get(i).getExprId();
                Alias aliasExpr = new Alias(exprId, matchedJoinSlot, matchedJoinSlot.toSql());
                newProjectOutputs.add(aliasExpr);
            }
        }
        return newProjectOutputs;
    }

    private LogicalJoin makeNewJoin(LogicalUnion newUnionNode,
            LogicalCatalogRelation pullUpTable, PullUpContext context) {
        List<Expression> newHashJoinConjuncts = new ArrayList<>();
        Slot unionSideExpr = newUnionNode.getOutput().get(context.replacedColumnIndex);

        Slot pullUpSidePkSlot = context.pullUpTableToPkSlotMap.get(pullUpTable);
        if (pullUpSidePkSlot == null) {
            return null;
        }
        EqualTo pullUpJoinCondition = new EqualTo(unionSideExpr, pullUpSidePkSlot);
        newHashJoinConjuncts.add(pullUpJoinCondition);

        // new a join with the newUnion and the pulled up table
        return new LogicalJoin<>(
                JoinType.INNER_JOIN,
                newHashJoinConjuncts,
                ExpressionUtils.EMPTY_CONDITION,
                new DistributeHint(DistributeType.NONE),
                Optional.empty(),
                newUnionNode,
                pullUpTable, null);
    }

    private LogicalUnion makeNewUnionNode(LogicalUnion origUnion,
            List<LogicalCatalogRelation> pullUpTableList, PullUpContext context) {
        List<Plan> newUnionChildren = new ArrayList<>();
        for (int i = 0; i < origUnion.children().size(); i++) {
            Plan unionChild = origUnion.child(i);
            context.setPullUpTable(pullUpTableList.get(i));
            context.setReplacedColumn(context.replaceColumns.get(i));
            Plan newChild = doPullUpJoinFromUnionAll(unionChild, context);
            newUnionChildren.add(newChild);
        }

        // wrap the replaced column with a shared alias which is exposed to outside
        List<Plan> formalizedNewUnionChildren = doWrapReplaceColumnForUnionChildren(newUnionChildren, context);
        List<List<SlotReference>> childrenOutputs = formalizedNewUnionChildren.stream()
                .map(j -> j.getOutput().stream()
                        .map(SlotReference.class::cast)
                        .collect(ImmutableList.toImmutableList()))
                .collect(ImmutableList.toImmutableList());

        LogicalUnion newUnionNode = new LogicalUnion(Qualifier.ALL, formalizedNewUnionChildren);
        newUnionNode = (LogicalUnion) newUnionNode.withChildrenAndTheirOutputs(
                formalizedNewUnionChildren, childrenOutputs);
        List<NamedExpression> newOutputs = newUnionNode.buildNewOutputs();
        newUnionNode = newUnionNode.withNewOutputs(newOutputs);

        // set up origin child 0 output to new union output mapping
        List<SlotReference> origChild0Output = childrenOutputs.get(0);
        for (int i = 0; i < origChild0Output.size(); i++) {
            SlotReference slot = origChild0Output.get(i);
            NamedExpression newExpr = newOutputs.get(i);
            context.origChild0ToNewUnionOutputMap.put(newExpr, slot);
        }

        return newUnionNode;
    }

    private static class PullUpRewriter extends DefaultPlanRewriter<PullUpContext> implements CustomRewriter {
        public static final PullUpRewriter INSTANCE = new PullUpRewriter();

        @Override
        public Plan rewriteRoot(Plan plan, JobContext context) {
            return null;
        }

        public Plan rewrite(Plan plan, PullUpContext context) {
            return plan.accept(this, context);
        }

        @Override
        public Plan visitLogicalAggregate(LogicalAggregate<? extends Plan> agg, PullUpContext context) {
            Plan input = agg.child().accept(this, context);

            LogicalCatalogRelation pullTable = context.pullUpTable;
            SlotReference replaceColumn = context.replaceColumn;

            // eliminate group by keys
            List<Expression> groupByExprList = new ArrayList<>();
            for (Expression expr : agg.getGroupByExpressions()) {
                // expr has been checked before
                if (!pullTable.getOutputExprIds().contains(((NamedExpression) expr).getExprId())) {
                    groupByExprList.add(expr);
                }
            }
            // add replaced group by key
            groupByExprList.add(replaceColumn);

            // eliminate outputs keys
            List<NamedExpression> outputExprList = new ArrayList<>();
            for (NamedExpression expr : agg.getOutputExpressions()) {
                if (!pullTable.getOutputExprIds().contains(expr.getExprId())) {
                    outputExprList.add(expr);
                }
            }
            // add replaced group by key
            outputExprList.add(replaceColumn);

            return new LogicalAggregate<>(groupByExprList, outputExprList, input);
        }

        public Plan visitLogicalJoin(LogicalJoin<? extends Plan, ? extends Plan> join, PullUpContext context) {
            Plan leftChild = join.child(0).accept(this, context);
            Plan rightChild = join.child(1).accept(this, context);
            LogicalCatalogRelation pullUpTable = context.pullUpTable;

            // no filter on pull up table, which has been checked before
            if (leftChild instanceof LogicalCatalogRelation
                    && leftChild.equals(pullUpTable)) {
                context.setNeedAddReplaceColumn(true);
                return rightChild;
            } else if (rightChild instanceof LogicalCatalogRelation
                    && rightChild.equals(pullUpTable)) {
                context.setNeedAddReplaceColumn(true);
                return leftChild;
            } else if (leftChild instanceof LogicalProject
                    && leftChild.child(0) instanceof LogicalCatalogRelation
                    && leftChild.child(0).equals(pullUpTable)) {
                context.setNeedAddReplaceColumn(true);
                return rightChild;
            } else if (rightChild instanceof LogicalProject
                    && rightChild.child(0) instanceof LogicalCatalogRelation
                    && rightChild.child(0).equals(pullUpTable)) {
                context.setNeedAddReplaceColumn(true);
                return leftChild;
            } else {
                return new LogicalJoin(JoinType.INNER_JOIN,
                        join.getHashJoinConjuncts(),
                        join.getOtherJoinConjuncts(),
                        new DistributeHint(DistributeType.NONE),
                        Optional.empty(),
                        leftChild, rightChild, null);
            }
        }

        @Override
        public Plan visitLogicalProject(LogicalProject<? extends Plan> project, PullUpContext context) {
            Plan input = project.child().accept(this, context);
            List<NamedExpression> outputs = input.getOutput().stream()
                    .map(e -> (NamedExpression) e).collect(Collectors.toList());
            for (NamedExpression expr : project.getProjects()) {
                // handle alias
                if (expr instanceof Alias && expr.child(0) instanceof Literal) {
                    outputs.add(expr);
                }
            }
            return new LogicalProject<>(outputs, input);
        }

        @Override
        public Plan visitLogicalFilter(LogicalFilter<? extends Plan> filter, PullUpContext context) {
            Plan input = filter.child().accept(this, context);
            return new LogicalFilter<>(filter.getConjuncts(), input);
        }
    }
}