Examples with Symbol io.trino.sql.planner.Symbol used on opensource projects

Example 16 with Symbol

use of io.trino.sql.planner.Symbol 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 17 with Symbol

use of io.trino.sql.planner.Symbol 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 18 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class PruneApplyColumns method pushDownProjectOff.

@Override
protected Optional<PlanNode> pushDownProjectOff(Context context, ApplyNode applyNode, Set<Symbol> referencedOutputs) {
    // remove unused apply node
    if (intersection(applyNode.getSubqueryAssignments().getSymbols(), referencedOutputs).isEmpty()) {
        return Optional.of(applyNode.getInput());
    }
    // extract referenced assignments
    ImmutableSet.Builder<Symbol> requiredAssignmentsSymbols = ImmutableSet.builder();
    Assignments.Builder newSubqueryAssignments = Assignments.builder();
    for (Map.Entry<Symbol, Expression> entry : applyNode.getSubqueryAssignments().entrySet()) {
        if (referencedOutputs.contains(entry.getKey())) {
            requiredAssignmentsSymbols.addAll(extractUnique(entry.getValue()));
            newSubqueryAssignments.put(entry);
        }
    }
    // prune subquery symbols
    Optional<PlanNode> newSubquery = restrictOutputs(context.getIdAllocator(), applyNode.getSubquery(), requiredAssignmentsSymbols.build());
    // prune input symbols
    Set<Symbol> requiredInputSymbols = ImmutableSet.<Symbol>builder().addAll(referencedOutputs).addAll(applyNode.getCorrelation()).addAll(requiredAssignmentsSymbols.build()).build();
    Optional<PlanNode> newInput = restrictOutputs(context.getIdAllocator(), applyNode.getInput(), requiredInputSymbols);
    boolean pruned = newSubquery.isPresent() || newInput.isPresent() || newSubqueryAssignments.build().size() < applyNode.getSubqueryAssignments().size();
    if (pruned) {
        return Optional.of(new ApplyNode(applyNode.getId(), newInput.orElse(applyNode.getInput()), newSubquery.orElse(applyNode.getSubquery()), newSubqueryAssignments.build(), applyNode.getCorrelation(), applyNode.getOriginSubquery()));
    }
    return Optional.empty();
}

Example 19 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class PushLimitThroughProject method apply.

@Override
public Result apply(LimitNode parent, Captures captures, Context context) {
    ProjectNode projectNode = captures.get(CHILD);
    // Do not push down if the projection is made up of symbol references and exclusive dereferences. This prevents
    // undoing of PushDownDereferencesThroughLimit. We still push limit in the case of overlapping dereferences since
    // it enables PushDownDereferencesThroughLimit rule to push optimal dereferences.
    Set<Expression> projections = ImmutableSet.copyOf(projectNode.getAssignments().getExpressions());
    if (!extractRowSubscripts(projections, false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes()).isEmpty() && exclusiveDereferences(projections, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes())) {
        return Result.empty();
    }
    // for a LimitNode without ties and pre-sorted inputs, simply reorder the nodes
    if (!parent.isWithTies() && !parent.requiresPreSortedInputs()) {
        return Result.ofPlanNode(transpose(parent, projectNode));
    }
    // for a LimitNode with ties, the tiesResolvingScheme must be rewritten in terms of symbols before projection
    SymbolMapper.Builder symbolMapper = SymbolMapper.builder();
    Set<Symbol> symbolsForRewrite = ImmutableSet.<Symbol>builder().addAll(parent.getPreSortedInputs()).addAll(parent.getTiesResolvingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of())).build();
    for (Symbol symbol : symbolsForRewrite) {
        Expression expression = projectNode.getAssignments().get(symbol);
        // if a symbol results from some computation, the translation fails
        if (!(expression instanceof SymbolReference)) {
            return Result.empty();
        }
        symbolMapper.put(symbol, Symbol.from(expression));
    }
    LimitNode mappedLimitNode = symbolMapper.build().map(parent, projectNode.getSource());
    return Result.ofPlanNode(projectNode.replaceChildren(ImmutableList.of(mappedLimitNode)));
}

Example 20 with Symbol

use of io.trino.sql.planner.Symbol in project trino by trinodb.

the class PushPartialAggregationThroughExchange method split.

private PlanNode split(AggregationNode node, Context context) {
    // otherwise, add a partial and final with an exchange in between
    Map<Symbol, AggregationNode.Aggregation> intermediateAggregation = new HashMap<>();
    Map<Symbol, AggregationNode.Aggregation> finalAggregation = new HashMap<>();
    for (Map.Entry<Symbol, AggregationNode.Aggregation> entry : node.getAggregations().entrySet()) {
        AggregationNode.Aggregation originalAggregation = entry.getValue();
        ResolvedFunction resolvedFunction = originalAggregation.getResolvedFunction();
        AggregationFunctionMetadata functionMetadata = plannerContext.getMetadata().getAggregationFunctionMetadata(resolvedFunction);
        List<Type> intermediateTypes = functionMetadata.getIntermediateTypes().stream().map(plannerContext.getTypeManager()::getType).collect(toImmutableList());
        Type intermediateType = intermediateTypes.size() == 1 ? intermediateTypes.get(0) : RowType.anonymous(intermediateTypes);
        Symbol intermediateSymbol = context.getSymbolAllocator().newSymbol(resolvedFunction.getSignature().getName(), intermediateType);
        checkState(originalAggregation.getOrderingScheme().isEmpty(), "Aggregate with ORDER BY does not support partial aggregation");
        intermediateAggregation.put(intermediateSymbol, new AggregationNode.Aggregation(resolvedFunction, originalAggregation.getArguments(), originalAggregation.isDistinct(), originalAggregation.getFilter(), originalAggregation.getOrderingScheme(), originalAggregation.getMask()));
        // rewrite final aggregation in terms of intermediate function
        finalAggregation.put(entry.getKey(), new AggregationNode.Aggregation(resolvedFunction, ImmutableList.<Expression>builder().add(intermediateSymbol.toSymbolReference()).addAll(originalAggregation.getArguments().stream().filter(LambdaExpression.class::isInstance).collect(toImmutableList())).build(), false, Optional.empty(), Optional.empty(), Optional.empty()));
    }
    PlanNode partial = new AggregationNode(context.getIdAllocator().getNextId(), node.getSource(), intermediateAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
    ImmutableList.of(), PARTIAL, node.getHashSymbol(), node.getGroupIdSymbol());
    return new AggregationNode(node.getId(), partial, finalAggregation, node.getGroupingSets(), // through the exchange may or may not preserve these properties. Hence, it is safest to drop preGroupedSymbols here.
    ImmutableList.of(), FINAL, node.getHashSymbol(), node.getGroupIdSymbol());
}