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

Example 1 with WindowNode

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

the class PushDownDereferencesThroughWindow method apply.

@Override
public Result apply(ProjectNode projectNode, Captures captures, Context context) {
    WindowNode windowNode = captures.get(CHILD);
    // Extract dereferences for pushdown
    Set<SubscriptExpression> dereferences = extractRowSubscripts(ImmutableList.<Expression>builder().addAll(projectNode.getAssignments().getExpressions()).addAll(windowNode.getWindowFunctions().values().stream().flatMap(function -> function.getArguments().stream()).collect(toImmutableList())).build(), false, context.getSession(), typeAnalyzer, context.getSymbolAllocator().getTypes());
    WindowNode.Specification specification = windowNode.getSpecification();
    dereferences = dereferences.stream().filter(expression -> {
        Symbol symbol = getBase(expression);
        // Exclude partitionBy, orderBy and synthesized symbols
        return !specification.getPartitionBy().contains(symbol) && !specification.getOrderingScheme().map(OrderingScheme::getOrderBy).orElse(ImmutableList.of()).contains(symbol) && !windowNode.getCreatedSymbols().contains(symbol);
    }).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(), new WindowNode(windowNode.getId(), new ProjectNode(context.getIdAllocator().getNextId(), windowNode.getSource(), Assignments.builder().putIdentities(windowNode.getSource().getOutputSymbols()).putAll(dereferenceAssignments).build()), windowNode.getSpecification(), // Replace dereference expressions in functions
    windowNode.getWindowFunctions().entrySet().stream().collect(toImmutableMap(Map.Entry::getKey, entry -> {
        WindowNode.Function oldFunction = entry.getValue();
        return new WindowNode.Function(oldFunction.getResolvedFunction(), oldFunction.getArguments().stream().map(expression -> replaceExpression(expression, mappings)).collect(toImmutableList()), oldFunction.getFrame(), oldFunction.isIgnoreNulls());
    })), windowNode.getHashSymbol(), windowNode.getPrePartitionedInputs(), windowNode.getPreSortedOrderPrefix()), newAssignments));
}

Example 2 with WindowNode

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

the class PruneWindowColumns method pushDownProjectOff.

@Override
protected Optional<PlanNode> pushDownProjectOff(Context context, WindowNode windowNode, Set<Symbol> referencedOutputs) {
    Map<Symbol, WindowNode.Function> referencedFunctions = Maps.filterKeys(windowNode.getWindowFunctions(), referencedOutputs::contains);
    if (referencedFunctions.isEmpty()) {
        return Optional.of(windowNode.getSource());
    }
    ImmutableSet.Builder<Symbol> referencedInputs = ImmutableSet.<Symbol>builder().addAll(windowNode.getSource().getOutputSymbols().stream().filter(referencedOutputs::contains).iterator()).addAll(windowNode.getPartitionBy());
    windowNode.getOrderingScheme().ifPresent(orderingScheme -> orderingScheme.getOrderBy().forEach(referencedInputs::add));
    windowNode.getHashSymbol().ifPresent(referencedInputs::add);
    for (WindowNode.Function windowFunction : referencedFunctions.values()) {
        referencedInputs.addAll(SymbolsExtractor.extractUnique(windowFunction));
    }
    PlanNode prunedWindowNode = new WindowNode(windowNode.getId(), restrictOutputs(context.getIdAllocator(), windowNode.getSource(), referencedInputs.build()).orElse(windowNode.getSource()), windowNode.getSpecification(), referencedFunctions, windowNode.getHashSymbol(), windowNode.getPrePartitionedInputs(), windowNode.getPreSortedOrderPrefix());
    if (prunedWindowNode.getOutputSymbols().size() == windowNode.getOutputSymbols().size()) {
        // Neither function pruning nor input pruning was successful.
        return Optional.empty();
    }
    return Optional.of(prunedWindowNode);
}

Example 3 with WindowNode

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

the class SetOperationNodeTranslator method appendCounts.

private WindowNode appendCounts(UnionNode sourceNode, List<Symbol> originalColumns, List<Symbol> markers, List<Symbol> countOutputs, Symbol rowNumberSymbol) {
    ImmutableMap.Builder<Symbol, WindowNode.Function> functions = ImmutableMap.builder();
    WindowNode.Frame defaultFrame = new WindowNode.Frame(ROWS, UNBOUNDED_PRECEDING, Optional.empty(), Optional.empty(), UNBOUNDED_FOLLOWING, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
    for (int i = 0; i < markers.size(); i++) {
        Symbol output = countOutputs.get(i);
        functions.put(output, new WindowNode.Function(countFunction, ImmutableList.of(markers.get(i).toSymbolReference()), defaultFrame, false));
    }
    functions.put(rowNumberSymbol, new WindowNode.Function(rowNumberFunction, ImmutableList.of(), defaultFrame, false));
    return new WindowNode(idAllocator.getNextId(), sourceNode, new Specification(originalColumns, Optional.empty()), functions.buildOrThrow(), Optional.empty(), ImmutableSet.of(), 0);
}

Example 4 with WindowNode

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

the class PushdownFilterIntoWindow method apply.

@Override
public Result apply(FilterNode node, Captures captures, Context context) {
    Session session = context.getSession();
    TypeProvider types = context.getSymbolAllocator().getTypes();
    WindowNode windowNode = captures.get(childCapture);
    DomainTranslator.ExtractionResult extractionResult = DomainTranslator.getExtractionResult(plannerContext, session, node.getPredicate(), types);
    TupleDomain<Symbol> tupleDomain = extractionResult.getTupleDomain();
    Optional<RankingType> rankingType = toTopNRankingType(windowNode);
    Symbol rankingSymbol = getOnlyElement(windowNode.getWindowFunctions().keySet());
    OptionalInt upperBound = extractUpperBound(tupleDomain, rankingSymbol);
    if (upperBound.isEmpty()) {
        return Result.empty();
    }
    if (upperBound.getAsInt() <= 0) {
        return Result.ofPlanNode(new ValuesNode(node.getId(), node.getOutputSymbols(), ImmutableList.of()));
    }
    TopNRankingNode newSource = new TopNRankingNode(windowNode.getId(), windowNode.getSource(), windowNode.getSpecification(), rankingType.get(), rankingSymbol, upperBound.getAsInt(), false, Optional.empty());
    if (!allRowNumberValuesInDomain(tupleDomain, rankingSymbol, upperBound.getAsInt())) {
        return Result.ofPlanNode(new FilterNode(node.getId(), newSource, node.getPredicate()));
    }
    // Remove the row number 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(session, newTupleDomain));
    if (newPredicate.equals(BooleanLiteral.TRUE_LITERAL)) {
        return Result.ofPlanNode(newSource);
    }
    return Result.ofPlanNode(new FilterNode(node.getId(), newSource, newPredicate));
}

Example 5 with WindowNode

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

the class TestTypeValidator method testValidWindow.

@Test
public void testValidWindow() {
    Symbol windowSymbol = symbolAllocator.newSymbol("sum", DOUBLE);
    ResolvedFunction resolvedFunction = functionResolution.resolveFunction(QualifiedName.of("sum"), fromTypes(DOUBLE));
    WindowNode.Frame frame = new WindowNode.Frame(WindowFrame.Type.RANGE, FrameBound.Type.UNBOUNDED_PRECEDING, Optional.empty(), Optional.empty(), FrameBound.Type.UNBOUNDED_FOLLOWING, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty());
    WindowNode.Function function = new WindowNode.Function(resolvedFunction, ImmutableList.of(columnC.toSymbolReference()), frame, false);
    WindowNode.Specification specification = new WindowNode.Specification(ImmutableList.of(), Optional.empty());
    PlanNode node = new WindowNode(newId(), baseTableScan, specification, ImmutableMap.of(windowSymbol, function), Optional.empty(), ImmutableSet.of(), 0);
    assertTypesValid(node);
}