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

Example 1 with TopNNode

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

the class MergeLimitOverProjectWithSort method apply.

@Override
public Result apply(LimitNode parent, Captures captures, Context context) {
    ProjectNode project = captures.get(PROJECT);
    SortNode sort = captures.get(SORT);
    return Result.ofPlanNode(project.replaceChildren(ImmutableList.of(new TopNNode(parent.getId(), sort.getSource(), parent.getCount(), sort.getOrderingScheme(), parent.isPartial() ? PARTIAL : SINGLE))));
}

Example 2 with TopNNode

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

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

the class TestUnion method testUnionUnderTopN.

@Test
public void testUnionUnderTopN() {
    Plan plan = plan("SELECT * FROM (" + "   SELECT regionkey FROM nation " + "   UNION ALL " + "   SELECT nationkey FROM nation" + ") t(a) " + "ORDER BY a LIMIT 1", OPTIMIZED_AND_VALIDATED, false);
    List<PlanNode> remotes = searchFrom(plan.getRoot()).where(TestUnion::isRemoteExchange).findAll();
    assertEquals(remotes.size(), 1, "There should be exactly one RemoteExchange");
    assertEquals(((ExchangeNode) Iterables.getOnlyElement(remotes)).getType(), GATHER);
    int numberOfpartialTopN = searchFrom(plan.getRoot()).where(planNode -> planNode instanceof TopNNode && ((TopNNode) planNode).getStep() == TopNNode.Step.PARTIAL).count();
    assertEquals(numberOfpartialTopN, 2, "There should be exactly two partial TopN nodes");
    assertPlanIsFullyDistributed(plan);
}

Example 4 with TopNNode

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

the class TopNStatsRule method doCalculate.

@Override
protected Optional<PlanNodeStatsEstimate> doCalculate(TopNNode node, StatsProvider statsProvider, Lookup lookup, Session session, TypeProvider types) {
    PlanNodeStatsEstimate sourceStats = statsProvider.getStats(node.getSource());
    double rowCount = sourceStats.getOutputRowCount();
    if (node.getStep() != TopNNode.Step.SINGLE) {
        return Optional.empty();
    }
    if (rowCount <= node.getCount()) {
        return Optional.of(sourceStats);
    }
    long limitCount = node.getCount();
    PlanNodeStatsEstimate resultStats = PlanNodeStatsEstimate.buildFrom(sourceStats).setOutputRowCount(limitCount).build();
    if (limitCount == 0) {
        return Optional.of(resultStats);
    }
    // augment null fraction estimation for first ORDER BY symbol
    // Assuming not empty list
    Symbol firstOrderSymbol = node.getOrderingScheme().getOrderBy().get(0);
    SortOrder sortOrder = node.getOrderingScheme().getOrdering(firstOrderSymbol);
    resultStats = resultStats.mapSymbolColumnStatistics(firstOrderSymbol, symbolStats -> {
        SymbolStatsEstimate.Builder newStats = SymbolStatsEstimate.buildFrom(symbolStats);
        double nullCount = rowCount * symbolStats.getNullsFraction();
        if (sortOrder.isNullsFirst()) {
            if (nullCount > limitCount) {
                newStats.setNullsFraction(1.0);
            } else {
                newStats.setNullsFraction(nullCount / limitCount);
            }
        } else {
            double nonNullCount = (rowCount - nullCount);
            if (nonNullCount > limitCount) {
                newStats.setNullsFraction(0.0);
            } else {
                newStats.setNullsFraction((limitCount - nonNullCount) / limitCount);
            }
        }
        return newStats.build();
    });
    // TopN actually limits (or when there was no row count estimated for source)
    return Optional.of(resultStats);
}

Example 5 with TopNNode

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

the class TopNMatcher method detailMatches.

@Override
public MatchResult detailMatches(PlanNode node, StatsProvider stats, Session session, Metadata metadata, SymbolAliases symbolAliases) {
    checkState(shapeMatches(node), "Plan testing framework error: shapeMatches returned false in detailMatches in %s", this.getClass().getName());
    TopNNode topNNode = (TopNNode) node;
    if (topNNode.getCount() != count) {
        return NO_MATCH;
    }
    if (!orderingSchemeMatches(orderBy, topNNode.getOrderingScheme(), symbolAliases)) {
        return NO_MATCH;
    }
    if (topNNode.getStep() != step) {
        return NO_MATCH;
    }
    return match();
}