Examples with Aggregation io.trino.sql.planner.plan.AggregationNode.Aggregation used on opensource projects

Example 1 with Aggregation

use of io.trino.sql.planner.plan.AggregationNode.Aggregation in project trino by trinodb.

the class AggregationStatsRule method groupBy.

public static PlanNodeStatsEstimate groupBy(PlanNodeStatsEstimate sourceStats, Collection<Symbol> groupBySymbols, Map<Symbol, Aggregation> aggregations) {
    PlanNodeStatsEstimate.Builder result = PlanNodeStatsEstimate.builder();
    for (Symbol groupBySymbol : groupBySymbols) {
        SymbolStatsEstimate symbolStatistics = sourceStats.getSymbolStatistics(groupBySymbol);
        result.addSymbolStatistics(groupBySymbol, symbolStatistics.mapNullsFraction(nullsFraction -> {
            if (nullsFraction == 0.0) {
                return 0.0;
            }
            return 1.0 / (symbolStatistics.getDistinctValuesCount() + 1);
        }));
    }
    double rowsCount = getRowsCount(sourceStats, groupBySymbols);
    result.setOutputRowCount(min(rowsCount, sourceStats.getOutputRowCount()));
    for (Map.Entry<Symbol, Aggregation> aggregationEntry : aggregations.entrySet()) {
        result.addSymbolStatistics(aggregationEntry.getKey(), estimateAggregationStats(aggregationEntry.getValue(), sourceStats));
    }
    return result.build();
}

Example 2 with Aggregation

use of io.trino.sql.planner.plan.AggregationNode.Aggregation in project trino by trinodb.

the class ImplementFilteredAggregations method apply.

@Override
public Result apply(AggregationNode aggregationNode, Captures captures, Context context) {
    Assignments.Builder newAssignments = Assignments.builder();
    ImmutableMap.Builder<Symbol, Aggregation> aggregations = ImmutableMap.builder();
    ImmutableList.Builder<Expression> maskSymbols = ImmutableList.builder();
    boolean aggregateWithoutFilterOrMaskPresent = false;
    for (Map.Entry<Symbol, Aggregation> entry : aggregationNode.getAggregations().entrySet()) {
        Symbol output = entry.getKey();
        // strip the filters
        Aggregation aggregation = entry.getValue();
        Optional<Symbol> mask = aggregation.getMask();
        if (aggregation.getFilter().isPresent()) {
            Symbol filter = aggregation.getFilter().get();
            if (mask.isPresent()) {
                Symbol newMask = context.getSymbolAllocator().newSymbol("mask", BOOLEAN);
                Expression expression = and(mask.get().toSymbolReference(), filter.toSymbolReference());
                newAssignments.put(newMask, expression);
                mask = Optional.of(newMask);
                maskSymbols.add(newMask.toSymbolReference());
            } else {
                mask = Optional.of(filter);
                maskSymbols.add(filter.toSymbolReference());
            }
        } else if (mask.isPresent()) {
            maskSymbols.add(mask.get().toSymbolReference());
        } else {
            aggregateWithoutFilterOrMaskPresent = true;
        }
        aggregations.put(output, new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), aggregation.isDistinct(), Optional.empty(), aggregation.getOrderingScheme(), mask));
    }
    Expression predicate = TRUE_LITERAL;
    if (!aggregationNode.hasNonEmptyGroupingSet() && !aggregateWithoutFilterOrMaskPresent) {
        predicate = combineDisjunctsWithDefault(metadata, maskSymbols.build(), TRUE_LITERAL);
    }
    // identity projection for all existing inputs
    newAssignments.putIdentities(aggregationNode.getSource().getOutputSymbols());
    return Result.ofPlanNode(new AggregationNode(context.getIdAllocator().getNextId(), new FilterNode(context.getIdAllocator().getNextId(), new ProjectNode(context.getIdAllocator().getNextId(), aggregationNode.getSource(), newAssignments.build()), predicate), aggregations.buildOrThrow(), aggregationNode.getGroupingSets(), ImmutableList.of(), aggregationNode.getStep(), aggregationNode.getHashSymbol(), aggregationNode.getGroupIdSymbol()));
}

Example 3 with Aggregation

use of io.trino.sql.planner.plan.AggregationNode.Aggregation in project trino by trinodb.

the class AggregationDecorrelation method rewriteWithMasks.

public static Map<Symbol, Aggregation> rewriteWithMasks(Map<Symbol, Aggregation> aggregations, Map<Symbol, Symbol> masks) {
    ImmutableMap.Builder<Symbol, Aggregation> rewritten = ImmutableMap.builder();
    for (Map.Entry<Symbol, Aggregation> entry : aggregations.entrySet()) {
        Symbol symbol = entry.getKey();
        Aggregation aggregation = entry.getValue();
        rewritten.put(symbol, new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), aggregation.isDistinct(), aggregation.getFilter(), aggregation.getOrderingScheme(), Optional.of(masks.get(symbol))));
    }
    return rewritten.buildOrThrow();
}

Example 4 with Aggregation

use of io.trino.sql.planner.plan.AggregationNode.Aggregation in project trino by trinodb.

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<Symbol>, Symbol> markers = new HashMap<>();
    Map<Symbol, Aggregation> newAggregations = new HashMap<>();
    PlanNode subPlan = parent.getSource();
    for (Map.Entry<Symbol, Aggregation> entry : parent.getAggregations().entrySet()) {
        Aggregation aggregation = entry.getValue();
        if (aggregation.isDistinct() && aggregation.getFilter().isEmpty() && aggregation.getMask().isEmpty()) {
            Set<Symbol> inputs = aggregation.getArguments().stream().map(Symbol::from).collect(toSet());
            Symbol marker = markers.get(inputs);
            if (marker == null) {
                marker = context.getSymbolAllocator().newSymbol(Iterables.getLast(inputs).getName(), BOOLEAN, "distinct");
                markers.put(inputs, marker);
                ImmutableSet.Builder<Symbol> distinctSymbols = ImmutableSet.<Symbol>builder().addAll(parent.getGroupingKeys()).addAll(inputs);
                parent.getGroupIdSymbol().ifPresent(distinctSymbols::add);
                subPlan = new MarkDistinctNode(context.getIdAllocator().getNextId(), subPlan, marker, ImmutableList.copyOf(distinctSymbols.build()), Optional.empty());
            }
            // remove the distinct flag and set the distinct marker
            newAggregations.put(entry.getKey(), new Aggregation(aggregation.getResolvedFunction(), aggregation.getArguments(), false, aggregation.getFilter(), aggregation.getOrderingScheme(), Optional.of(marker)));
        } else {
            newAggregations.put(entry.getKey(), aggregation);
        }
    }
    return Result.ofPlanNode(new AggregationNode(parent.getId(), subPlan, newAggregations, parent.getGroupingSets(), ImmutableList.of(), parent.getStep(), parent.getHashSymbol(), parent.getGroupIdSymbol()));
}

Example 5 with Aggregation

use of io.trino.sql.planner.plan.AggregationNode.Aggregation in project trino by trinodb.

the class SingleDistinctAggregationToGroupBy method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    List<Set<Expression>> argumentSets = extractArgumentSets(aggregation).collect(Collectors.toList());
    Set<Symbol> symbols = Iterables.getOnlyElement(argumentSets).stream().map(Symbol::from).collect(Collectors.toSet());
    return Result.ofPlanNode(new AggregationNode(aggregation.getId(), new AggregationNode(context.getIdAllocator().getNextId(), aggregation.getSource(), ImmutableMap.of(), singleGroupingSet(ImmutableList.<Symbol>builder().addAll(aggregation.getGroupingKeys()).addAll(symbols).build()), ImmutableList.of(), SINGLE, Optional.empty(), Optional.empty()), // remove DISTINCT flag from function calls
    aggregation.getAggregations().entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> removeDistinct(e.getValue()))), aggregation.getGroupingSets(), emptyList(), aggregation.getStep(), aggregation.getHashSymbol(), aggregation.getGroupIdSymbol()));
}