Examples with GroupingSetAnalysis io.trino.sql.analyzer.Analysis.GroupingSetAnalysis used on opensource projects

Example 1 with GroupingSetAnalysis

use of io.trino.sql.analyzer.Analysis.GroupingSetAnalysis in project trino by trinodb.

the class QueryPlanner method planGroupingSets.

private GroupingSetsPlan planGroupingSets(PlanBuilder subPlan, QuerySpecification node, GroupingSetAnalysis groupingSetAnalysis) {
    Map<Symbol, Symbol> groupingSetMappings = new LinkedHashMap<>();
    // Compute a set of artificial columns that will contain the values of the original columns
    // filtered by whether the column is included in the grouping set
    // This will become the basis for the scope for any column references
    Symbol[] fields = new Symbol[subPlan.getTranslations().getFieldSymbols().size()];
    for (FieldId field : groupingSetAnalysis.getAllFields()) {
        Symbol input = subPlan.getTranslations().getFieldSymbols().get(field.getFieldIndex());
        Symbol output = symbolAllocator.newSymbol(input, "gid");
        fields[field.getFieldIndex()] = output;
        groupingSetMappings.put(output, input);
    }
    Map<ScopeAware<Expression>, Symbol> complexExpressions = new HashMap<>();
    for (Expression expression : groupingSetAnalysis.getComplexExpressions()) {
        if (!complexExpressions.containsKey(scopeAwareKey(expression, analysis, subPlan.getScope()))) {
            Symbol input = subPlan.translate(expression);
            Symbol output = symbolAllocator.newSymbol(expression, analysis.getType(expression), "gid");
            complexExpressions.put(scopeAwareKey(expression, analysis, subPlan.getScope()), output);
            groupingSetMappings.put(output, input);
        }
    }
    // 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.
    // This tracks the grouping sets before complex expressions are considered.
    // It's also used to compute the descriptors needed to implement grouping()
    List<Set<FieldId>> columnOnlyGroupingSets = enumerateGroupingSets(groupingSetAnalysis);
    if (node.getGroupBy().isPresent() && node.getGroupBy().get().isDistinct()) {
        columnOnlyGroupingSets = columnOnlyGroupingSets.stream().distinct().collect(toImmutableList());
    }
    // translate from FieldIds to Symbols
    List<List<Symbol>> sets = columnOnlyGroupingSets.stream().map(set -> set.stream().map(FieldId::getFieldIndex).map(index -> fields[index]).collect(toImmutableList())).collect(toImmutableList());
    // combine (cartesian product) with complex expressions
    List<List<Symbol>> groupingSets = sets.stream().map(set -> ImmutableList.<Symbol>builder().addAll(set).addAll(complexExpressions.values()).build()).collect(toImmutableList());
    // Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
    PlanNode groupId;
    Optional<Symbol> groupIdSymbol = Optional.empty();
    if (groupingSets.size() > 1) {
        groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
        groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSets, groupingSetMappings, subPlan.getRoot().getOutputSymbols(), groupIdSymbol.get());
    } else {
        Assignments.Builder assignments = Assignments.builder();
        assignments.putIdentities(subPlan.getRoot().getOutputSymbols());
        groupingSetMappings.forEach((key, value) -> assignments.put(key, value.toSymbolReference()));
        groupId = new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), assignments.build());
    }
    subPlan = new PlanBuilder(subPlan.getTranslations().withNewMappings(complexExpressions, Arrays.asList(fields)), groupId);
    return new GroupingSetsPlan(subPlan, columnOnlyGroupingSets, groupingSets, groupIdSymbol);
}

Example 2 with GroupingSetAnalysis

use of io.trino.sql.analyzer.Analysis.GroupingSetAnalysis in project trino by trinodb.

the class QueryPlanner method aggregate.

private PlanBuilder aggregate(PlanBuilder subPlan, QuerySpecification node) {
    if (!analysis.isAggregation(node)) {
        return subPlan;
    }
    ImmutableList.Builder<Expression> inputBuilder = ImmutableList.builder();
    analysis.getAggregates(node).stream().map(FunctionCall::getArguments).flatMap(List::stream).filter(// lambda expression is generated at execution time
    expression -> !(expression instanceof LambdaExpression)).forEach(inputBuilder::add);
    analysis.getAggregates(node).stream().map(FunctionCall::getOrderBy).map(NodeUtils::getSortItemsFromOrderBy).flatMap(List::stream).map(SortItem::getSortKey).forEach(inputBuilder::add);
    // filter expressions need to be projected first
    analysis.getAggregates(node).stream().map(FunctionCall::getFilter).filter(Optional::isPresent).map(Optional::get).forEach(inputBuilder::add);
    GroupingSetAnalysis groupingSetAnalysis = analysis.getGroupingSets(node);
    inputBuilder.addAll(groupingSetAnalysis.getComplexExpressions());
    List<Expression> inputs = inputBuilder.build();
    subPlan = subqueryPlanner.handleSubqueries(subPlan, inputs, analysis.getSubqueries(node));
    subPlan = subPlan.appendProjections(inputs, symbolAllocator, idAllocator);
    // Add projection to coerce inputs to their site-specific types.
    // This is important because the same lexical expression may need to be coerced
    // in different ways if it's referenced by multiple arguments to the window function.
    // For example, given v::integer,
    // avg(v)
    // Needs to be rewritten as
    // avg(CAST(v AS double))
    PlanAndMappings coercions = coerce(subPlan, inputs, analysis, idAllocator, symbolAllocator, typeCoercion);
    subPlan = coercions.getSubPlan();
    GroupingSetsPlan groupingSets = planGroupingSets(subPlan, node, groupingSetAnalysis);
    subPlan = planAggregation(groupingSets.getSubPlan(), groupingSets.getGroupingSets(), groupingSets.getGroupIdSymbol(), analysis.getAggregates(node), coercions::get);
    return planGroupingOperations(subPlan, node, groupingSets.getGroupIdSymbol(), groupingSets.getColumnOnlyGroupingSets());
}