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

Example 1 with EnforceSingleRowNode

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

the class SubqueryPlanner method planScalarSubquery.

private PlanBuilder planScalarSubquery(PlanBuilder subPlan, Cluster<SubqueryExpression> cluster) {
    // Plan one of the predicates from the cluster
    SubqueryExpression scalarSubquery = cluster.getRepresentative();
    RelationPlan relationPlan = planSubquery(scalarSubquery, subPlan.getTranslations());
    PlanBuilder subqueryPlan = newPlanBuilder(relationPlan, analysis, lambdaDeclarationToSymbolMap);
    PlanNode root = new EnforceSingleRowNode(idAllocator.getNextId(), subqueryPlan.getRoot());
    Type type = analysis.getType(scalarSubquery);
    RelationType descriptor = relationPlan.getDescriptor();
    List<Symbol> fieldMappings = relationPlan.getFieldMappings();
    Symbol column;
    if (descriptor.getVisibleFieldCount() > 1) {
        column = symbolAllocator.newSymbol("row", type);
        ImmutableList.Builder<Expression> fields = ImmutableList.builder();
        for (int i = 0; i < descriptor.getAllFieldCount(); i++) {
            Field field = descriptor.getFieldByIndex(i);
            if (!field.isHidden()) {
                fields.add(fieldMappings.get(i).toSymbolReference());
            }
        }
        Expression expression = new Cast(new Row(fields.build()), TypeSignatureTranslator.toSqlType(type));
        root = new ProjectNode(idAllocator.getNextId(), root, Assignments.of(column, expression));
    } else {
        column = getOnlyElement(fieldMappings);
    }
    return appendCorrelatedJoin(subPlan, root, scalarSubquery.getQuery(), CorrelatedJoinNode.Type.INNER, TRUE_LITERAL, mapAll(cluster, subPlan.getScope(), column));
}

Example 2 with EnforceSingleRowNode

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

the class DecorrelateUnnest method apply.

@Override
public Result apply(CorrelatedJoinNode correlatedJoinNode, Captures captures, Context context) {
    // determine shape of the subquery
    PlanNode searchRoot = correlatedJoinNode.getSubquery();
    // 1. find EnforceSingleRowNode in the subquery
    Optional<EnforceSingleRowNode> enforceSingleRow = PlanNodeSearcher.searchFrom(searchRoot, context.getLookup()).where(EnforceSingleRowNode.class::isInstance).recurseOnlyWhen(planNode -> false).findFirst();
    if (enforceSingleRow.isPresent()) {
        searchRoot = enforceSingleRow.get().getSource();
    }
    // 2. find correlated UnnestNode in the subquery
    Optional<UnnestNode> subqueryUnnest = PlanNodeSearcher.searchFrom(searchRoot, context.getLookup()).where(node -> isSupportedUnnest(node, correlatedJoinNode.getCorrelation(), context.getLookup())).recurseOnlyWhen(node -> node instanceof ProjectNode || (node instanceof LimitNode && ((LimitNode) node).getCount() > 0) || (node instanceof TopNNode && ((TopNNode) node).getCount() > 0)).findFirst();
    if (subqueryUnnest.isEmpty()) {
        return Result.empty();
    }
    UnnestNode unnestNode = subqueryUnnest.get();
    // assign unique id to input rows
    Symbol uniqueSymbol = context.getSymbolAllocator().newSymbol("unique", BIGINT);
    PlanNode input = new AssignUniqueId(context.getIdAllocator().getNextId(), correlatedJoinNode.getInput(), uniqueSymbol);
    // pre-project unnest symbols if they were pre-projected in subquery
    // The correlated UnnestNode either unnests correlation symbols directly, or unnests symbols produced by a projection that uses only correlation symbols.
    // Here, any underlying projection that was a source of the correlated UnnestNode, is appended as a source of the rewritten UnnestNode.
    // If the projection is not necessary for UnnestNode (i.e. it does not produce any unnest symbols), it should be pruned afterwards.
    PlanNode unnestSource = context.getLookup().resolve(unnestNode.getSource());
    if (unnestSource instanceof ProjectNode) {
        ProjectNode sourceProjection = (ProjectNode) unnestSource;
        input = new ProjectNode(sourceProjection.getId(), input, Assignments.builder().putIdentities(input.getOutputSymbols()).putAll(sourceProjection.getAssignments()).build());
    }
    // determine join type for rewritten UnnestNode
    Type unnestJoinType = LEFT;
    if (enforceSingleRow.isEmpty() && correlatedJoinNode.getType() == CorrelatedJoinNode.Type.INNER && unnestNode.getJoinType() == INNER) {
        unnestJoinType = INNER;
    }
    // make sure that the rewritten node is with ordinality, which might be necessary to restore inner unnest semantics after rewrite.
    Symbol ordinalitySymbol = unnestNode.getOrdinalitySymbol().orElseGet(() -> context.getSymbolAllocator().newSymbol("ordinality", BIGINT));
    // rewrite correlated join to UnnestNode.
    UnnestNode rewrittenUnnest = new UnnestNode(context.getIdAllocator().getNextId(), input, input.getOutputSymbols(), unnestNode.getMappings(), Optional.of(ordinalitySymbol), unnestJoinType, Optional.empty());
    // restore all nodes from the subquery
    PlanNode rewrittenPlan = Rewriter.rewriteNodeSequence(correlatedJoinNode.getSubquery(), input.getOutputSymbols(), ordinalitySymbol, uniqueSymbol, rewrittenUnnest, context.getSession(), metadata, context.getLookup(), context.getIdAllocator(), context.getSymbolAllocator());
    // between unnested rows and synthetic rows added by left unnest.
    if (unnestNode.getJoinType() == INNER && rewrittenUnnest.getJoinType() == LEFT) {
        Assignments.Builder assignments = Assignments.builder().putIdentities(correlatedJoinNode.getInput().getOutputSymbols());
        for (Symbol subquerySymbol : correlatedJoinNode.getSubquery().getOutputSymbols()) {
            assignments.put(subquerySymbol, new IfExpression(new IsNullPredicate(ordinalitySymbol.toSymbolReference()), new Cast(new NullLiteral(), toSqlType(context.getSymbolAllocator().getTypes().get(subquerySymbol))), subquerySymbol.toSymbolReference()));
        }
        rewrittenPlan = new ProjectNode(context.getIdAllocator().getNextId(), rewrittenPlan, assignments.build());
    }
    // restrict outputs
    return Result.ofPlanNode(restrictOutputs(context.getIdAllocator(), rewrittenPlan, ImmutableSet.copyOf(correlatedJoinNode.getOutputSymbols())).orElse(rewrittenPlan));
}

Example 3 with EnforceSingleRowNode

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

the class TestCostCalculator method testEnforceSingleRow.

@Test
public void testEnforceSingleRow() {
    TableScanNode ts1 = tableScan("ts1", "orderkey");
    EnforceSingleRowNode singleRow = new EnforceSingleRowNode(new PlanNodeId("singleRow"), ts1);
    Map<String, PlanNodeStatsEstimate> stats = ImmutableMap.of("ts1", statsEstimate(ts1, 4000), "singleRow", // 1 * Average Row Size
    statsEstimate(ts1, 8));
    Map<String, PlanCostEstimate> costs = ImmutableMap.of("ts1", cpuCost(1000));
    Map<String, Type> types = ImmutableMap.of("orderkey", BIGINT);
    assertCost(singleRow, costs, stats, types).cpu(// Only count the accumulated cost of source nodes
    1000).memory(0).network(0);
    assertCostEstimatedExchanges(singleRow, costs, stats, types).cpu(1000).memory(0).network(0);
}