Examples with OrderBy io.trino.sql.tree.OrderBy used on opensource projects

Example 11 with OrderBy

use of io.trino.sql.tree.OrderBy in project trino by trinodb.

the class QueryPlanner method planAggregation.

private PlanBuilder planAggregation(PlanBuilder subPlan, List<List<Symbol>> groupingSets, Optional<Symbol> groupIdSymbol, List<FunctionCall> aggregates, Function<Expression, Symbol> coercions) {
    ImmutableList.Builder<AggregationAssignment> aggregateMappingBuilder = ImmutableList.builder();
    // deduplicate based on scope-aware equality
    for (FunctionCall function : scopeAwareDistinct(subPlan, aggregates)) {
        Symbol symbol = symbolAllocator.newSymbol(function, analysis.getType(function));
        // TODO: for ORDER BY arguments, rewrite them such that they match the actual arguments to the function. This is necessary to maintain the semantics of DISTINCT + ORDER BY,
        // which requires that ORDER BY be a subset of arguments
        // What can happen currently is that if the argument requires a coercion, the argument will take a different input that the ORDER BY clause, which is undefined behavior
        Aggregation aggregation = new Aggregation(analysis.getResolvedFunction(function), function.getArguments().stream().map(argument -> {
            if (argument instanceof LambdaExpression) {
                return subPlan.rewrite(argument);
            }
            return coercions.apply(argument).toSymbolReference();
        }).collect(toImmutableList()), function.isDistinct(), function.getFilter().map(coercions), function.getOrderBy().map(orderBy -> translateOrderingScheme(orderBy.getSortItems(), coercions)), Optional.empty());
        aggregateMappingBuilder.add(new AggregationAssignment(symbol, function, aggregation));
    }
    List<AggregationAssignment> aggregateMappings = aggregateMappingBuilder.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<Symbol> groupingKeys = ImmutableList.builder();
    groupingSets.stream().flatMap(List::stream).distinct().forEach(groupingKeys::add);
    groupIdSymbol.ifPresent(groupingKeys::add);
    AggregationNode aggregationNode = new AggregationNode(idAllocator.getNextId(), subPlan.getRoot(), aggregateMappings.stream().collect(toImmutableMap(AggregationAssignment::getSymbol, AggregationAssignment::getRewritten)), groupingSets(groupingKeys.build(), groupingSets.size(), globalGroupingSets.build()), ImmutableList.of(), AggregationNode.Step.SINGLE, Optional.empty(), groupIdSymbol);
    return new PlanBuilder(subPlan.getTranslations().withAdditionalMappings(aggregateMappings.stream().collect(toImmutableMap(assignment -> scopeAwareKey(assignment.getAstExpression(), analysis, subPlan.getScope()), AggregationAssignment::getSymbol))), aggregationNode);
}