924060929 commented on code in PR #14827: URL: https://github.com/apache/doris/pull/14827#discussion_r1050649218
########## fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/AggregateStrategies.java: ########## @@ -0,0 +1,1234 @@ +// 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.Column; +import org.apache.doris.catalog.KeysType; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.common.Pair; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.annotation.DependsRules; +import org.apache.doris.nereids.pattern.PatternDescriptor; +import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; +import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.properties.RequireProperties; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.rules.expression.rewrite.ExpressionRewriteContext; +import org.apache.doris.nereids.rules.expression.rewrite.rules.FoldConstantRuleOnFE; +import org.apache.doris.nereids.rules.expression.rewrite.rules.TypeCoercion; +import org.apache.doris.nereids.rules.implementation.ImplementationRuleFactory; +import org.apache.doris.nereids.rules.implementation.LogicalOlapScanToPhysicalOlapScan; +import org.apache.doris.nereids.rules.rewrite.logical.NormalizeAggregate; +import org.apache.doris.nereids.trees.expressions.AggregateExpression; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.IsNull; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; +import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; +import org.apache.doris.nereids.trees.expressions.functions.agg.Count; +import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctCount; +import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctSum; +import org.apache.doris.nereids.trees.expressions.functions.agg.Sum; +import org.apache.doris.nereids.trees.expressions.functions.scalar.If; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.plans.AggMode; +import org.apache.doris.nereids.trees.plans.AggPhase; +import org.apache.doris.nereids.trees.plans.GroupPlan; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.Project; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.physical.PhysicalHashAggregate; +import org.apache.doris.nereids.trees.plans.physical.PhysicalOlapScan; +import org.apache.doris.nereids.trees.plans.physical.PhysicalStorageLayerAggregate; +import org.apache.doris.nereids.trees.plans.physical.PhysicalStorageLayerAggregate.PushDownAggOp; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.qe.ConnectContext; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; + +/** AggregateStrategies */ +@DependsRules({ + NormalizeAggregate.class, + FoldConstantRuleOnFE.class +}) +public class AggregateStrategies implements ImplementationRuleFactory { + + @Override + public List<Rule> buildRules() { + PatternDescriptor<LogicalAggregate<GroupPlan>> basePattern = logicalAggregate() + .when(LogicalAggregate::isNormalized); + + return ImmutableList.of( + RuleType.STORAGE_LAYER_AGGREGATE_WITHOUT_PROJECT.build( + logicalAggregate( + logicalOlapScan() + ) + .when(agg -> agg.isNormalized() && enablePushDownNoGroupAgg()) + .thenApply(ctx -> storageLayerAggregate(ctx.root, null, ctx.root.child(), ctx.cascadesContext)) + ), + RuleType.STORAGE_LAYER_AGGREGATE_WITH_PROJECT.build( + logicalAggregate( + logicalProject( + logicalOlapScan() + ) + ) + .when(agg -> agg.isNormalized() && enablePushDownNoGroupAgg()) + .thenApply(ctx -> { + LogicalAggregate<LogicalProject<LogicalOlapScan>> agg = ctx.root; + LogicalProject<LogicalOlapScan> project = agg.child(); + LogicalOlapScan olapScan = project.child(); + return storageLayerAggregate(agg, project, olapScan, ctx.cascadesContext); + }) + ), + RuleType.ONE_PHASE_AGGREGATE_WITHOUT_DISTINCT.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 0) + .thenApplyMulti(ctx -> onePhaseAggregateWithoutDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.TWO_PHASE_AGGREGATE_WITHOUT_DISTINCT.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 0) + .thenApplyMulti(ctx -> twoPhaseAggregateWithoutDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.TWO_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI.build( + basePattern + .when(this::containsCountDistinctMultiExpr) + .thenApplyMulti(ctx -> twoPhaseAggregateWithCountDistinctMulti(ctx.root, ctx.connectContext)) + ), + RuleType.THREE_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI.build( + basePattern + .when(this::containsCountDistinctMultiExpr) + .thenApplyMulti(ctx -> threePhaseAggregateWithCountDistinctMulti(ctx.root, ctx.connectContext)) + ), + RuleType.TWO_PHASE_AGGREGATE_WITH_DISTINCT.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 1) + .thenApplyMulti(ctx -> twoPhaseAggregateWithDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.ONE_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 1 && enableSingleDistinctColumnOpt()) + .thenApplyMulti(ctx -> onePhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.TWO_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 1 && enableSingleDistinctColumnOpt()) + .thenApplyMulti(ctx -> twoPhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.THREE_PHASE_AGGREGATE_WITH_DISTINCT.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() == 1) + .thenApplyMulti(ctx -> threePhaseAggregateWithDistinct(ctx.root, ctx.connectContext)) + ), + RuleType.TWO_PHASE_AGGREGATE_WITH_MULTI_DISTINCT.build( + basePattern + .when(agg -> agg.getDistinctArguments().size() > 1 && !containsCountDistinctMultiExpr(agg)) + .thenApplyMulti(ctx -> twoPhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext)) + ) + ); + } + + /** + * sql: select count(*) from tbl + * + * before: + * + * LogicalAggregate(groupBy=[], output=[count(*)]) + * | + * LogicalOlapScan(table=tbl) + * + * after: + * + * LogicalAggregate(groupBy=[], output=[count(*)]) + * | + * PhysicalStorageLayerAggregate(pushAggOp=COUNT, table=PhysicalOlapScan(table=tbl)) + * + */ + private LogicalAggregate<? extends Plan> storageLayerAggregate( + LogicalAggregate<? extends Plan> aggregate, + @Nullable LogicalProject<? extends Plan> project, + LogicalOlapScan olapScan, CascadesContext cascadesContext) { + final LogicalAggregate<? extends Plan> canNotPush = aggregate; + + KeysType keysType = olapScan.getTable().getKeysType(); + if (keysType != KeysType.AGG_KEYS && keysType != KeysType.DUP_KEYS) { + return canNotPush; + } + + List<Expression> groupByExpressions = aggregate.getGroupByExpressions(); + if (!groupByExpressions.isEmpty() || !aggregate.getDistinctArguments().isEmpty()) { + return canNotPush; + } + + Set<AggregateFunction> aggregateFunctions = aggregate.getAggregateFunctions(); + Set<Class<? extends AggregateFunction>> functionClasses = aggregateFunctions + .stream() + .map(AggregateFunction::getClass) + .collect(Collectors.toSet()); + + Map<Class, PushDownAggOp> supportedAgg = PushDownAggOp.supportedFunctions(); + if (!supportedAgg.keySet().containsAll(functionClasses)) { + return canNotPush; + } + if (functionClasses.contains(Count.class) && keysType != KeysType.DUP_KEYS) { + return canNotPush; + } + if (aggregateFunctions.stream().anyMatch(fun -> fun.arity() > 1)) { + return canNotPush; + } + + // we already normalize the arguments to slotReference + List<Expression> argumentsOfAggregateFunction = aggregateFunctions.stream() + .flatMap(aggregateFunction -> aggregateFunction.getArguments().stream()) + .collect(ImmutableList.toImmutableList()); + + if (project != null) { + argumentsOfAggregateFunction = Project.findProject( + (List<SlotReference>) (List) argumentsOfAggregateFunction, project.getProjects()) + .stream() + .map(p -> p instanceof Alias ? p.child(0) : p) + .collect(ImmutableList.toImmutableList()); + } + + boolean onlyContainsSlotOrNumericCastSlot = argumentsOfAggregateFunction + .stream() + .allMatch(argument -> { + if (argument instanceof SlotReference) { + return true; + } + if (argument instanceof Cast) { + return argument.child(0) instanceof SlotReference + && argument.getDataType().isNumericType() + && argument.child(0).getDataType().isNumericType(); + } + return false; + }); + if (!onlyContainsSlotOrNumericCastSlot) { + return canNotPush; + } + + Set<PushDownAggOp> pushDownAggOps = functionClasses.stream() + .map(supportedAgg::get) + .collect(Collectors.toSet()); + + PushDownAggOp mergeOp = pushDownAggOps.size() == 1 + ? pushDownAggOps.iterator().next() + : PushDownAggOp.MIX; + + Set<SlotReference> aggUsedSlots = + ExpressionUtils.collect(argumentsOfAggregateFunction, SlotReference.class::isInstance); + + List<SlotReference> usedSlotInTable = (List<SlotReference>) (List) Project.findProject(aggUsedSlots, + (List<NamedExpression>) (List) olapScan.getOutput()); + + for (SlotReference slot : usedSlotInTable) { + Column column = slot.getColumn().get(); + if (keysType == KeysType.AGG_KEYS && !column.isKey()) { + return canNotPush; + } + // The zone map max length of CharFamily is 512, do not + // over the length: https://github.com/apache/doris/pull/6293 + if (mergeOp == PushDownAggOp.MIN_MAX || mergeOp == PushDownAggOp.MIX) { + PrimitiveType colType = column.getType().getPrimitiveType(); + if (colType.isArrayType() || colType.isComplexType() || colType == PrimitiveType.STRING) { + return canNotPush; + } + if (colType.isCharFamily() && mergeOp != PushDownAggOp.COUNT && column.getType().getLength() > 512) { + return canNotPush; + } + } + if (mergeOp == PushDownAggOp.COUNT || mergeOp == PushDownAggOp.MIX) { + // NULL value behavior in `count` function is zero, so + // we should not use row_count to speed up query. the col + // must be not null + if (column.isAllowNull()) { + return canNotPush; + } + } + } + + PhysicalOlapScan physicalOlapScan = (PhysicalOlapScan) new LogicalOlapScanToPhysicalOlapScan() + .build() + .transform(olapScan, cascadesContext) + .get(0); + + return aggregate.withChildren(ImmutableList.of( + new PhysicalStorageLayerAggregate(physicalOlapScan, mergeOp) + )); + } + + /** + * sql: select count(*) from tbl group by id + * + * before: + * + * LogicalAggregate(groupBy=[id], output=[count(*)]) + * | + * LogicalOlapScan(table=tbl) + * + * after: + * + * single node aggregate: + * + * PhysicalHashAggregate(groupBy=[id], output=[count(*)]) + * | + * PhysicalDistribute(distributionSpec=GATHER) + * | + * LogicalOlapScan(table=tbl) + * + * distribute node aggregate: + * + * PhysicalHashAggregate(groupBy=[id], output=[count(*)]) + * | + * LogicalOlapScan(table=tbl, **already distribute by id**) + * + */ + private List<PhysicalHashAggregate<Plan>> onePhaseAggregateWithoutDistinct( + LogicalAggregate<? extends Plan> logicalAgg, ConnectContext connectContext) { + RequireProperties requireGather = RequireProperties.of(PhysicalProperties.GATHER); Review Comment: RequireProperties has a `List<PhysicalProperties> properties` match to every require physical properties of children. so `RequireProperties.GATHER` is only used for one child and not very general -- 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]
