Search in sources :

Example 11 with WindowNode

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

the class ImplementLimitWithTies method rewriteLimitWithTiesWithPartitioning.

/**
 * Rewrite LimitNode with ties to WindowNode and FilterNode, with partitioning defined by partitionBy.
 * <p>
 * This method does not prune outputs of the rewritten plan. After the rewrite, the output consists of
 * source's output symbols and the newly created rankSymbol.
 * Passing all input symbols is intentional, because this method is used for de-correlation in the scenario
 * where the original LimitNode is in the correlated subquery, and the rewrite result is placed on top of
 * de-correlated join.
 * It is the responsibility of the caller to prune redundant outputs.
 */
public static PlanNode rewriteLimitWithTiesWithPartitioning(LimitNode limitNode, PlanNode source, Session session, Metadata metadata, PlanNodeIdAllocator idAllocator, SymbolAllocator symbolAllocator, List<Symbol> partitionBy) {
    checkArgument(limitNode.isWithTies(), "Expected LimitNode with ties");
    Symbol rankSymbol = symbolAllocator.newSymbol("rank_num", BIGINT);
    WindowNode.Function rankFunction = new WindowNode.Function(metadata.resolveFunction(session, QualifiedName.of("rank"), ImmutableList.of()), ImmutableList.of(), DEFAULT_FRAME, false);
    WindowNode windowNode = new WindowNode(idAllocator.getNextId(), source, new WindowNode.Specification(partitionBy, limitNode.getTiesResolvingScheme()), ImmutableMap.of(rankSymbol, rankFunction), Optional.empty(), ImmutableSet.of(), 0);
    return new FilterNode(idAllocator.getNextId(), windowNode, new ComparisonExpression(ComparisonExpression.Operator.LESS_THAN_OR_EQUAL, rankSymbol.toSymbolReference(), new GenericLiteral("BIGINT", Long.toString(limitNode.getCount()))));
}
Also used : WindowNode(io.trino.sql.planner.plan.WindowNode) ComparisonExpression(io.trino.sql.tree.ComparisonExpression) Symbol(io.trino.sql.planner.Symbol) FilterNode(io.trino.sql.planner.plan.FilterNode) GenericLiteral(io.trino.sql.tree.GenericLiteral)

Example 12 with WindowNode

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

the class WindowMatcher 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());
    WindowNode windowNode = (WindowNode) node;
    if (!prePartitionedInputs.map(expectedInputs -> expectedInputs.stream().map(alias -> alias.toSymbol(symbolAliases)).collect(toImmutableSet()).equals(windowNode.getPrePartitionedInputs())).orElse(true)) {
        return NO_MATCH;
    }
    if (!specification.map(expectedSpecification -> expectedSpecification.getExpectedValue(symbolAliases).equals(windowNode.getSpecification())).orElse(true)) {
        return NO_MATCH;
    }
    if (!preSortedOrderPrefix.map(Integer.valueOf(windowNode.getPreSortedOrderPrefix())::equals).orElse(true)) {
        return NO_MATCH;
    }
    if (!hashSymbol.map(expectedHashSymbol -> expectedHashSymbol.map(alias -> alias.toSymbol(symbolAliases)).equals(windowNode.getHashSymbol())).orElse(true)) {
        return NO_MATCH;
    }
    /*
         * Window functions produce a symbol (the result of the function call) that we might
         * want to bind to an alias so we can reference it further up the tree. As such,
         * they need to be matched with an Alias matcher so we can bind the symbol if desired.
         */
    return match();
}
Also used : MatchResult.match(io.trino.sql.planner.assertions.MatchResult.match) ResolvedFunction(io.trino.metadata.ResolvedFunction) Set(java.util.Set) StatsProvider(io.trino.cost.StatsProvider) SortOrder(io.trino.spi.connector.SortOrder) PlanNode(io.trino.sql.planner.plan.PlanNode) Preconditions.checkState(com.google.common.base.Preconditions.checkState) PlanMatchPattern.node(io.trino.sql.planner.assertions.PlanMatchPattern.node) List(java.util.List) Map(java.util.Map) Objects.requireNonNull(java.util.Objects.requireNonNull) Metadata(io.trino.metadata.Metadata) Optional(java.util.Optional) ImmutableSet.toImmutableSet(com.google.common.collect.ImmutableSet.toImmutableSet) NO_MATCH(io.trino.sql.planner.assertions.MatchResult.NO_MATCH) WindowNode(io.trino.sql.planner.plan.WindowNode) LinkedList(java.util.LinkedList) FunctionCall(io.trino.sql.tree.FunctionCall) Session(io.trino.Session) MoreObjects.toStringHelper(com.google.common.base.MoreObjects.toStringHelper) WindowNode(io.trino.sql.planner.plan.WindowNode)

Example 13 with WindowNode

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

the class PushPredicateThroughProjectIntoWindow method apply.

@Override
public Result apply(FilterNode filter, Captures captures, Context context) {
    ProjectNode project = captures.get(PROJECT);
    WindowNode window = captures.get(WINDOW);
    Symbol rankingSymbol = getOnlyElement(window.getWindowFunctions().keySet());
    if (!project.getAssignments().getSymbols().contains(rankingSymbol)) {
        return Result.empty();
    }
    DomainTranslator.ExtractionResult extractionResult = DomainTranslator.getExtractionResult(plannerContext, context.getSession(), filter.getPredicate(), context.getSymbolAllocator().getTypes());
    TupleDomain<Symbol> tupleDomain = extractionResult.getTupleDomain();
    OptionalInt upperBound = extractUpperBound(tupleDomain, rankingSymbol);
    if (upperBound.isEmpty()) {
        return Result.empty();
    }
    if (upperBound.getAsInt() <= 0) {
        return Result.ofPlanNode(new ValuesNode(filter.getId(), filter.getOutputSymbols(), ImmutableList.of()));
    }
    RankingType rankingType = toTopNRankingType(window).orElseThrow();
    project = (ProjectNode) project.replaceChildren(ImmutableList.of(new TopNRankingNode(window.getId(), window.getSource(), window.getSpecification(), rankingType, rankingSymbol, upperBound.getAsInt(), false, Optional.empty())));
    if (!allRankingValuesInDomain(tupleDomain, rankingSymbol, upperBound.getAsInt())) {
        return Result.ofPlanNode(filter.replaceChildren(ImmutableList.of(project)));
    }
    // Remove the ranking domain because it is absorbed into the node
    TupleDomain<Symbol> newTupleDomain = tupleDomain.filter((symbol, domain) -> !symbol.equals(rankingSymbol));
    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));
}
Also used : WindowNode(io.trino.sql.planner.plan.WindowNode) ValuesNode(io.trino.sql.planner.plan.ValuesNode) Util.toTopNRankingType(io.trino.sql.planner.iterative.rule.Util.toTopNRankingType) RankingType(io.trino.sql.planner.plan.TopNRankingNode.RankingType) TopNRankingNode(io.trino.sql.planner.plan.TopNRankingNode) Expression(io.trino.sql.tree.Expression) Symbol(io.trino.sql.planner.Symbol) FilterNode(io.trino.sql.planner.plan.FilterNode) DomainTranslator(io.trino.sql.planner.DomainTranslator) ProjectNode(io.trino.sql.planner.plan.ProjectNode) OptionalInt(java.util.OptionalInt)

Example 14 with WindowNode

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

the class SetOperationNodeTranslator method makeSetContainmentPlanForAll.

public TranslationResult makeSetContainmentPlanForAll(SetOperationNode node) {
    checkArgument(!(node instanceof UnionNode), "Cannot simplify a UnionNode");
    List<Symbol> markers = allocateSymbols(node.getSources().size(), MARKER, BOOLEAN);
    // identity projection for all the fields in each of the sources plus marker columns
    List<PlanNode> withMarkers = appendMarkers(markers, node.getSources(), node);
    // add a union over all the rewritten sources
    List<Symbol> outputs = node.getOutputSymbols();
    UnionNode union = union(withMarkers, ImmutableList.copyOf(concat(outputs, markers)));
    // add counts and row number
    List<Symbol> countOutputs = allocateSymbols(markers.size(), "count", BIGINT);
    Symbol rowNumberSymbol = symbolAllocator.newSymbol("row_number", BIGINT);
    WindowNode window = appendCounts(union, outputs, markers, countOutputs, rowNumberSymbol);
    // prune markers
    ProjectNode project = new ProjectNode(idAllocator.getNextId(), window, Assignments.identity(ImmutableList.copyOf(concat(outputs, countOutputs, ImmutableList.of(rowNumberSymbol)))));
    return new TranslationResult(project, countOutputs, Optional.of(rowNumberSymbol));
}
Also used : WindowNode(io.trino.sql.planner.plan.WindowNode) PlanNode(io.trino.sql.planner.plan.PlanNode) UnionNode(io.trino.sql.planner.plan.UnionNode) Symbol(io.trino.sql.planner.Symbol) ProjectNode(io.trino.sql.planner.plan.ProjectNode)

Example 15 with WindowNode

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

the class PushdownLimitIntoWindow method apply.

@Override
public Result apply(LimitNode node, Captures captures, Context context) {
    WindowNode source = captures.get(childCapture);
    Optional<RankingType> rankingType = toTopNRankingType(source);
    int limit = toIntExact(node.getCount());
    TopNRankingNode topNRowNumberNode = new TopNRankingNode(source.getId(), source.getSource(), source.getSpecification(), rankingType.get(), getOnlyElement(source.getWindowFunctions().keySet()), limit, false, Optional.empty());
    if (rankingType.get() == ROW_NUMBER && source.getPartitionBy().isEmpty()) {
        return Result.ofPlanNode(topNRowNumberNode);
    }
    return Result.ofPlanNode(replaceChildren(node, ImmutableList.of(topNRowNumberNode)));
}
Also used : WindowNode(io.trino.sql.planner.plan.WindowNode) Util.toTopNRankingType(io.trino.sql.planner.iterative.rule.Util.toTopNRankingType) RankingType(io.trino.sql.planner.plan.TopNRankingNode.RankingType) TopNRankingNode(io.trino.sql.planner.plan.TopNRankingNode)

Aggregations

WindowNode (io.trino.sql.planner.plan.WindowNode)17 ResolvedFunction (io.trino.metadata.ResolvedFunction)9 Symbol (io.trino.sql.planner.Symbol)9 PlanNode (io.trino.sql.planner.plan.PlanNode)8 Expression (io.trino.sql.tree.Expression)6 WindowFrame (io.trino.sql.tree.WindowFrame)5 FilterNode (io.trino.sql.planner.plan.FilterNode)4 ComparisonExpression (io.trino.sql.tree.ComparisonExpression)4 Map (java.util.Map)4 Test (org.testng.annotations.Test)4 ImmutableList (com.google.common.collect.ImmutableList)3 ImmutableMap (com.google.common.collect.ImmutableMap)3 ImmutableSet.toImmutableSet (com.google.common.collect.ImmutableSet.toImmutableSet)3 Session (io.trino.Session)3 SortOrder (io.trino.spi.connector.SortOrder)3 Util.toTopNRankingType (io.trino.sql.planner.iterative.rule.Util.toTopNRankingType)3 ProjectNode (io.trino.sql.planner.plan.ProjectNode)3 Preconditions.checkState (com.google.common.base.Preconditions.checkState)2 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)2 ImmutableMap.toImmutableMap (com.google.common.collect.ImmutableMap.toImmutableMap)2