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

Example 6 with AggregationNode

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

the class QueryPlanner method aggregate.

private PlanBuilder aggregate(PlanBuilder subPlan, QuerySpecification node) {
    List<List<Expression>> groupingSets = analysis.getGroupingSets(node);
    if (groupingSets.isEmpty()) {
        return subPlan;
    }
    // 1. Pre-project all scalar inputs (arguments and non-trivial group by expressions)
    Set<Expression> distinctGroupingColumns = groupingSets.stream().flatMap(Collection::stream).collect(toImmutableSet());
    ImmutableList.Builder<Expression> arguments = ImmutableList.builder();
    analysis.getAggregates(node).stream().map(FunctionCall::getArguments).flatMap(List::stream).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(distinctGroupingColumns, 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, lambdaDeclarationToSymbolMap);
    ImmutableMap.Builder<Symbol, Symbol> argumentMappingBuilder = ImmutableMap.builder();
    for (Expression argument : arguments.build()) {
        Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), argument);
        argumentTranslations.addIntermediateMapping(argument, parametersReplaced);
        Symbol input = subPlan.translate(parametersReplaced);
        if (!argumentTranslations.containsSymbol(parametersReplaced)) {
            Symbol output = symbolAllocator.newSymbol(parametersReplaced, analysis.getTypeWithCoercions(parametersReplaced), "arg");
            argumentMappingBuilder.put(output, input);
            argumentTranslations.put(parametersReplaced, output);
        }
    }
    Map<Symbol, Symbol> argumentMappings = argumentMappingBuilder.build();
    // 2.b. Rewrite grouping columns
    TranslationMap groupingTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis, lambdaDeclarationToSymbolMap);
    Map<Symbol, Symbol> groupingSetMappings = new HashMap<>();
    List<List<Symbol>> groupingSymbols = new ArrayList<>();
    for (List<Expression> groupingSet : groupingSets) {
        ImmutableList.Builder<Symbol> symbols = ImmutableList.builder();
        for (Expression expression : groupingSet) {
            Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), expression);
            groupingTranslations.addIntermediateMapping(expression, parametersReplaced);
            Symbol input = subPlan.translate(expression);
            Symbol output;
            if (!groupingTranslations.containsSymbol(parametersReplaced)) {
                output = symbolAllocator.newSymbol(parametersReplaced, analysis.getTypeWithCoercions(expression), "gid");
                groupingTranslations.put(parametersReplaced, output);
            } else {
                output = groupingTranslations.get(parametersReplaced);
            }
            groupingSetMappings.put(output, input);
            symbols.add(output);
        }
        groupingSymbols.add(symbols.build());
    }
    // 2.c. Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
    Optional<Symbol> groupIdSymbol = Optional.empty();
    if (groupingSets.size() > 1) {
        groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
        GroupIdNode groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSymbols, groupingSetMappings, argumentMappings, groupIdSymbol.get());
        subPlan = new PlanBuilder(groupingTranslations, groupId, analysis.getParameters());
    } else {
        Assignments.Builder assignments = Assignments.builder();
        for (Symbol output : argumentMappings.keySet()) {
            assignments.put(output, argumentMappings.get(output).toSymbolReference());
        }
        for (Symbol output : groupingSetMappings.keySet()) {
            assignments.put(output, groupingSetMappings.get(output).toSymbolReference());
        }
        ProjectNode project = new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), assignments.build());
        subPlan = new PlanBuilder(groupingTranslations, project, analysis.getParameters());
    }
    TranslationMap aggregationTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis, lambdaDeclarationToSymbolMap);
    aggregationTranslations.copyMappingsFrom(groupingTranslations);
    // 2.d. Rewrite aggregates
    ImmutableMap.Builder<Symbol, FunctionCall> aggregationAssignments = ImmutableMap.builder();
    ImmutableMap.Builder<Symbol, Signature> functions = ImmutableMap.builder();
    boolean needPostProjectionCoercion = false;
    for (FunctionCall aggregate : analysis.getAggregates(node)) {
        Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), aggregate);
        aggregationTranslations.addIntermediateMapping(aggregate, parametersReplaced);
        Expression rewritten = argumentTranslations.rewrite(parametersReplaced);
        Symbol newSymbol = symbolAllocator.newSymbol(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;
        }
        aggregationAssignments.put(newSymbol, (FunctionCall) rewritten);
        aggregationTranslations.put(parametersReplaced, newSymbol);
        functions.put(newSymbol, analysis.getFunctionSignature(aggregate));
    }
    // 2.e. Mark distinct rows for each aggregate that has DISTINCT
    // Map from aggregate function arguments to marker symbols, so that we can reuse the markers, if two aggregates have the same argument
    Map<Set<Expression>, Symbol> argumentMarkers = new HashMap<>();
    // Map from aggregate functions to marker symbols
    Map<Symbol, Symbol> masks = new HashMap<>();
    for (FunctionCall aggregate : Iterables.filter(analysis.getAggregates(node), FunctionCall::isDistinct)) {
        Set<Expression> args = ImmutableSet.copyOf(aggregate.getArguments());
        Symbol marker = argumentMarkers.get(args);
        Symbol aggregateSymbol = aggregationTranslations.get(aggregate);
        if (marker == null) {
            if (args.size() == 1) {
                marker = symbolAllocator.newSymbol(getOnlyElement(args), BOOLEAN, "distinct");
            } else {
                marker = symbolAllocator.newSymbol(aggregateSymbol.getName(), BOOLEAN, "distinct");
            }
            argumentMarkers.put(args, marker);
        }
        masks.put(aggregateSymbol, marker);
    }
    for (Map.Entry<Set<Expression>, Symbol> entry : argumentMarkers.entrySet()) {
        ImmutableList.Builder<Symbol> builder = ImmutableList.builder();
        builder.addAll(groupingSymbols.stream().flatMap(Collection::stream).distinct().collect(Collectors.toList()));
        groupIdSymbol.ifPresent(builder::add);
        for (Expression expression : entry.getKey()) {
            builder.add(argumentTranslations.get(expression));
        }
        subPlan = subPlan.withNewRoot(new MarkDistinctNode(idAllocator.getNextId(), subPlan.getRoot(), entry.getValue(), builder.build(), Optional.empty()));
    }
    AggregationNode aggregationNode = new AggregationNode(idAllocator.getNextId(), subPlan.getRoot(), aggregationAssignments.build(), functions.build(), masks, groupingSymbols, AggregationNode.Step.SINGLE, Optional.empty(), groupIdSymbol);
    subPlan = new PlanBuilder(aggregationTranslations, aggregationNode, analysis.getParameters());
    // TODO: this is a hack, we should change type coercions to coerce the inputs to functions/operators instead of coercing the output
    if (needPostProjectionCoercion) {
        return explicitCoercionFields(subPlan, distinctGroupingColumns, analysis.getAggregates(node));
    }
    return subPlan;
}

Example 7 with AggregationNode

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

the class SingleMarkDistinctToGroupBy method apply.

@Override
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator) {
    if (!(node instanceof AggregationNode)) {
        return Optional.empty();
    }
    AggregationNode parent = (AggregationNode) node;
    PlanNode source = lookup.resolve(parent.getSource());
    if (!(source instanceof MarkDistinctNode)) {
        return Optional.empty();
    }
    MarkDistinctNode child = (MarkDistinctNode) source;
    boolean hasFilters = parent.getAggregations().values().stream().map(FunctionCall::getFilter).anyMatch(Optional::isPresent);
    if (hasFilters) {
        return Optional.empty();
    }
    // optimize if and only if
    // all aggregation functions have a single common distinct mask symbol
    // AND all aggregation functions have mask
    Set<Symbol> masks = ImmutableSet.copyOf(parent.getMasks().values());
    if (masks.size() != 1 || parent.getMasks().size() != parent.getAggregations().size()) {
        return Optional.empty();
    }
    Symbol mask = Iterables.getOnlyElement(masks);
    if (!child.getMarkerSymbol().equals(mask)) {
        return Optional.empty();
    }
    return Optional.of(new AggregationNode(idAllocator.getNextId(), new AggregationNode(idAllocator.getNextId(), child.getSource(), Collections.emptyMap(), ImmutableList.of(child.getDistinctSymbols()), SINGLE, child.getHashSymbol(), Optional.empty()), // remove DISTINCT flag from function calls
    parent.getAssignments().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> removeDistinct(e.getValue()))), parent.getGroupingSets(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol()));
}

Example 8 with AggregationNode

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

the class TransformExistsApplyToScalarApply method apply.

@Override
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator) {
    if (!(node instanceof ApplyNode)) {
        return Optional.empty();
    }
    ApplyNode parent = (ApplyNode) node;
    if (parent.getSubqueryAssignments().size() != 1) {
        return Optional.empty();
    }
    Expression expression = getOnlyElement(parent.getSubqueryAssignments().getExpressions());
    if (!(expression instanceof ExistsPredicate)) {
        return Optional.empty();
    }
    Symbol count = symbolAllocator.newSymbol(COUNT.toString(), BIGINT);
    Symbol exists = getOnlyElement(parent.getSubqueryAssignments().getSymbols());
    return Optional.of(new ApplyNode(node.getId(), parent.getInput(), new ProjectNode(idAllocator.getNextId(), new AggregationNode(idAllocator.getNextId(), new LimitNode(idAllocator.getNextId(), parent.getSubquery(), 1, false), ImmutableMap.of(count, COUNT_CALL), ImmutableMap.of(count, countSignature), ImmutableMap.of(), ImmutableList.of(ImmutableList.of()), AggregationNode.Step.SINGLE, Optional.empty(), Optional.empty()), Assignments.of(exists, new ComparisonExpression(GREATER_THAN, count.toSymbolReference(), new Cast(new LongLiteral("0"), BIGINT.toString())))), Assignments.of(exists, exists.toSymbolReference()), parent.getCorrelation()));
}

Example 9 with AggregationNode

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

the class ImplementFilteredAggregations method apply.

@Override
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator) {
    if (!(node instanceof AggregationNode)) {
        return Optional.empty();
    }
    AggregationNode aggregation = (AggregationNode) node;
    boolean hasFilters = aggregation.getAggregations().entrySet().stream().anyMatch(e -> e.getValue().getFilter().isPresent() && // can't handle filtered aggregations with DISTINCT (conservatively, if they have a mask)
    !aggregation.getMasks().containsKey(e.getKey()));
    if (!hasFilters) {
        return Optional.empty();
    }
    Assignments.Builder newAssignments = Assignments.builder();
    ImmutableMap.Builder<Symbol, Symbol> masks = ImmutableMap.<Symbol, Symbol>builder().putAll(aggregation.getMasks());
    ImmutableMap.Builder<Symbol, FunctionCall> calls = ImmutableMap.builder();
    for (Map.Entry<Symbol, FunctionCall> entry : aggregation.getAggregations().entrySet()) {
        Symbol output = entry.getKey();
        // strip the filters
        FunctionCall call = entry.getValue();
        calls.put(output, new FunctionCall(call.getName(), call.getWindow(), Optional.empty(), call.isDistinct(), call.getArguments()));
        if (call.getFilter().isPresent()) {
            Expression filter = entry.getValue().getFilter().get();
            Symbol symbol = symbolAllocator.newSymbol(filter, BOOLEAN);
            newAssignments.put(symbol, filter);
            masks.put(output, symbol);
        }
    }
    // identity projection for all existing inputs
    newAssignments.putIdentities(aggregation.getSource().getOutputSymbols());
    return Optional.of(new AggregationNode(idAllocator.getNextId(), new ProjectNode(idAllocator.getNextId(), aggregation.getSource(), newAssignments.build()), calls.build(), aggregation.getFunctions(), masks.build(), aggregation.getGroupingSets(), aggregation.getStep(), aggregation.getHashSymbol(), aggregation.getGroupIdSymbol()));
}

Example 10 with AggregationNode

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

the class MergeLimitWithDistinct method apply.

@Override
public Optional<PlanNode> apply(PlanNode node, Lookup lookup, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator) {
    if (!(node instanceof LimitNode)) {
        return Optional.empty();
    }
    LimitNode parent = (LimitNode) node;
    PlanNode input = lookup.resolve(parent.getSource());
    if (!(input instanceof AggregationNode)) {
        return Optional.empty();
    }
    AggregationNode child = (AggregationNode) input;
    if (isDistinct(child)) {
        return Optional.empty();
    }
    return Optional.of(new DistinctLimitNode(parent.getId(), child.getSource(), parent.getCount(), false, child.getHashSymbol()));
}