Examples with Aggregation com.facebook.presto.spi.plan.AggregationNode.Aggregation used on opensource projects

Example 11 with Aggregation

use of com.facebook.presto.spi.plan.AggregationNode.Aggregation in project presto by prestodb.

the class AggregationStatsRule method groupBy.

public static PlanNodeStatsEstimate groupBy(PlanNodeStatsEstimate sourceStats, Collection<VariableReferenceExpression> groupByVariables, Map<VariableReferenceExpression, Aggregation> aggregations) {
    PlanNodeStatsEstimate.Builder result = PlanNodeStatsEstimate.builder();
    for (VariableReferenceExpression groupByVariable : groupByVariables) {
        VariableStatsEstimate symbolStatistics = sourceStats.getVariableStatistics(groupByVariable);
        result.addVariableStatistics(groupByVariable, symbolStatistics.mapNullsFraction(nullsFraction -> {
            if (nullsFraction == 0.0) {
                return 0.0;
            }
            return 1.0 / (symbolStatistics.getDistinctValuesCount() + 1);
        }));
    }
    double rowsCount = 1;
    for (VariableReferenceExpression groupByVariable : groupByVariables) {
        VariableStatsEstimate symbolStatistics = sourceStats.getVariableStatistics(groupByVariable);
        int nullRow = (symbolStatistics.getNullsFraction() == 0.0) ? 0 : 1;
        rowsCount *= symbolStatistics.getDistinctValuesCount() + nullRow;
    }
    result.setOutputRowCount(min(rowsCount, sourceStats.getOutputRowCount()));
    for (Map.Entry<VariableReferenceExpression, Aggregation> aggregationEntry : aggregations.entrySet()) {
        result.addVariableStatistics(aggregationEntry.getKey(), estimateAggregationStats(aggregationEntry.getValue(), sourceStats));
    }
    return result.build();
}

Example 12 with Aggregation

use of com.facebook.presto.spi.plan.AggregationNode.Aggregation in project presto by prestodb.

the class MultipleDistinctAggregationToMarkDistinct method apply.

@Override
public Result apply(AggregationNode parent, Captures captures, Context context) {
    if (!SystemSessionProperties.useMarkDistinct(context.getSession())) {
        return Result.empty();
    }
    // the distinct marker for the given set of input columns
    Map<Set<VariableReferenceExpression>, VariableReferenceExpression> markers = new HashMap<>();
    Map<VariableReferenceExpression, Aggregation> newAggregations = new HashMap<>();
    PlanNode subPlan = parent.getSource();
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : parent.getAggregations().entrySet()) {
        Aggregation aggregation = entry.getValue();
        if (aggregation.isDistinct() && !aggregation.getFilter().isPresent() && !aggregation.getMask().isPresent()) {
            Set<VariableReferenceExpression> inputs = aggregation.getArguments().stream().map(OriginalExpressionUtils::castToExpression).map(context.getVariableAllocator()::toVariableReference).collect(toSet());
            VariableReferenceExpression marker = markers.get(inputs);
            if (marker == null) {
                marker = context.getVariableAllocator().newVariable(Iterables.getLast(inputs).getName(), BOOLEAN, "distinct");
                markers.put(inputs, marker);
                ImmutableSet.Builder<VariableReferenceExpression> distinctVariables = ImmutableSet.<VariableReferenceExpression>builder().addAll(parent.getGroupingKeys()).addAll(inputs);
                parent.getGroupIdVariable().ifPresent(distinctVariables::add);
                subPlan = new MarkDistinctNode(subPlan.getSourceLocation(), context.getIdAllocator().getNextId(), subPlan, marker, ImmutableList.copyOf(distinctVariables.build()), Optional.empty());
            }
            // remove the distinct flag and set the distinct marker
            newAggregations.put(entry.getKey(), new Aggregation(aggregation.getCall(), aggregation.getFilter(), aggregation.getOrderBy(), false, Optional.of(marker)));
        } else {
            newAggregations.put(entry.getKey(), aggregation);
        }
    }
    return Result.ofPlanNode(new AggregationNode(parent.getSourceLocation(), parent.getId(), subPlan, newAggregations, parent.getGroupingSets(), ImmutableList.of(), parent.getStep(), parent.getHashVariable(), parent.getGroupIdVariable()));
}

Example 13 with Aggregation

use of com.facebook.presto.spi.plan.AggregationNode.Aggregation in project presto by prestodb.

the class RewriteSpatialPartitioningAggregation method apply.

@Override
public Result apply(AggregationNode node, Captures captures, Context context) {
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    VariableReferenceExpression partitionCountVariable = context.getVariableAllocator().newVariable("partition_count", INTEGER);
    ImmutableMap.Builder<VariableReferenceExpression, RowExpression> envelopeAssignments = ImmutableMap.builder();
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : node.getAggregations().entrySet()) {
        Aggregation aggregation = entry.getValue();
        QualifiedObjectName name = metadata.getFunctionAndTypeManager().getFunctionMetadata(aggregation.getFunctionHandle()).getName();
        Type geometryType = metadata.getType(GEOMETRY_TYPE_SIGNATURE);
        if (name.equals(NAME) && aggregation.getArguments().size() == 1) {
            RowExpression geometry = getOnlyElement(aggregation.getArguments());
            VariableReferenceExpression envelopeVariable = context.getVariableAllocator().newVariable(aggregation.getCall().getSourceLocation(), "envelope", geometryType);
            if (isFunctionNameMatch(geometry, "ST_Envelope")) {
                envelopeAssignments.put(envelopeVariable, geometry);
            } else {
                envelopeAssignments.put(envelopeVariable, castToRowExpression(new FunctionCall(QualifiedName.of("ST_Envelope"), ImmutableList.of(castToExpression(geometry)))));
            }
            aggregations.put(entry.getKey(), new Aggregation(new CallExpression(envelopeVariable.getSourceLocation(), name.getObjectName(), metadata.getFunctionAndTypeManager().lookupFunction(NAME.getObjectName(), fromTypes(geometryType, INTEGER)), entry.getKey().getType(), ImmutableList.of(castToRowExpression(asSymbolReference(envelopeVariable)), castToRowExpression(asSymbolReference(partitionCountVariable)))), Optional.empty(), Optional.empty(), false, aggregation.getMask()));
        } else {
            aggregations.put(entry);
        }
    }
    return Result.ofPlanNode(new AggregationNode(node.getSourceLocation(), node.getId(), new ProjectNode(context.getIdAllocator().getNextId(), node.getSource(), Assignments.builder().putAll(identitiesAsSymbolReferences(node.getSource().getOutputVariables())).put(partitionCountVariable, castToRowExpression(new LongLiteral(Integer.toString(getHashPartitionCount(context.getSession()))))).putAll(envelopeAssignments.build()).build()), aggregations.build(), node.getGroupingSets(), node.getPreGroupedVariables(), node.getStep(), node.getHashVariable(), node.getGroupIdVariable()));
}

Example 14 with Aggregation

use of com.facebook.presto.spi.plan.AggregationNode.Aggregation 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()));
}

Example 15 with Aggregation

use of com.facebook.presto.spi.plan.AggregationNode.Aggregation in project presto by prestodb.

the class ImplementFilteredAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Assignments.Builder newAssignments = Assignments.builder();
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    ImmutableList.Builder<Expression> maskSymbols = ImmutableList.builder();
    boolean aggregateWithoutFilterPresent = false;
    for (Map.Entry<VariableReferenceExpression, Aggregation> entry : aggregation.getAggregations().entrySet()) {
        VariableReferenceExpression output = entry.getKey();
        // strip the filters
        Optional<VariableReferenceExpression> mask = entry.getValue().getMask();
        if (entry.getValue().getFilter().isPresent()) {
            // TODO remove cast once assignment can be RowExpression
            Expression filter = OriginalExpressionUtils.castToExpression(entry.getValue().getFilter().get());
            VariableReferenceExpression variable = context.getVariableAllocator().newVariable(filter, BOOLEAN);
            verify(!mask.isPresent(), "Expected aggregation without mask symbols, see Rule pattern");
            newAssignments.put(variable, castToRowExpression(filter));
            mask = Optional.of(variable);
            maskSymbols.add(createSymbolReference(variable));
        } else {
            aggregateWithoutFilterPresent = true;
        }
        aggregations.put(output, new Aggregation(entry.getValue().getCall(), Optional.empty(), entry.getValue().getOrderBy(), entry.getValue().isDistinct(), mask));
    }
    Expression predicate = TRUE_LITERAL;
    if (!aggregation.hasNonEmptyGroupingSet() && !aggregateWithoutFilterPresent) {
        predicate = combineDisjunctsWithDefault(maskSymbols.build(), TRUE_LITERAL);
    }
    // identity projection for all existing inputs
    newAssignments.putAll(identitiesAsSymbolReferences(aggregation.getSource().getOutputVariables()));
    return Result.ofPlanNode(new AggregationNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), new FilterNode(aggregation.getSourceLocation(), context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), aggregation.getSource(), newAssignments.build()), castToRowExpression(predicate)), aggregations.build(), aggregation.getGroupingSets(), ImmutableList.of(), aggregation.getStep(), aggregation.getHashVariable(), aggregation.getGroupIdVariable()));
}