xzj7019 commented on code in PR #45958: URL: https://github.com/apache/doris/pull/45958#discussion_r1901622903
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownEncodeSlot.java: ########## @@ -0,0 +1,660 @@ +// 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.common.Pair; +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.ComparisonPredicate; +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.functions.scalar.DecodeAsVarchar; +import org.apache.doris.nereids.trees.expressions.functions.scalar.EncodeString; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalLeaf; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalRepeat; +import org.apache.doris.nereids.trees.plans.logical.LogicalSetOperation; +import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.nereids.types.coercion.CharacterType; +import org.apache.doris.qe.ConnectContext; + +import com.google.common.base.Preconditions; +import com.google.common.collect.BiMap; +import com.google.common.collect.HashBiMap; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Lists; +import com.google.common.collect.Sets; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * push down encode_as_int(slot) down + * example: + * group by x + * -->project(encode_as_int(A) as x) + * -->Any(A) + * -->project(A) + * --> scan + * => + * group by x + * -->project(x) + * -->Any(x) + * --> project(encode_as_int(A) as x) + * -->scan + * Note: + * do not push down encode if encode.child() is not slot, + * example + * group by encode_as_int(A + B) + * --> any(A, B) + */ +public class PushDownEncodeSlot extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalProject() + .when(topN -> ConnectContext.get() != null + && ConnectContext.get().getSessionVariable().enableCompressMaterialize) + .whenNot(project -> project.child() instanceof LogicalRepeat) + .whenNot(project -> (project.child() instanceof LogicalLeaf)) + .then(this::pushDownEncodeSlot) + .toRule(RuleType.PUSH_DOWN_ENCODE_SLOT); + } + + private List<Alias> collectEncodeAliases(LogicalProject<? extends Plan> project) { + List<Alias> encodeAliases = new ArrayList<>(); + Set<Slot> computingSlots = new HashSet<>(); + for (NamedExpression e : project.getProjects()) { + if (e instanceof Alias) { + Expression aliasBody = e.child(0); + if (!(aliasBody instanceof SlotReference) && !(aliasBody instanceof EncodeString)) { + computingSlots.addAll(e.getInputSlots()); + } + } + } + for (NamedExpression e : project.getProjects()) { + if (e instanceof Alias && e.child(0) instanceof EncodeString + && e.child(0).child(0) instanceof SlotReference + && !computingSlots.contains(e.child(0).child(0))) { + encodeAliases.add((Alias) e); + } + } + return encodeAliases; + } + + /** + * case 1 + * project(encode(A) as B) + * --> any(A) + * => + * project(B) + * -->any(A): push "encode(A) as B" + * + * case 2 + * project(A, encode(A) as B) + * -->any(A) + * => + * project(decode(B) as A, B) + * -->any(A): push "encode(A) as B" + * + * case 3 + * project(A as C, encode(A) as B) + * -->any(A) + * => + * project(decode(B) as C, B) + * -->any(A): push "encode(A) as B" + */ + private LogicalProject<? extends Plan> rewriteRootProject(LogicalProject<? extends Plan> project, + List<Alias> pushedEncodeAlias) { + if (pushedEncodeAlias.isEmpty()) { + return project; + } + Map<Expression, Alias> encodeBodyToEncodeAlias = new HashMap<>(); + for (Alias alias : pushedEncodeAlias) { + Expression encodeBody = alias.child().child(0); + encodeBodyToEncodeAlias.put(encodeBody, alias); + } + List<NamedExpression> projections = Lists.newArrayListWithCapacity(project.getProjects().size()); + for (NamedExpression e : project.getProjects()) { + if (pushedEncodeAlias.contains(e)) { + // case 1 + projections.add(e.toSlot()); + } else if (encodeBodyToEncodeAlias.containsKey(e)) { + // case 2 + ExprId id = e.getExprId(); + DecodeAsVarchar decode = new DecodeAsVarchar(encodeBodyToEncodeAlias.get(e).toSlot()); + Alias alias = new Alias(id, decode, decode.toSql()); + projections.add(alias); + } else if (e instanceof Alias && encodeBodyToEncodeAlias.containsKey(e.child(0))) { + // case 3 + Alias alias = (Alias) e; + DecodeAsVarchar decode = new DecodeAsVarchar(encodeBodyToEncodeAlias.get(e.child(0)).toSlot()); + Alias newAlias = (Alias) alias.withChildren(decode); + projections.add(newAlias); + } else { + projections.add(e); + } + } + return project.withProjects(projections); + + } + + private LogicalProject<? extends Plan> pushDownEncodeSlot(LogicalProject<? extends Plan> project) { + List<Alias> encodeAliases = collectEncodeAliases(project); + if (encodeAliases.isEmpty()) { + return project; + } + + PushDownContext ctx = new PushDownContext(project, encodeAliases); + ctx.prepare(); + if (ctx.notPushed.size() == encodeAliases.size()) { + return project; + } + Plan child = project.child(); + PushDownContext childContext = new PushDownContext(child, ctx.toBePushedToChild.get(child)); + Plan newChild = child.accept(EncodeSlotPushDownVisitor.INSTANCE, childContext); + List<Alias> pushed = ctx.toBePused; + if (child != newChild) { + if (newChild instanceof LogicalProject) { + pushed.removeAll(childContext.notPushed); + newChild = ((LogicalProject<?>) newChild).child(); + } + project = (LogicalProject<? extends Plan>) project.withChildren(newChild); + project = rewriteRootProject(project, pushed); + } + return project; + } + + /** + * push down encode slot context + */ + public static class PushDownContext { + public Plan plan; + + public List<Alias> encodeAliases; + // encode_as_int(slot1) as slot2 + // replaceMap: + // slot1 -> slot2 + Map<Expression, SlotReference> replaceMap = new HashMap<>(); + // child plan -> aliases in encodeAliases which can be pushed down to child plan + Map<Plan, List<Alias>> toBePushedToChild = new HashMap<>(); + List<Alias> toBePused = new ArrayList<>(); + // the aliases that cannot be pushed down to any child plan + // for example: + // encode(A+B) as x, where plan is a join, and A, B comes from join's left and right child respectively + List<Alias> notPushed = new ArrayList<>(); + + public PushDownContext(Plan plan, List<Alias> encodeAliases) { + this.plan = plan; + this.encodeAliases = encodeAliases; + } + + // init replaceMap/toBePushed/notPushed + private void prepare() { + List<Set<Slot>> childrenPassThroughSlots = + plan.children().stream().map(n -> getPassThroughSlots(n)).collect(Collectors.toList()); + for (int i = 0; i < plan.children().size(); i++) { + Plan child = plan.children().get(i); + if (child instanceof LogicalJoin) { + LogicalJoin<?, ?> join = (LogicalJoin<?, ?>) child; + BiMap<SlotReference, SlotReference> compareSlots = EncodeSlotPushDownVisitor + .getEncodeCandidateSlotsAndShouldNotPushSlotsFromJoinCondition(join).first; + childrenPassThroughSlots.get(i).addAll(compareSlots.keySet()); + } + } + for (Alias alias : encodeAliases) { + EncodeString encode = (EncodeString) alias.child(); + Expression strExpr = encode.child(); + boolean pushed = false; + Preconditions.checkArgument(strExpr instanceof SlotReference, + "expect encode_as_xxx(slot), but " + alias); + + for (int i = 0; i < childrenPassThroughSlots.size(); i++) { + if (childrenPassThroughSlots.get(i).contains(strExpr)) { + toBePushedToChild.putIfAbsent(plan.child(i), new ArrayList<>()); + toBePushedToChild.get(plan.child(i)).add(alias); + toBePused.add(alias); + replaceMap.put(alias.child().child(0), (SlotReference) alias.toSlot()); + pushed = true; + break; + } + } + if (!pushed) { + notPushed.add(alias); + } + } + } + + /** + * expandEncodeAliasForJoin + */ + public void expandEncodeAliasForJoin(BiMap<SlotReference, SlotReference> equalSlots) { + List<Alias> expanded = new ArrayList<>(); + for (Alias alias : encodeAliases) { + if (alias.child().child(0) instanceof SlotReference) { + SlotReference slot = (SlotReference) alias.child().child(0); + SlotReference otherHand = equalSlots.get(slot); + if (otherHand != null) { + EncodeString encodeOtherHand = (EncodeString) alias.child().withChildren(otherHand); + Alias encodeOtherHandAlias = new Alias(encodeOtherHand, encodeOtherHand.toSql()); + if (!encodeAliases.contains(encodeOtherHandAlias)) { + expanded.add(encodeOtherHandAlias); + } + } + } + } + encodeAliases.addAll(expanded); + } + + // the child of alias is a slot reference. for example: slotA as B + // + private boolean isSlotAlias(Expression expr) { + return expr instanceof Alias && expr.child(0) instanceof SlotReference; + } + + private Set<Slot> getPassThroughSlots(Plan plan) { + Set<Slot> outputSlots = Sets.newHashSet(plan.getOutputSet()); + Set<Slot> keySlots = Sets.newHashSet(); + for (Expression e : plan.getExpressions()) { + if (!(e instanceof SlotReference) && !isSlotAlias(e)) { + keySlots.addAll(e.getInputSlots()); + } + } + outputSlots.removeAll(keySlots); + return outputSlots; + } + } + + /** + * push down encode slot + */ + public static class EncodeSlotPushDownVisitor extends PlanVisitor<Plan, PushDownContext> { Review Comment: how about window function visitor? is there any necessary to handle? -- 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
