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

Example 1 with RowNumberNode

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

the class ImplementOffset method apply.

@Override
public Result apply(OffsetNode parent, Captures captures, Context context) {
    Symbol rowNumberSymbol = context.getSymbolAllocator().newSymbol("row_number", BIGINT);
    RowNumberNode rowNumberNode = new RowNumberNode(context.getIdAllocator().getNextId(), parent.getSource(), ImmutableList.of(), true, rowNumberSymbol, Optional.empty(), Optional.empty());
    FilterNode filterNode = new FilterNode(context.getIdAllocator().getNextId(), rowNumberNode, new ComparisonExpression(ComparisonExpression.Operator.GREATER_THAN, rowNumberSymbol.toSymbolReference(), new GenericLiteral("BIGINT", Long.toString(parent.getCount()))));
    ProjectNode projectNode = new ProjectNode(context.getIdAllocator().getNextId(), filterNode, Assignments.identity(parent.getOutputSymbols()));
    return Result.ofPlanNode(projectNode);
}

Example 2 with RowNumberNode

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

the class PushdownLimitIntoRowNumber method apply.

@Override
public Result apply(LimitNode node, Captures captures, Context context) {
    RowNumberNode source = captures.get(CHILD);
    RowNumberNode rowNumberNode = mergeLimit(source, node);
    if (rowNumberNode.getPartitionBy().isEmpty()) {
        return Result.ofPlanNode(rowNumberNode);
    }
    if (source.getMaxRowCountPerPartition().isPresent()) {
        if (rowNumberNode.getMaxRowCountPerPartition().equals(source.getMaxRowCountPerPartition())) {
            // Source node unchanged
            return Result.empty();
        }
    }
    return Result.ofPlanNode(replaceChildren(node, ImmutableList.of(rowNumberNode)));
}

Example 3 with RowNumberNode

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

the class RowNumberMatcher 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());
    RowNumberNode rowNumberNode = (RowNumberNode) node;
    if (partitionBy.isPresent()) {
        List<Symbol> expected = partitionBy.get().stream().map(alias -> alias.toSymbol(symbolAliases)).collect(toImmutableList());
        if (!expected.equals(rowNumberNode.getPartitionBy())) {
            return NO_MATCH;
        }
    }
    if (rowNumberSymbol.isPresent()) {
        Symbol expected = rowNumberSymbol.get().toSymbol(symbolAliases);
        if (!expected.equals(rowNumberNode.getRowNumberSymbol())) {
            return NO_MATCH;
        }
    }
    if (maxRowCountPerPartition.isPresent()) {
        if (!maxRowCountPerPartition.get().equals(rowNumberNode.getMaxRowCountPerPartition())) {
            return NO_MATCH;
        }
    }
    if (hashSymbol.isPresent()) {
        Optional<Symbol> expected = hashSymbol.get().map(alias -> alias.toSymbol(symbolAliases));
        if (!expected.equals(rowNumberNode.getHashSymbol())) {
            return NO_MATCH;
        }
    }
    if (orderSensitive.isPresent()) {
        if (!orderSensitive.get().equals(rowNumberNode.isOrderSensitive())) {
            return NO_MATCH;
        }
    }
    return match();
}

Example 4 with RowNumberNode

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

the class PushPredicateThroughProjectIntoRowNumber method apply.

@Override
public Result apply(FilterNode filter, Captures captures, Context context) {
    ProjectNode project = captures.get(PROJECT);
    RowNumberNode rowNumber = captures.get(ROW_NUMBER);
    Symbol rowNumberSymbol = rowNumber.getRowNumberSymbol();
    if (!project.getAssignments().getSymbols().contains(rowNumberSymbol)) {
        return Result.empty();
    }
    ExtractionResult extractionResult = DomainTranslator.getExtractionResult(plannerContext, context.getSession(), filter.getPredicate(), context.getSymbolAllocator().getTypes());
    TupleDomain<Symbol> tupleDomain = extractionResult.getTupleDomain();
    OptionalInt upperBound = extractUpperBound(tupleDomain, rowNumberSymbol);
    if (upperBound.isEmpty()) {
        return Result.empty();
    }
    if (upperBound.getAsInt() <= 0) {
        return Result.ofPlanNode(new ValuesNode(filter.getId(), filter.getOutputSymbols(), ImmutableList.of()));
    }
    boolean updatedMaxRowCountPerPartition = false;
    if (rowNumber.getMaxRowCountPerPartition().isEmpty() || rowNumber.getMaxRowCountPerPartition().get() > upperBound.getAsInt()) {
        rowNumber = new RowNumberNode(rowNumber.getId(), rowNumber.getSource(), rowNumber.getPartitionBy(), rowNumber.isOrderSensitive(), rowNumber.getRowNumberSymbol(), Optional.of(upperBound.getAsInt()), rowNumber.getHashSymbol());
        project = (ProjectNode) project.replaceChildren(ImmutableList.of(rowNumber));
        updatedMaxRowCountPerPartition = true;
    }
    if (!allRowNumberValuesInDomain(tupleDomain, rowNumberSymbol, rowNumber.getMaxRowCountPerPartition().get())) {
        if (updatedMaxRowCountPerPartition) {
            return Result.ofPlanNode(filter.replaceChildren(ImmutableList.of(project)));
        }
        return Result.empty();
    }
    // Remove the row number domain because it is absorbed into the node
    TupleDomain<Symbol> newTupleDomain = tupleDomain.filter((symbol, domain) -> !symbol.equals(rowNumberSymbol));
    Expression newPredicate = ExpressionUtils.combineConjuncts(plannerContext.getMetadata(), extractionResult.getRemainingExpression(), new DomainTranslator(plannerContext).toPredicate(context.getSession(), newTupleDomain));
    if (newPredicate.equals(TRUE_LITERAL)) {
        return Result.ofPlanNode(project);
    }
    return Result.ofPlanNode(new FilterNode(filter.getId(), project, newPredicate));
}

Example 5 with RowNumberNode

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

the class PushDownDereferencesThroughRowNumber method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    RowNumberNode rowNumberNode = captures.get(CHILD);
    // Extract dereferences from project node assignments for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(projectNode.getAssignments().getExpressions(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    // Exclude dereferences on symbols being used in partitionBy
    dereferences = dereferences.stream().filter(expression -> !rowNumberNode.getPartitionBy().contains(getBase(expression))).collect(toImmutableSet());
    if (dereferences.isEmpty()) {
        return Result.empty();
    }
    // Create new symbols for dereference expressions
    Assignments dereferenceAssignments = Assignments.of(dereferences, context.getSession(), context.getSymbolAllocator(), typeAnalyzer);
    // Rewrite project node assignments using new symbols for dereference expressions
    Map<Expression, SymbolReference> mappings = HashBiMap.create(dereferenceAssignments.getMap()).inverse().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
    Assignments newAssignments = projectNode.getAssignments().rewrite(expression -> replaceExpression(expression, mappings));
    return Result.ofPlanNode(new ProjectNode(context.getIdAllocator().getNextId(), rowNumberNode.replaceChildren(ImmutableList.of(new ProjectNode(context.getIdAllocator().getNextId(), rowNumberNode.getSource(), Assignments.builder().putIdentities(rowNumberNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build()))), newAssignments));
}