Search in sources :

Example 1 with AggregationIfToFilterRewriteStrategy

use of com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationIfToFilterRewriteStrategy in project presto by prestodb.

the class RewriteAggregationIfToFilter method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    ProjectNode sourceProject = captures.get(CHILD);
    Set<Aggregation> aggregationsToRewrite = aggregationNode.getAggregations().values().stream().filter(aggregation -> shouldRewriteAggregation(aggregation, sourceProject)).collect(toImmutableSet());
    if (aggregationsToRewrite.isEmpty()) {
        return Result.empty();
    }
    context.getSession().getRuntimeStats().addMetricValue(REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED, 1);
    // Get the corresponding assignments in the input project.
    // The aggregationReferences only has the aggregations to rewrite, thus the sourceAssignments only has IF/CAST(IF) expressions with NULL false results.
    // Multiple aggregations may reference the same input. We use a map to dedup them based on the VariableReferenceExpression, so that we only do the rewrite once per input
    // IF expression.
    // The order of sourceAssignments determines the order of generating the new variables for the IF conditions and results. We use a sorted map to get a deterministic
    // order based on the name of the VariableReferenceExpressions.
    Map<VariableReferenceExpression, RowExpression> sourceAssignments = aggregationsToRewrite.stream().map(aggregation -> (VariableReferenceExpression) aggregation.getArguments().get(0)).collect(toImmutableSortedMap(VariableReferenceExpression::compareTo, identity(), variable -> sourceProject.getAssignments().get(variable), (left, right) -> left));
    Assignments.Builder newAssignments = Assignments.builder();
    newAssignments.putAll(sourceProject.getAssignments());
    // Map from the aggregation reference to the IF condition reference which will be put in the mask.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToConditionReference = new HashMap<>();
    // Map from the aggregation reference to the IF result reference. This only contains the aggregates where the IF can be safely unwrapped.
    // E.g., SUM(IF(CARDINALITY(array) > 0, array[1])) will not be included in this map as array[1] can return errors if we unwrap the IF.
    Map<VariableReferenceExpression, VariableReferenceExpression> aggregationReferenceToIfResultReference = new HashMap<>();
    AggregationIfToFilterRewriteStrategy rewriteStrategy = getAggregationIfToFilterRewriteStrategy(context.getSession());
    for (Map.Entry<VariableReferenceExpression, RowExpression> entry : sourceAssignments.entrySet()) {
        VariableReferenceExpression outputVariable = entry.getKey();
        RowExpression rowExpression = entry.getValue();
        SpecialFormExpression ifExpression = (SpecialFormExpression) ((rowExpression instanceof CallExpression) ? ((CallExpression) rowExpression).getArguments().get(0) : rowExpression);
        RowExpression condition = ifExpression.getArguments().get(0);
        VariableReferenceExpression conditionReference = context.getVariableAllocator().newVariable(condition);
        newAssignments.put(conditionReference, condition);
        aggregationReferenceToConditionReference.put(outputVariable, conditionReference);
        if (canUnwrapIf(ifExpression, rewriteStrategy)) {
            RowExpression trueResult = ifExpression.getArguments().get(1);
            if (rowExpression instanceof CallExpression) {
                // Wrap the result with CAST().
                trueResult = new CallExpression(((CallExpression) rowExpression).getDisplayName(), ((CallExpression) rowExpression).getFunctionHandle(), rowExpression.getType(), ImmutableList.of(trueResult));
            }
            VariableReferenceExpression ifResultReference = context.getVariableAllocator().newVariable(trueResult);
            newAssignments.put(ifResultReference, trueResult);
            aggregationReferenceToIfResultReference.put(outputVariable, ifResultReference);
        }
    }
    // Build new aggregations.
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    // Stores the masks used to build the filter predicates. Use set to dedup the predicates.
    ImmutableSortedSet.Builder<VariableReferenceExpression> masks = ImmutableSortedSet.naturalOrder();
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        VariableReferenceExpression output = entry.getKey();
        Aggregation aggregation = entry.getValue();
        if (!aggregationsToRewrite.contains(aggregation)) {
            aggregations.put(output, aggregation);
            continue;
        }
        VariableReferenceExpression aggregationReference = (VariableReferenceExpression) aggregation.getArguments().get(0);
        CallExpression callExpression = aggregation.getCall();
        VariableReferenceExpression ifResultReference = aggregationReferenceToIfResultReference.get(aggregationReference);
        if (ifResultReference != null) {
            callExpression = new CallExpression(callExpression.getSourceLocation(), callExpression.getDisplayName(), callExpression.getFunctionHandle(), callExpression.getType(), ImmutableList.of(ifResultReference));
        }
        VariableReferenceExpression mask = aggregationReferenceToConditionReference.get(aggregationReference);
        aggregations.put(output, new Aggregation(callExpression, Optional.empty(), aggregation.getOrderBy(), aggregation.isDistinct(), Optional.of(aggregationReferenceToConditionReference.get(aggregationReference))));
        masks.add(mask);
    }
    RowExpression predicate = TRUE_CONSTANT;
    if (!aggregationNode.hasNonEmptyGroupingSet() && aggregationsToRewrite.size() == aggregationNode.getAggregations().size()) {
        // All aggregations are rewritten by this rule. We can add a filter with all the masks to make the query more efficient.
        predicate = or(masks.build());
    }
    return Result.ofPlanNode(new AggregationNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new FilterNode(aggregationNode.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), sourceProject.getSource(), newAssignments.build()), predicate), aggregations.build(), aggregationNode.getGroupingSets(), aggregationNode.getPreGroupedVariables(), aggregationNode.getStep(), aggregationNode.getHashVariable(), aggregationNode.getGroupIdVariable()));
}
Also used : FunctionAndTypeManager(com.facebook.presto.metadata.FunctionAndTypeManager) AggregationNode(com.facebook.presto.spi.plan.AggregationNode) Captures(com.facebook.presto.matching.Captures) Assignments(com.facebook.presto.spi.plan.Assignments) AtomicBoolean(java.util.concurrent.atomic.AtomicBoolean) ImmutableSortedMap.toImmutableSortedMap(com.google.common.collect.ImmutableSortedMap.toImmutableSortedMap) VariableReferenceExpression(com.facebook.presto.spi.relation.VariableReferenceExpression) HashMap(java.util.HashMap) RowExpressionDeterminismEvaluator(com.facebook.presto.sql.relational.RowExpressionDeterminismEvaluator) SystemSessionProperties.getAggregationIfToFilterRewriteStrategy(com.facebook.presto.SystemSessionProperties.getAggregationIfToFilterRewriteStrategy) StandardFunctionResolution(com.facebook.presto.spi.function.StandardFunctionResolution) Pattern(com.facebook.presto.matching.Pattern) FilterNode(com.facebook.presto.spi.plan.FilterNode) Capture(com.facebook.presto.matching.Capture) ImmutableList(com.google.common.collect.ImmutableList) UNWRAP_IF(com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationIfToFilterRewriteStrategy.UNWRAP_IF) IF(com.facebook.presto.spi.relation.SpecialFormExpression.Form.IF) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) ImmutableSet.toImmutableSet(com.google.common.collect.ImmutableSet.toImmutableSet) Expressions(com.facebook.presto.sql.relational.Expressions) FunctionResolution(com.facebook.presto.sql.relational.FunctionResolution) CallExpression(com.facebook.presto.spi.relation.CallExpression) SpecialFormExpression(com.facebook.presto.spi.relation.SpecialFormExpression) VariablesExtractor(com.facebook.presto.sql.planner.VariablesExtractor) RowExpression(com.facebook.presto.spi.relation.RowExpression) ImmutableSortedSet(com.google.common.collect.ImmutableSortedSet) Patterns.aggregation(com.facebook.presto.sql.planner.plan.Patterns.aggregation) ImmutableMap(com.google.common.collect.ImmutableMap) Session(com.facebook.presto.Session) Rule(com.facebook.presto.sql.planner.iterative.Rule) Patterns.project(com.facebook.presto.sql.planner.plan.Patterns.project) Set(java.util.Set) LambdaDefinitionExpression(com.facebook.presto.spi.relation.LambdaDefinitionExpression) OperatorType(com.facebook.presto.common.function.OperatorType) AggregationIfToFilterRewriteStrategy(com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationIfToFilterRewriteStrategy) TRUE_CONSTANT(com.facebook.presto.expressions.LogicalRowExpressions.TRUE_CONSTANT) Patterns.source(com.facebook.presto.sql.planner.plan.Patterns.source) FILTER_WITH_IF(com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationIfToFilterRewriteStrategy.FILTER_WITH_IF) DefaultRowExpressionTraversalVisitor(com.facebook.presto.expressions.DefaultRowExpressionTraversalVisitor) ProjectNode(com.facebook.presto.spi.plan.ProjectNode) LogicalRowExpressions.or(com.facebook.presto.expressions.LogicalRowExpressions.or) Capture.newCapture(com.facebook.presto.matching.Capture.newCapture) Aggregation(com.facebook.presto.spi.plan.AggregationNode.Aggregation) Function.identity(java.util.function.Function.identity) Optional(java.util.Optional) REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED(com.facebook.presto.common.RuntimeMetricName.REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED) SystemSessionProperties.getAggregationIfToFilterRewriteStrategy(com.facebook.presto.SystemSessionProperties.getAggregationIfToFilterRewriteStrategy) AggregationIfToFilterRewriteStrategy(com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationIfToFilterRewriteStrategy) HashMap(java.util.HashMap) FilterNode(com.facebook.presto.spi.plan.FilterNode) Assignments(com.facebook.presto.spi.plan.Assignments) RowExpression(com.facebook.presto.spi.relation.RowExpression) AggregationNode(com.facebook.presto.spi.plan.AggregationNode) ImmutableMap(com.google.common.collect.ImmutableMap) Aggregation(com.facebook.presto.spi.plan.AggregationNode.Aggregation) VariableReferenceExpression(com.facebook.presto.spi.relation.VariableReferenceExpression) ImmutableSortedSet(com.google.common.collect.ImmutableSortedSet) ProjectNode(com.facebook.presto.spi.plan.ProjectNode) ImmutableSortedMap.toImmutableSortedMap(com.google.common.collect.ImmutableSortedMap.toImmutableSortedMap) HashMap(java.util.HashMap) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) SpecialFormExpression(com.facebook.presto.spi.relation.SpecialFormExpression) CallExpression(com.facebook.presto.spi.relation.CallExpression)

Aggregations

Session (com.facebook.presto.Session)1 SystemSessionProperties.getAggregationIfToFilterRewriteStrategy (com.facebook.presto.SystemSessionProperties.getAggregationIfToFilterRewriteStrategy)1 REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED (com.facebook.presto.common.RuntimeMetricName.REWRITE_AGGREGATION_IF_TO_FILTER_APPLIED)1 OperatorType (com.facebook.presto.common.function.OperatorType)1 DefaultRowExpressionTraversalVisitor (com.facebook.presto.expressions.DefaultRowExpressionTraversalVisitor)1 TRUE_CONSTANT (com.facebook.presto.expressions.LogicalRowExpressions.TRUE_CONSTANT)1 LogicalRowExpressions.or (com.facebook.presto.expressions.LogicalRowExpressions.or)1 Capture (com.facebook.presto.matching.Capture)1 Capture.newCapture (com.facebook.presto.matching.Capture.newCapture)1 Captures (com.facebook.presto.matching.Captures)1 Pattern (com.facebook.presto.matching.Pattern)1 FunctionAndTypeManager (com.facebook.presto.metadata.FunctionAndTypeManager)1 StandardFunctionResolution (com.facebook.presto.spi.function.StandardFunctionResolution)1 AggregationNode (com.facebook.presto.spi.plan.AggregationNode)1 Aggregation (com.facebook.presto.spi.plan.AggregationNode.Aggregation)1 Assignments (com.facebook.presto.spi.plan.Assignments)1 FilterNode (com.facebook.presto.spi.plan.FilterNode)1 ProjectNode (com.facebook.presto.spi.plan.ProjectNode)1 CallExpression (com.facebook.presto.spi.relation.CallExpression)1 LambdaDefinitionExpression (com.facebook.presto.spi.relation.LambdaDefinitionExpression)1