linrrzqqq commented on code in PR #66968: URL: https://github.com/apache/doris/pull/66968#discussion_r3829820828
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AddProjectForMapLambdaInput.java: ########## @@ -0,0 +1,786 @@ +// 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.exceptions.AnalysisException; +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.ArrayItemReference; +import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot; +import org.apache.doris.nereids.trees.expressions.Cast; +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.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.expressions.functions.Function; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapEntryArrayMap; +import org.apache.doris.nereids.trees.expressions.functions.scalar.MapLambdaValidator; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalGenerate; +import org.apache.doris.nereids.trees.plans.logical.LogicalHaving; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.JoinUtils; + +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 com.google.common.collect.Sets; + +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Optional; +import java.util.Set; + +/** + * Materialize computed Map inputs used by {@link MapEntryArrayMap}. + * + * <p>A Map entry lambda takes {@code map_keys(computedMap)} and + * {@code map_values(computedMap)} as its two input arrays. rule evaThisluates + * {@code computedMap} in a child Project and replaces all its occurrences with the same Slot: + * + * <pre> + * before: + * Project[map_from_arrays( + * map_keys(computedMap), + * MapEntryArrayMap( + * (mapKey, mapValue) -> valueExpression, + * map_keys(computedMap), map_values(computedMap)))] + * child + * + * after: + * Project[map_from_arrays( + * map_keys(materializedMapSlot), + * MapEntryArrayMap( + * (mapKey, mapValue) -> valueExpression, + * map_keys(materializedMapSlot), map_values(materializedMapSlot)))] + * Project[child.*, computedMap AS materializedMapSlot] + * child + * </pre> + * + * <p> Besides the basic rewrite above, this rule handles + * repeated entry arrays, nested lambdas, and Join children through dedicated helper methods below. + */ +public class AddProjectForMapLambdaInput implements RewriteRuleFactory { + + @Override + public List<Rule> buildRules() { + return ImmutableList.of( + new GenerateRewrite().build(), + new OneRowRelationRewrite().build(), + new ProjectRewrite().build(), + new FilterRewrite().build(), + new HavingRewrite().build(), + new AggregateRewrite().build(), + new JoinRewrite().build() + ); + } + + private class GenerateRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalGenerate().thenApply(ctx -> { + LogicalGenerate<Plan> generate = ctx.root; + List<Function> generators = materializeNestedMapInputs(generate.getGenerators()); + Optional<Pair<List<Function>, LogicalProject<Plan>>> + rewrittenOpt = rewriteExpressions(generate, generators); + if (rewrittenOpt.isPresent()) { + return generate.withGenerators(rewrittenOpt.get().first) + .withChildren(rewrittenOpt.get().second); + } else if (!generators.equals(generate.getGenerators())) { + return generate.withGenerators(generators); + } else { + return generate; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class OneRowRelationRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalOneRowRelation().thenApply(ctx -> { + LogicalOneRowRelation oneRowRelation = ctx.root; + List<NamedExpression> projects = materializeNestedMapInputs(oneRowRelation.getProjects()); + List<NamedExpression> mapInputAliases = tryGenMapInputAliases(projects); + List<NamedExpression> rewrittenProjects = replaceExpressions(projects, mapInputAliases); + List<NamedExpression> entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenProjects); + if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { + return projects.equals(oneRowRelation.getProjects()) + ? oneRowRelation : oneRowRelation.withProjects(projects); + } + + // A OneRowRelation has no child on which to install the usual materialization + // Project. Use the relation itself as the lowest projection, then stack the shared + // entry-array Project and the original output Project above it. + Plan child; + if (mapInputAliases.isEmpty()) { + child = oneRowRelation.withProjects(entryArrayAliases); + } else { + child = oneRowRelation.withProjects(mapInputAliases); + if (!entryArrayAliases.isEmpty()) { + child = appendProject(child, entryArrayAliases); + } + } + rewrittenProjects = replaceExpressions(rewrittenProjects, entryArrayAliases); + return new LogicalProject<>(rewrittenProjects, child); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class ProjectRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalProject().thenApply(ctx -> { + LogicalProject<Plan> project = ctx.root; + List<NamedExpression> projects = materializeNestedMapInputs(project.getProjects()); + Optional<Pair<List<NamedExpression>, LogicalProject<Plan>>> + rewrittenOpt = rewriteExpressions(project, projects); + if (rewrittenOpt.isPresent()) { + return project.withProjectsAndChild(rewrittenOpt.get().first, rewrittenOpt.get().second); + } else if (!projects.equals(project.getProjects())) { + return project.withProjects(projects); + } else { + return project; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class FilterRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalFilter().thenApply(ctx -> { + LogicalFilter<Plan> filter = ctx.root; + List<Expression> conjuncts = materializeNestedMapInputs(filter.getConjuncts()); + Optional<Pair<List<Expression>, LogicalProject<Plan>>> + rewrittenOpt = rewriteExpressions(filter, conjuncts); + if (rewrittenOpt.isPresent()) { + return filter.withConjunctsAndChild( + ImmutableSet.copyOf(rewrittenOpt.get().first), + rewrittenOpt.get().second); + } else if (!ImmutableSet.copyOf(conjuncts).equals(filter.getConjuncts())) { + return filter.withConjuncts(ImmutableSet.copyOf(conjuncts)); + } else { + return filter; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class HavingRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalHaving().thenApply(ctx -> { + LogicalHaving<Plan> having = ctx.root; + List<Expression> conjuncts = materializeNestedMapInputs(having.getConjuncts()); + Optional<Pair<List<Expression>, LogicalProject<Plan>>> + rewrittenOpt = rewriteExpressions(having, conjuncts); + if (rewrittenOpt.isPresent()) { + return having.withConjuncts(ImmutableSet.copyOf(rewrittenOpt.get().first)) + .withChildren(rewrittenOpt.get().second); + } else if (!ImmutableSet.copyOf(conjuncts).equals(having.getConjuncts())) { + return having.withConjuncts(ImmutableSet.copyOf(conjuncts)); + } else { + return having; + } + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class AggregateRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalAggregate().thenApply(ctx -> { + LogicalAggregate<Plan> aggregate = ctx.root; + List<Expression> originalTargets = Lists.newArrayList(); + originalTargets.addAll(aggregate.getGroupByExpressions()); + originalTargets.addAll(aggregate.getOutputExpressions()); + List<Expression> targets = materializeNestedMapInputs(originalTargets); + Optional<Pair<List<Expression>, LogicalProject<Plan>>> rewrittenOpt + = rewriteExpressions(aggregate, targets); + Plan newChild = rewrittenOpt.isPresent() + ? rewrittenOpt.get().second : aggregate.child(); + List<Expression> newTargets = rewrittenOpt.isPresent() + ? rewrittenOpt.get().first : targets; + if (!rewrittenOpt.isPresent() && newTargets.equals(originalTargets)) { + return aggregate; + } + // rewriteExpressions treats group-by expressions and outputs as one ordered list + // so a common Map input is materialized only once. Restore the two original lists + // after replacement. + int groupBySize = aggregate.getGroupByExpressions().size(); + ImmutableList<Expression> newGroupBy = ImmutableList.copyOf( + newTargets.subList(0, groupBySize)); + ImmutableList.Builder<NamedExpression> newOutputBuilder + = ImmutableList.builderWithExpectedSize(aggregate.getOutputExpressions().size()); + for (int i = groupBySize; i < newTargets.size(); i++) { + newOutputBuilder.add((NamedExpression) newTargets.get(i)); + } + return aggregate.withChildGroupByAndOutput(newGroupBy, newOutputBuilder.build(), newChild); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + private class JoinRewrite extends OneRewriteRuleFactory { + @Override + public Rule build() { + return logicalJoin().thenApply(ctx -> { + LogicalJoin<Plan, Plan> join = ctx.root; + int hashOtherConjunctsSize = join.getHashJoinConjuncts().size() + + join.getOtherJoinConjuncts().size(); + int totalConjunctsSize = hashOtherConjunctsSize + join.getMarkJoinConjuncts().size(); + List<Expression> allConjuncts = Lists.newArrayListWithExpectedSize(totalConjunctsSize); + allConjuncts.addAll(join.getHashJoinConjuncts()); + allConjuncts.addAll(join.getOtherJoinConjuncts()); + allConjuncts.addAll(join.getMarkJoinConjuncts()); + List<Expression> originalAllConjuncts = ImmutableList.copyOf(allConjuncts); + allConjuncts = materializeNestedMapInputs(allConjuncts); + Optional<JoinRewriteResult> rewrittenOpt = rewriteJoinExpressions(join, allConjuncts); + if (!rewrittenOpt.isPresent() && allConjuncts.equals(originalAllConjuncts)) { + return join; + } + + Plan newLeftChild = rewrittenOpt.map(result -> result.left).orElse(join.left()); + Plan newRightChild = rewrittenOpt.map(result -> result.right).orElse(join.right()); + List<Expression> newAllConjuncts = rewrittenOpt + .map(result -> result.newConjuncts).orElse(allConjuncts); + List<Expression> newHashOtherConjuncts = newAllConjuncts.subList(0, hashOtherConjunctsSize); + List<Expression> newMarkJoinConjuncts = ImmutableList.copyOf( + newAllConjuncts.subList(hashOtherConjunctsSize, totalConjunctsSize)); + + Pair<List<Expression>, List<Expression>> pair = JoinUtils.extractExpressionForHashTable( + newLeftChild.getOutput(), newRightChild.getOutput(), newHashOtherConjuncts); + List<Expression> newHashJoinConjuncts = pair.first; + List<Expression> newOtherJoinConjuncts = pair.second; + JoinType joinType = join.getJoinType(); + if (joinType == JoinType.CROSS_JOIN && !newHashJoinConjuncts.isEmpty()) { + joinType = JoinType.INNER_JOIN; + } + return new LogicalJoin<>(joinType, + newHashJoinConjuncts, + newOtherJoinConjuncts, + newMarkJoinConjuncts, + join.getDistributeHint(), + join.getMarkJoinSlotReference(), + ImmutableList.of(newLeftChild, newRightChild), + join.getJoinReorderContext()); + }).toRule(RuleType.ADD_PROJECT_FOR_MAP_LAMBDA_INPUT); + } + } + + /** + * Rewrite expressions owned by a single-child plan and install their materialization Projects. + * + * <p>It first materializes computed Map inputs and replaces them in {@code targets}. It then + * materializes any {@link MapEntryArrayMap} still used more than once. These are separate + * Project layers because the second expression can depend on a Map Slot created by the first. + * The returned pair contains the rewritten targets and the top materialization Project. + */ + private <T extends Expression> Optional<Pair<List<T>, LogicalProject<Plan>>> rewriteExpressions( + LogicalPlan plan, Collection<T> targets) { + // computed map materialized + List<NamedExpression> mapInputAliases = tryGenMapInputAliases(targets); + List<T> rewrittenTargets = replaceExpressions(targets, mapInputAliases); + // MapEntryArrayMap merteialized + List<NamedExpression> entryArrayAliases = tryGenSharedEntryArrayAliases(rewrittenTargets); + if (mapInputAliases.isEmpty() && entryArrayAliases.isEmpty()) { + return Optional.empty(); + } + + Plan child = plan.child(0); + if (!mapInputAliases.isEmpty()) { + child = appendProject(child, mapInputAliases); + } + if (!entryArrayAliases.isEmpty()) { + child = appendProject(child, entryArrayAliases); + rewrittenTargets = replaceExpressions(rewrittenTargets, entryArrayAliases); + } + + return Optional.of(Pair.of(rewrittenTargets, (LogicalProject<Plan>) child)); + } + + /** Add aliases without hiding any output already produced by {@code child}. */ + private LogicalProject<Plan> appendProject(Plan child, List<NamedExpression> aliases) { + List<NamedExpression> projects = ImmutableList.<NamedExpression>builder() + .addAll(child.getOutput()) + .addAll(aliases) + .build(); + return new LogicalProject<>(projects, child); + } + + /** Replace each aliased expression by its Slot in all target expression trees. */ + private <T extends Expression> List<T> replaceExpressions( + Collection<T> expressions, List<NamedExpression> aliases) { + if (aliases.isEmpty()) { + return ImmutableList.copyOf(expressions); + } + Map<Expression, Slot> replaceMap = Maps.newHashMap(); + for (NamedExpression alias : aliases) { + replaceMap.put(alias.child(0), alias.toSlot()); + } + ImmutableList.Builder<T> builder = ImmutableList.builderWithExpectedSize(expressions.size()); + for (T expression : expressions) { + builder.add((T) ExpressionUtils.replace(expression, replaceMap)); + } + return builder.build(); + } + + /** + * Rewrite Join conjuncts using the same two materialization stages as + * {@link #rewriteExpressions(LogicalPlan, Collection)}. + * + * <p>Unlike a single-child plan, each generated alias must be attached to the Join child that + * contains all its input Slots. An expression referencing both children cannot be evaluated in + * either child Project, so a deterministic expression is left unchanged and a volatile one is + * rejected. Entry-array aliases are assigned after Map aliases because they may use new Slots. + */ + private Optional<JoinRewriteResult> rewriteJoinExpressions(LogicalJoin<Plan, Plan> join, + Collection<Expression> targets) { + List<Expression> rewrittenTargets = ImmutableList.copyOf(targets); + Plan left = join.left(); + Plan right = join.right(); + + Map<Expression, Set<Slot>> mapInputSlots = Maps.newLinkedHashMap(); + for (Expression target : rewrittenTargets) { + Set<Expression> mapInputs = Sets.newLinkedHashSet(); + collectMapInputs(target, mapInputs); + for (Expression mapInput : mapInputs) { + Set<Slot> inputSlots = mapInput.getInputSlots(); + mapInputSlots.computeIfAbsent(mapInput, ignored -> Sets.newLinkedHashSet()) + .addAll(inputSlots.isEmpty() ? target.getInputSlots() : inputSlots); + } + } + + ImmutableList.Builder<NamedExpression> leftAliases = ImmutableList.builder(); + ImmutableList.Builder<NamedExpression> rightAliases = ImmutableList.builder(); + Map<Expression, Slot> replaceMap = Maps.newHashMap(); + Set<Slot> leftOutputSet = left.getOutputSet(); + Set<Slot> rightOutputSet = right.getOutputSet(); + for (Entry<Expression, Set<Slot>> entry : mapInputSlots.entrySet()) { + Set<Slot> inputSlots = entry.getValue(); + Set<Slot> mapInputExpressionSlots = entry.getKey().getInputSlots(); + if (!mapInputExpressionSlots.isEmpty() + && !leftOutputSet.containsAll(inputSlots) + && !rightOutputSet.containsAll(inputSlots)) { + // No child Project can reference Slots from both sides. Recalculation is safe for + // a deterministic expression, but a volatile Map would no longer have one stable + // value shared by map_keys and map_values. + if (entry.getKey().containsVolatileExpression()) { + throw new AnalysisException( + "A computed Map input containing a volatile expression cannot " + + "reference both sides of a join"); + } + continue; + } + ExprId exprId = StatementScopeIdGenerator.newExprId(); Review Comment: doris 目前不支持在左右侧生成之后，执行 ON 的 other conjunct 之前插入 project，这块行为暂时先保持这样吧 -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
