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

Example 1 with PlanNodeIdAllocator

use of io.trino.sql.planner.PlanNodeIdAllocator 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 PlanNodeIdAllocator

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

the class SqlQueryExecution method doPlanQuery.

private PlanRoot doPlanQuery() {
    // plan query
    PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
    LogicalPlanner logicalPlanner = new LogicalPlanner(stateMachine.getSession(), planOptimizers, idAllocator, plannerContext, typeAnalyzer, statsCalculator, costCalculator, stateMachine.getWarningCollector());
    Plan plan = logicalPlanner.plan(analysis);
    queryPlan.set(plan);
    // fragment the plan
    SubPlan fragmentedPlan = planFragmenter.createSubPlans(stateMachine.getSession(), plan, false, stateMachine.getWarningCollector());
    // extract inputs
    List<Input> inputs = new InputExtractor(plannerContext.getMetadata(), stateMachine.getSession()).extractInputs(fragmentedPlan);
    stateMachine.setInputs(inputs);
    stateMachine.setOutput(analysis.getTarget());
    boolean explainAnalyze = analysis.getStatement() instanceof ExplainAnalyze;
    return new PlanRoot(fragmentedPlan, !explainAnalyze);
}

Example 3 with PlanNodeIdAllocator

use of io.trino.sql.planner.PlanNodeIdAllocator 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 PlanNodeIdAllocator

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

the class TestDetermineJoinDistributionType method testGetSourceTablesSizeInBytes.

@Test
public void testGetSourceTablesSizeInBytes() {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), tester.getMetadata(), tester.getSession());
    Symbol symbol = planBuilder.symbol("col");
    Symbol sourceSymbol1 = planBuilder.symbol("source1");
    Symbol sourceSymbol2 = planBuilder.symbol("soruce2");
    // missing source stats
    assertEquals(getSourceTablesSizeInBytes(planBuilder.values(symbol), noLookup(), node -> PlanNodeStatsEstimate.unknown(), planBuilder.getTypes()), NaN);
    // two source plan nodes
    PlanNodeStatsEstimate sourceStatsEstimate1 = PlanNodeStatsEstimate.builder().setOutputRowCount(10).build();
    PlanNodeStatsEstimate sourceStatsEstimate2 = PlanNodeStatsEstimate.builder().setOutputRowCount(20).build();
    assertEquals(getSourceTablesSizeInBytes(planBuilder.union(ImmutableListMultimap.<Symbol, Symbol>builder().put(symbol, sourceSymbol1).put(symbol, sourceSymbol2).build(), ImmutableList.of(planBuilder.tableScan(ImmutableList.of(sourceSymbol1), ImmutableMap.of(sourceSymbol1, new TestingColumnHandle("col"))), planBuilder.values(new PlanNodeId("valuesNode"), sourceSymbol2))), noLookup(), node -> {
        if (node instanceof TableScanNode) {
            return sourceStatsEstimate1;
        }
        if (node instanceof ValuesNode) {
            return sourceStatsEstimate2;
        }
        return PlanNodeStatsEstimate.unknown();
    }, planBuilder.getTypes()), 270.0);
    // join node
    assertEquals(getSourceTablesSizeInBytes(planBuilder.join(INNER, planBuilder.values(sourceSymbol1), planBuilder.values(sourceSymbol2)), noLookup(), node -> sourceStatsEstimate1, planBuilder.getTypes()), NaN);
}

Example 5 with PlanNodeIdAllocator

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

the class TestEliminateCrossJoins method testDoNotReorderCrossJoins.

@Test
public void testDoNotReorderCrossJoins() {
    Session session = testSessionBuilder().build();
    PlanNode plan = joinNode(joinNode(values("a"), values("b")), values("c"), "b", "c");
    JoinGraph joinGraph = JoinGraph.buildFrom(tester().getPlannerContext(), plan, noLookup(), new PlanNodeIdAllocator(), session, createTestingTypeAnalyzer(tester().getPlannerContext()), TypeProvider.empty());
    assertEquals(getJoinOrder(joinGraph), ImmutableList.of(0, 1, 2));
}