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

Example 26 with Symbol

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

the class PruneCountAggregationOverScalar method apply.

@Override
public Result apply(AggregationNode parent, Captures captures, Context context) {
    if (!parent.hasDefaultOutput() || parent.getOutputSymbols().size() != 1) {
        return Result.empty();
    }
    FunctionId countFunctionId = metadata.resolveFunction(context.getSession(), QualifiedName.of("count"), ImmutableList.of()).getFunctionId();
    Map<Symbol, AggregationNode.Aggregation> assignments = parent.getAggregations();
    for (Map.Entry<Symbol, AggregationNode.Aggregation> entry : assignments.entrySet()) {
        AggregationNode.Aggregation aggregation = entry.getValue();
        requireNonNull(aggregation, "aggregation is null");
        ResolvedFunction resolvedFunction = aggregation.getResolvedFunction();
        if (!countFunctionId.equals(resolvedFunction.getFunctionId())) {
            return Result.empty();
        }
    }
    if (!assignments.isEmpty() && isScalar(parent.getSource(), context.getLookup())) {
        return Result.ofPlanNode(new ValuesNode(parent.getId(), parent.getOutputSymbols(), ImmutableList.of(new Row(ImmutableList.of(new GenericLiteral("BIGINT", "1"))))));
    }
    return Result.empty();
}

Example 27 with Symbol

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

the class PruneTableWriterSourceColumns method apply.

@Override
public Result apply(TableWriterNode tableWriterNode, Captures captures, Context context) {
    ImmutableSet.Builder<Symbol> requiredInputs = ImmutableSet.<Symbol>builder().addAll(tableWriterNode.getColumns());
    if (tableWriterNode.getPartitioningScheme().isPresent()) {
        PartitioningScheme partitioningScheme = tableWriterNode.getPartitioningScheme().get();
        partitioningScheme.getPartitioning().getColumns().forEach(requiredInputs::add);
        partitioningScheme.getHashColumn().ifPresent(requiredInputs::add);
    }
    if (tableWriterNode.getStatisticsAggregation().isPresent()) {
        StatisticAggregations aggregations = tableWriterNode.getStatisticsAggregation().get();
        requiredInputs.addAll(aggregations.getGroupingSymbols());
        aggregations.getAggregations().values().stream().map(SymbolsExtractor::extractUnique).forEach(requiredInputs::addAll);
    }
    return restrictChildOutputs(context.getIdAllocator(), tableWriterNode, requiredInputs.build()).map(Result::ofPlanNode).orElse(Result.empty());
}

Example 28 with Symbol

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

the class PruneWindowColumns method pushDownProjectOff.

@Override
protected Optional<PlanNode> pushDownProjectOff(Context context, WindowNode windowNode, Set<Symbol> referencedOutputs) {
    Map<Symbol, WindowNode.Function> referencedFunctions = Maps.filterKeys(windowNode.getWindowFunctions(), referencedOutputs::contains);
    if (referencedFunctions.isEmpty()) {
        return Optional.of(windowNode.getSource());
    }
    ImmutableSet.Builder<Symbol> referencedInputs = ImmutableSet.<Symbol>builder().addAll(windowNode.getSource().getOutputSymbols().stream().filter(referencedOutputs::contains).iterator()).addAll(windowNode.getPartitionBy());
    windowNode.getOrderingScheme().ifPresent(orderingScheme -> orderingScheme.getOrderBy().forEach(referencedInputs::add));
    windowNode.getHashSymbol().ifPresent(referencedInputs::add);
    for (WindowNode.Function windowFunction : referencedFunctions.values()) {
        referencedInputs.addAll(SymbolsExtractor.extractUnique(windowFunction));
    }
    PlanNode prunedWindowNode = new WindowNode(windowNode.getId(), restrictOutputs(context.getIdAllocator(), windowNode.getSource(), referencedInputs.build()).orElse(windowNode.getSource()), windowNode.getSpecification(), referencedFunctions, windowNode.getHashSymbol(), windowNode.getPrePartitionedInputs(), windowNode.getPreSortedOrderPrefix());
    if (prunedWindowNode.getOutputSymbols().size() == windowNode.getOutputSymbols().size()) {
        // Neither function pruning nor input pruning was successful.
        return Optional.empty();
    }
    return Optional.of(prunedWindowNode);
}

Example 29 with Symbol

use of io.trino.sql.planner.Symbol 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()));
}

Example 30 with Symbol

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

the class TransformCorrelatedDistinctAggregationWithProjection method apply.

@Override
public Result apply(CorrelatedJoinNode correlatedJoinNode, Captures captures, Context context) {
    // decorrelate nested plan
    PlanNodeDecorrelator decorrelator = new PlanNodeDecorrelator(plannerContext, context.getSymbolAllocator(), context.getLookup());
    Optional<PlanNodeDecorrelator.DecorrelatedNode> decorrelatedSource = decorrelator.decorrelateFilters(captures.get(AGGREGATION).getSource(), correlatedJoinNode.getCorrelation());
    if (decorrelatedSource.isEmpty()) {
        return Result.empty();
    }
    PlanNode source = decorrelatedSource.get().getNode();
    // assign unique id on correlated join's input. It will be used to distinguish between original input rows after join
    PlanNode inputWithUniqueId = new AssignUniqueId(context.getIdAllocator().getNextId(), correlatedJoinNode.getInput(), context.getSymbolAllocator().newSymbol("unique", BIGINT));
    JoinNode join = new JoinNode(context.getIdAllocator().getNextId(), JoinNode.Type.LEFT, inputWithUniqueId, source, ImmutableList.of(), inputWithUniqueId.getOutputSymbols(), source.getOutputSymbols(), false, decorrelatedSource.get().getCorrelatedPredicates(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(), Optional.empty());
    // restore aggregation
    AggregationNode aggregation = captures.get(AGGREGATION);
    aggregation = new AggregationNode(aggregation.getId(), join, aggregation.getAggregations(), singleGroupingSet(ImmutableList.<Symbol>builder().addAll(join.getLeftOutputSymbols()).addAll(aggregation.getGroupingKeys()).build()), ImmutableList.of(), aggregation.getStep(), Optional.empty(), Optional.empty());
    // restrict outputs and apply projection
    Set<Symbol> outputSymbols = new HashSet<>(correlatedJoinNode.getOutputSymbols());
    List<Symbol> expectedAggregationOutputs = aggregation.getOutputSymbols().stream().filter(outputSymbols::contains).collect(toImmutableList());
    Assignments assignments = Assignments.builder().putIdentities(expectedAggregationOutputs).putAll(captures.get(PROJECTION).getAssignments()).build();
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), aggregation, assignments));
}