Examples with Lookup io.trino.sql.planner.iterative.Lookup used on opensource projects

Example 1 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class AddIntermediateAggregations method apply.

@Override
public Result apply(AggregationNode aggregation, Captures captures, Context context) {
    Lookup lookup = context.getLookup();
    PlanNodeIdAllocator idAllocator = context.getIdAllocator();
    Session session = context.getSession();
    Optional<PlanNode> rewrittenSource = recurseToPartial(lookup.resolve(aggregation.getSource()), lookup, idAllocator);
    if (rewrittenSource.isEmpty()) {
        return Result.empty();
    }
    PlanNode source = rewrittenSource.get();
    if (getTaskConcurrency(session) > 1) {
        source = ExchangeNode.partitionedExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source, new PartitioningScheme(Partitioning.create(FIXED_ARBITRARY_DISTRIBUTION, ImmutableList.of()), source.getOutputSymbols()));
        source = new AggregationNode(idAllocator.getNextId(), source, inputsAsOutputs(aggregation.getAggregations()), aggregation.getGroupingSets(), aggregation.getPreGroupedSymbols(), AggregationNode.Step.INTERMEDIATE, aggregation.getHashSymbol(), aggregation.getGroupIdSymbol());
        source = ExchangeNode.gatheringExchange(idAllocator.getNextId(), ExchangeNode.Scope.LOCAL, source);
    }
    return Result.ofPlanNode(aggregation.replaceChildren(ImmutableList.of(source)));
}

Example 2 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class PushProjectionThroughJoin method inlineProjections.

private static PlanNode inlineProjections(PlannerContext plannerContext, ProjectNode parentProjection, Lookup lookup, Session session, TypeAnalyzer typeAnalyzer, TypeProvider types) {
    PlanNode child = lookup.resolve(parentProjection.getSource());
    if (!(child instanceof ProjectNode)) {
        return parentProjection;
    }
    ProjectNode childProjection = (ProjectNode) child;
    return InlineProjections.inlineProjections(plannerContext, parentProjection, childProjection, session, typeAnalyzer, types).map(node -> inlineProjections(plannerContext, node, lookup, session, typeAnalyzer, types)).orElse(parentProjection);
}

Example 3 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class PushProjectionThroughJoin method pushProjectionThroughJoin.

public static Optional<PlanNode> pushProjectionThroughJoin(PlannerContext plannerContext, ProjectNode projectNode, Lookup lookup, PlanNodeIdAllocator planNodeIdAllocator, Session session, TypeAnalyzer typeAnalyzer, TypeProvider types) {
    if (!projectNode.getAssignments().getExpressions().stream().allMatch(expression -> isDeterministic(expression, plannerContext.getMetadata()))) {
        return Optional.empty();
    }
    PlanNode child = lookup.resolve(projectNode.getSource());
    if (!(child instanceof JoinNode)) {
        return Optional.empty();
    }
    JoinNode joinNode = (JoinNode) child;
    PlanNode leftChild = joinNode.getLeft();
    PlanNode rightChild = joinNode.getRight();
    if (joinNode.getType() != INNER) {
        return Optional.empty();
    }
    Assignments.Builder leftAssignmentsBuilder = Assignments.builder();
    Assignments.Builder rightAssignmentsBuilder = Assignments.builder();
    for (Map.Entry<Symbol, Expression> assignment : projectNode.getAssignments().entrySet()) {
        Expression expression = assignment.getValue();
        Set<Symbol> symbols = extractUnique(expression);
        if (leftChild.getOutputSymbols().containsAll(symbols)) {
            // expression is satisfied with left child symbols
            leftAssignmentsBuilder.put(assignment.getKey(), expression);
        } else if (rightChild.getOutputSymbols().containsAll(symbols)) {
            // expression is satisfied with right child symbols
            rightAssignmentsBuilder.put(assignment.getKey(), expression);
        } else {
            // expression is using symbols from both join sides
            return Optional.empty();
        }
    }
    // add projections for symbols required by the join itself
    Set<Symbol> joinRequiredSymbols = getJoinRequiredSymbols(joinNode);
    for (Symbol requiredSymbol : joinRequiredSymbols) {
        if (leftChild.getOutputSymbols().contains(requiredSymbol)) {
            leftAssignmentsBuilder.putIdentity(requiredSymbol);
        } else {
            checkState(rightChild.getOutputSymbols().contains(requiredSymbol));
            rightAssignmentsBuilder.putIdentity(requiredSymbol);
        }
    }
    Assignments leftAssignments = leftAssignmentsBuilder.build();
    Assignments rightAssignments = rightAssignmentsBuilder.build();
    List<Symbol> leftOutputSymbols = leftAssignments.getOutputs().stream().filter(ImmutableSet.copyOf(projectNode.getOutputSymbols())::contains).collect(toImmutableList());
    List<Symbol> rightOutputSymbols = rightAssignments.getOutputs().stream().filter(ImmutableSet.copyOf(projectNode.getOutputSymbols())::contains).collect(toImmutableList());
    return Optional.of(new JoinNode(joinNode.getId(), joinNode.getType(), inlineProjections(plannerContext, new ProjectNode(planNodeIdAllocator.getNextId(), leftChild, leftAssignments), lookup, session, typeAnalyzer, types), inlineProjections(plannerContext, new ProjectNode(planNodeIdAllocator.getNextId(), rightChild, rightAssignments), lookup, session, typeAnalyzer, types), joinNode.getCriteria(), leftOutputSymbols, rightOutputSymbols, joinNode.isMaySkipOutputDuplicates(), joinNode.getFilter(), joinNode.getLeftHashSymbol(), joinNode.getRightHashSymbol(), joinNode.getDistributionType(), joinNode.isSpillable(), joinNode.getDynamicFilters(), joinNode.getReorderJoinStatsAndCost()));
}

Example 4 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class SetOperationMerge method mergeFirstSource.

/**
 * Only merge first source node. This method is assumed to be used for EXCEPT, which is a non-associative set operation.
 * Provides a correct plan transformation also for associative set operations: UNION and INTERSECT.
 *
 * @return Merged plan node if applied.
 */
public Optional<SetOperationNode> mergeFirstSource() {
    Lookup lookup = context.getLookup();
    List<PlanNode> sources = node.getSources().stream().flatMap(lookup::resolveGroup).collect(Collectors.toList());
    PlanNode child = sources.get(0);
    // Determine if set operations can be merged and whether the resulting set operation is quantified DISTINCT or ALL
    Optional<Boolean> mergedQuantifier = mergedQuantifierIsDistinct(node, child);
    if (mergedQuantifier.isEmpty()) {
        return Optional.empty();
    }
    ImmutableListMultimap.Builder<Symbol, Symbol> newMappingsBuilder = ImmutableListMultimap.builder();
    // Merge all sources of the first source.
    addMergedMappings((SetOperationNode) child, 0, newMappingsBuilder);
    // Keep remaining as it is
    for (int i = 1; i < sources.size(); i++) {
        PlanNode source = sources.get(i);
        addOriginalMappings(source, i, newMappingsBuilder);
    }
    if (node instanceof UnionNode) {
        return Optional.of(new UnionNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols()));
    }
    if (node instanceof IntersectNode) {
        return Optional.of(new IntersectNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), mergedQuantifier.get()));
    }
    if (node instanceof ExceptNode) {
        return Optional.of(new ExceptNode(node.getId(), newSources, newMappingsBuilder.build(), node.getOutputSymbols(), mergedQuantifier.get()));
    }
    throw new IllegalArgumentException("unexpected node type: " + node.getClass().getSimpleName());
}

Example 5 with Lookup

use of io.trino.sql.planner.iterative.Lookup in project trino by trinodb.

the class TableScanStatsRule method doCalculate.

@Override
protected Optional<PlanNodeStatsEstimate> doCalculate(TableScanNode node, StatsProvider sourceStats, Lookup lookup, Session session, TypeProvider types) {
    if (isStatisticsPrecalculationForPushdownEnabled(session) && node.getStatistics().isPresent()) {
        return node.getStatistics();
    }
    // TODO Construct predicate like AddExchanges's LayoutConstraintEvaluator
    Constraint constraint = new Constraint(TupleDomain.all());
    TableStatistics tableStatistics = metadata.getTableStatistics(session, node.getTable(), constraint);
    Map<Symbol, SymbolStatsEstimate> outputSymbolStats = new HashMap<>();
    for (Map.Entry<Symbol, ColumnHandle> entry : node.getAssignments().entrySet()) {
        Symbol symbol = entry.getKey();
        Optional<ColumnStatistics> columnStatistics = Optional.ofNullable(tableStatistics.getColumnStatistics().get(entry.getValue()));
        SymbolStatsEstimate symbolStatistics = columnStatistics.map(statistics -> toSymbolStatistics(tableStatistics, statistics, types.get(symbol))).orElse(SymbolStatsEstimate.unknown());
        outputSymbolStats.put(symbol, symbolStatistics);
    }
    return Optional.of(PlanNodeStatsEstimate.builder().setOutputRowCount(tableStatistics.getRowCount().getValue()).addSymbolStatistics(outputSymbolStats).build());
}