linrrzqqq commented on code in PR #66968: URL: https://github.com/apache/doris/pull/66968#discussion_r3829970943
########## 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); Review Comment: 普通 project 不好放进一个标量的 if 表达式里面，先保证正确性 ########## be/src/exprs/lambda_function/varray_map_function.cpp: ########## @@ -265,11 +265,17 @@ class ArrayMapFunction : public LambdaFunction { const size_t lambda_batch_rows = _calculate_lambda_batch_size(children[0], lambda_datas, block, required_input_column_ids, has_row_dependent_captures); + // Reuse the nested input columns directly when they fit within eight regular lambda + // batches. Larger inputs use the base batch size, while a smaller byte-budget-derived + // batch remains authoritative. + const size_t lambda_fast_path_rows = lambda_batch_rows == _lambda_block_budget.max_rows Review Comment: 经验值，可以保留 -- 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]
