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

Example 16 with AggregationNode

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

the class TestEffectivePredicateExtractor method testAggregation.

@Test
public void testAggregation() {
    PlanNode node = new AggregationNode(Optional.empty(), newId(), filter(baseTableScan, and(equals(AV, DV), equals(BV, EV), equals(CV, FV), lessThan(DV, bigintLiteral(10)), lessThan(CV, DV), greaterThan(AV, bigintLiteral(2)), equals(EV, FV))), ImmutableMap.of(CV, count(metadata.getFunctionAndTypeManager()), DV, count(metadata.getFunctionAndTypeManager())), singleGroupingSet(ImmutableList.of(AV, BV, CV)), ImmutableList.of(), AggregationNode.Step.FINAL, Optional.empty(), Optional.empty());
    RowExpression effectivePredicate = effectivePredicateExtractor.extract(node);
    // Rewrite in terms of group by symbols
    assertEquals(normalizeConjuncts(effectivePredicate), normalizeConjuncts(lessThan(AV, bigintLiteral(10)), lessThan(BV, AV), greaterThan(AV, bigintLiteral(2)), equals(BV, CV)));
}

Example 17 with AggregationNode

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

the class QueryPlanner method aggregate.

private PlanBuilder aggregate(PlanBuilder subPlan, QuerySpecification node) {
    if (!analysis.isAggregation(node)) {
        return subPlan;
    }
    // 1. Pre-project all scalar inputs (arguments and non-trivial group by expressions)
    Set<Expression> groupByExpressions = ImmutableSet.copyOf(analysis.getGroupByExpressions(node));
    ImmutableList.Builder<Expression> arguments = ImmutableList.builder();
    analysis.getAggregates(node).stream().map(FunctionCall::getArguments).flatMap(List::stream).filter(// lambda expression is generated at execution time
    exp -> !(exp instanceof LambdaExpression)).forEach(arguments::add);
    analysis.getAggregates(node).stream().map(FunctionCall::getOrderBy).filter(Optional::isPresent).map(Optional::get).map(OrderBy::getSortItems).flatMap(List::stream).map(SortItem::getSortKey).forEach(arguments::add);
    // filter expressions need to be projected first
    analysis.getAggregates(node).stream().map(FunctionCall::getFilter).filter(Optional::isPresent).map(Optional::get).forEach(arguments::add);
    Iterable<Expression> inputs = Iterables.concat(groupByExpressions, arguments.build());
    subPlan = handleSubqueries(subPlan, node, inputs);
    if (!Iterables.isEmpty(inputs)) {
        // avoid an empty projection if the only aggregation is COUNT (which has no arguments)
        subPlan = project(subPlan, inputs);
    }
    // 2. Aggregate
    // 2.a. Rewrite aggregate arguments
    TranslationMap argumentTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis, lambdaDeclarationToVariableMap);
    ImmutableList.Builder<VariableReferenceExpression> aggregationArgumentsBuilder = ImmutableList.builder();
    for (Expression argument : arguments.build()) {
        VariableReferenceExpression variable = subPlan.translate(argument);
        argumentTranslations.put(argument, variable);
        aggregationArgumentsBuilder.add(variable);
    }
    List<VariableReferenceExpression> aggregationArguments = aggregationArgumentsBuilder.build();
    // 2.b. Rewrite grouping columns
    TranslationMap groupingTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis, lambdaDeclarationToVariableMap);
    Map<VariableReferenceExpression, VariableReferenceExpression> groupingSetMappings = new LinkedHashMap<>();
    for (Expression expression : groupByExpressions) {
        VariableReferenceExpression input = subPlan.translate(expression);
        VariableReferenceExpression output = variableAllocator.newVariable(expression, analysis.getTypeWithCoercions(expression), "gid");
        groupingTranslations.put(expression, output);
        groupingSetMappings.put(output, input);
    }
    // This tracks the grouping sets before complex expressions are considered (see comments below)
    // It's also used to compute the descriptors needed to implement grouping()
    List<Set<FieldId>> columnOnlyGroupingSets = ImmutableList.of(ImmutableSet.of());
    List<List<VariableReferenceExpression>> groupingSets = ImmutableList.of(ImmutableList.of());
    if (node.getGroupBy().isPresent()) {
        // For the purpose of "distinct", we need to canonicalize column references that may have varying
        // syntactic forms (e.g., "t.a" vs "a"). Thus we need to enumerate grouping sets based on the underlying
        // fieldId associated with each column reference expression.
        // The catch is that simple group-by expressions can be arbitrary expressions (this is a departure from the SQL specification).
        // But, they don't affect the number of grouping sets or the behavior of "distinct" . We can compute all the candidate
        // grouping sets in terms of fieldId, dedup as appropriate and then cross-join them with the complex expressions.
        Analysis.GroupingSetAnalysis groupingSetAnalysis = analysis.getGroupingSets(node);
        columnOnlyGroupingSets = enumerateGroupingSets(groupingSetAnalysis);
        if (node.getGroupBy().get().isDistinct()) {
            columnOnlyGroupingSets = columnOnlyGroupingSets.stream().distinct().collect(toImmutableList());
        }
        // add in the complex expressions an turn materialize the grouping sets in terms of plan columns
        ImmutableList.Builder<List<VariableReferenceExpression>> groupingSetBuilder = ImmutableList.builder();
        for (Set<FieldId> groupingSet : columnOnlyGroupingSets) {
            ImmutableList.Builder<VariableReferenceExpression> columns = ImmutableList.builder();
            groupingSetAnalysis.getComplexExpressions().stream().map(groupingTranslations::get).forEach(columns::add);
            groupingSet.stream().map(field -> groupingTranslations.get(new FieldReference(field.getFieldIndex()))).forEach(columns::add);
            groupingSetBuilder.add(columns.build());
        }
        groupingSets = groupingSetBuilder.build();
    }
    // 2.c. Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
    Optional<VariableReferenceExpression> groupIdVariable = Optional.empty();
    if (groupingSets.size() > 1) {
        groupIdVariable = Optional.of(variableAllocator.newVariable("groupId", BIGINT));
        GroupIdNode groupId = new GroupIdNode(subPlan.getRoot().getSourceLocation(), idAllocator.getNextId(), subPlan.getRoot(), groupingSets, groupingSetMappings, aggregationArguments, groupIdVariable.get());
        subPlan = new PlanBuilder(groupingTranslations, groupId);
    } else {
        Assignments.Builder assignments = Assignments.builder();
        aggregationArguments.stream().map(AssignmentUtils::identityAsSymbolReference).forEach(assignments::put);
        groupingSetMappings.forEach((key, value) -> assignments.put(key, castToRowExpression(asSymbolReference(value))));
        ProjectNode project = new ProjectNode(subPlan.getRoot().getSourceLocation(), idAllocator.getNextId(), subPlan.getRoot(), assignments.build(), LOCAL);
        subPlan = new PlanBuilder(groupingTranslations, project);
    }
    TranslationMap aggregationTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis, lambdaDeclarationToVariableMap);
    aggregationTranslations.copyMappingsFrom(groupingTranslations);
    // 2.d. Rewrite aggregates
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregationsBuilder = ImmutableMap.builder();
    boolean needPostProjectionCoercion = false;
    for (FunctionCall aggregate : analysis.getAggregates(node)) {
        Expression rewritten = argumentTranslations.rewrite(aggregate);
        VariableReferenceExpression newVariable = variableAllocator.newVariable(rewritten, analysis.getType(aggregate));
        // Therefore we can end up with this implicit cast, and have to move it into a post-projection
        if (rewritten instanceof Cast) {
            rewritten = ((Cast) rewritten).getExpression();
            needPostProjectionCoercion = true;
        }
        aggregationTranslations.put(aggregate, newVariable);
        FunctionCall rewrittenFunction = (FunctionCall) rewritten;
        aggregationsBuilder.put(newVariable, new Aggregation(new CallExpression(getSourceLocation(rewrittenFunction), aggregate.getName().getSuffix(), analysis.getFunctionHandle(aggregate), analysis.getType(aggregate), rewrittenFunction.getArguments().stream().map(OriginalExpressionUtils::castToRowExpression).collect(toImmutableList())), rewrittenFunction.getFilter().map(OriginalExpressionUtils::castToRowExpression), rewrittenFunction.getOrderBy().map(orderBy -> toOrderingScheme(orderBy, variableAllocator.getTypes())), rewrittenFunction.isDistinct(), Optional.empty()));
    }
    Map<VariableReferenceExpression, Aggregation> aggregations = aggregationsBuilder.build();
    ImmutableSet.Builder<Integer> globalGroupingSets = ImmutableSet.builder();
    for (int i = 0; i < groupingSets.size(); i++) {
        if (groupingSets.get(i).isEmpty()) {
            globalGroupingSets.add(i);
        }
    }
    ImmutableList.Builder<VariableReferenceExpression> groupingKeys = ImmutableList.builder();
    groupingSets.stream().flatMap(List::stream).distinct().forEach(groupingKeys::add);
    groupIdVariable.ifPresent(groupingKeys::add);
    AggregationNode aggregationNode = new AggregationNode(subPlan.getRoot().getSourceLocation(), idAllocator.getNextId(), subPlan.getRoot(), aggregations, groupingSets(groupingKeys.build(), groupingSets.size(), globalGroupingSets.build()), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), groupIdVariable);
    subPlan = new PlanBuilder(aggregationTranslations, aggregationNode);
    // TODO: this is a hack, we should change type coercions to coerce the inputs to functions/operators instead of coercing the output
    if (needPostProjectionCoercion) {
        ImmutableList.Builder<Expression> alreadyCoerced = ImmutableList.builder();
        alreadyCoerced.addAll(groupByExpressions);
        groupIdVariable.map(ExpressionTreeUtils::createSymbolReference).ifPresent(alreadyCoerced::add);
        subPlan = explicitCoercionFields(subPlan, alreadyCoerced.build(), analysis.getAggregates(node));
    }
    // 4. Project and re-write all grouping functions
    return handleGroupingOperations(subPlan, node, groupIdVariable, columnOnlyGroupingSets);
}

Example 18 with AggregationNode

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

the class CanonicalPlanGenerator method visitAggregation.

@Override
public Optional<PlanNode> visitAggregation(AggregationNode node, Map<VariableReferenceExpression, VariableReferenceExpression> context) {
    Optional<PlanNode> source = node.getSource().accept(this, context);
    if (!source.isPresent()) {
        return Optional.empty();
    }
    // Steps to get canonical aggregations:
    // 1. Transform aggregation into canonical form
    // 2. Sort based on canonical aggregation expression
    // 3. Get new variable reference for aggregation expression
    // 4. Record mapping from original variable reference to the new one
    List<AggregationReference> aggregationReferences = node.getAggregations().entrySet().stream().map(entry -> new AggregationReference(getCanonicalAggregation(entry.getValue(), context), entry.getKey())).sorted(comparing(aggregationReference -> aggregationReference.getAggregation().getCall().toString())).collect(toImmutableList());
    ImmutableMap.Builder<VariableReferenceExpression, Aggregation> aggregations = ImmutableMap.builder();
    for (AggregationReference aggregationReference : aggregationReferences) {
        VariableReferenceExpression reference = variableAllocator.newVariable(aggregationReference.getAggregation().getCall());
        context.put(aggregationReference.getVariableReferenceExpression(), reference);
        aggregations.put(reference, aggregationReference.getAggregation());
    }
    return Optional.of(new AggregationNode(node.getSourceLocation(), planNodeidAllocator.getNextId(), source.get(), aggregations.build(), getCanonicalGroupingSetDescriptor(node.getGroupingSets(), context), node.getPreGroupedVariables().stream().map(context::get).collect(toImmutableList()), node.getStep(), node.getHashVariable().map(ignored -> variableAllocator.newHashVariable()), node.getGroupIdVariable().map(context::get)));
}

Example 19 with AggregationNode

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

the class TestDruidPlanOptimizer method testUnionAll.

@Test
public void testUnionAll() {
    PlanBuilder planBuilder = createPlanBuilder(defaultSessionHolder);
    PlanVariableAllocator variableAllocator = new PlanVariableAllocator();
    AggregationNode aggregationOne = simpleAggregationSum(planBuilder, tableScan(planBuilder, druidTableOne, city, fare), variableAllocator, ImmutableList.of(city), fare);
    AggregationNode aggregationTwo = simpleAggregationSum(planBuilder, tableScan(planBuilder, druidTableTwo, city, fare), variableAllocator, ImmutableList.of(city), fare);
    VariableReferenceExpression groupByColumn = variableAllocator.newVariable(city.getColumnName(), city.getColumnType());
    VariableReferenceExpression sumColumn = variableAllocator.newVariable(fare.getColumnName(), fare.getColumnType());
    PlanNode originalPlan = new UnionNode(Optional.empty(), planBuilder.getIdAllocator().getNextId(), ImmutableList.of(aggregationOne, aggregationTwo), ImmutableList.of(groupByColumn, sumColumn), ImmutableMap.of(groupByColumn, Stream.concat(aggregationOne.getGroupingKeys().stream(), aggregationTwo.getGroupingKeys().stream()).collect(toImmutableList()), sumColumn, ImmutableList.of(Iterables.getOnlyElement(aggregationOne.getAggregations().keySet()), Iterables.getOnlyElement(aggregationTwo.getAggregations().keySet()))));
    PlanNode optimizedPlan = getOptimizedPlan(planBuilder, originalPlan);
    PlanMatchPattern tableScanMatcherOne = DruidTableScanMatcher.match(druidTableOne.getTableName(), "SELECT \"city\", sum(fare) FROM \"realtimeOnly\" GROUP BY \"city\"");
    PlanMatchPattern tableScanMatcherTwo = DruidTableScanMatcher.match(druidTableTwo.getTableName(), "SELECT \"city\", sum(fare) FROM \"hybrid\" GROUP BY \"city\"");
    assertPlanMatch(optimizedPlan, PlanMatchPattern.union(tableScanMatcherOne, tableScanMatcherTwo), typeProvider);
}

Example 20 with AggregationNode

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

the class AggregationMatcher method detailMatches.

@Override
public MatchResult detailMatches(PlanNode node, StatsProvider stats, Session session, Metadata metadata, SymbolAliases symbolAliases) {
    checkState(shapeMatches(node), "Plan testing framework error: shapeMatches returned false in detailMatches in %s", this.getClass().getName());
    AggregationNode aggregationNode = (AggregationNode) node;
    if (groupId.isPresent() != aggregationNode.getGroupIdVariable().isPresent()) {
        return NO_MATCH;
    }
    if (!matches(groupingSets.getGroupingKeys(), aggregationNode.getGroupingKeys(), symbolAliases)) {
        return NO_MATCH;
    }
    if (groupingSets.getGroupingSetCount() != aggregationNode.getGroupingSetCount()) {
        return NO_MATCH;
    }
    if (!groupingSets.getGlobalGroupingSets().equals(aggregationNode.getGlobalGroupingSets())) {
        return NO_MATCH;
    }
    List<VariableReferenceExpression> aggregationsWithMask = aggregationNode.getAggregations().entrySet().stream().filter(entry -> entry.getValue().getMask().isPresent()).map(Map.Entry::getKey).collect(Collectors.toList());
    if (aggregationsWithMask.size() != masks.keySet().size()) {
        return NO_MATCH;
    }
    for (VariableReferenceExpression variable : aggregationsWithMask) {
        if (!masks.keySet().contains(new Symbol(variable.getName()))) {
            return NO_MATCH;
        }
    }
    if (step != aggregationNode.getStep()) {
        return NO_MATCH;
    }
    if (!matches(preGroupedSymbols, aggregationNode.getPreGroupedVariables(), symbolAliases)) {
        return NO_MATCH;
    }
    return match();
}