Examples with OrderingScheme io.trino.sql.planner.OrderingScheme used on opensource projects

Example 1 with OrderingScheme

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

use of io.trino.sql.planner.OrderingScheme in project trino by trinodb.

the class TestReplaceWindowWithRowNumber method testDoNotFire.

@Test
public void testDoNotFire() {
    ResolvedFunction rank = tester().getMetadata().resolveFunction(tester().getSession(), QualifiedName.of("rank"), fromTypes());
    tester().assertThat(new ReplaceWindowWithRowNumber(tester().getMetadata())).on(p -> {
        Symbol a = p.symbol("a");
        Symbol rank1 = p.symbol("rank_1");
        return p.window(new WindowNode.Specification(ImmutableList.of(a), Optional.empty()), ImmutableMap.of(rank1, newWindowNodeFunction(rank)), p.values(a));
    }).doesNotFire();
    ResolvedFunction rowNumber = tester().getMetadata().resolveFunction(tester().getSession(), QualifiedName.of("row_number"), fromTypes());
    tester().assertThat(new ReplaceWindowWithRowNumber(tester().getMetadata())).on(p -> {
        Symbol a = p.symbol("a");
        Symbol rowNumber1 = p.symbol("row_number_1");
        Symbol rank1 = p.symbol("rank_1");
        return p.window(new WindowNode.Specification(ImmutableList.of(a), Optional.empty()), ImmutableMap.of(rowNumber1, newWindowNodeFunction(rowNumber), rank1, newWindowNodeFunction(rank)), p.values(a));
    }).doesNotFire();
    tester().assertThat(new ReplaceWindowWithRowNumber(tester().getMetadata())).on(p -> {
        Symbol a = p.symbol("a");
        OrderingScheme orderingScheme = new OrderingScheme(ImmutableList.of(a), ImmutableMap.of(a, SortOrder.ASC_NULLS_FIRST));
        Symbol rowNumber1 = p.symbol("row_number_1");
        return p.window(new WindowNode.Specification(ImmutableList.of(a), Optional.of(orderingScheme)), ImmutableMap.of(rowNumber1, newWindowNodeFunction(rowNumber)), p.values(a));
    }).doesNotFire();
}

Example 3 with OrderingScheme

use of io.trino.sql.planner.OrderingScheme in project trino by trinodb.

the class TestPushdownLimitIntoWindow method testConvertToTopNRowNumber.

@Test
public void testConvertToTopNRowNumber() {
    ResolvedFunction ranking = tester().getMetadata().resolveFunction(tester().getSession(), QualifiedName.of("row_number"), fromTypes());
    tester().assertThat(new PushdownLimitIntoWindow()).on(p -> {
        Symbol a = p.symbol("a");
        Symbol rowNumberSymbol = p.symbol("row_number_1");
        OrderingScheme orderingScheme = new OrderingScheme(ImmutableList.of(a), ImmutableMap.of(a, SortOrder.ASC_NULLS_FIRST));
        return p.limit(3, p.window(new WindowNode.Specification(ImmutableList.of(), Optional.of(orderingScheme)), ImmutableMap.of(rowNumberSymbol, newWindowNodeFunction(ranking, a)), p.values(a)));
    }).matches(topNRanking(pattern -> pattern.specification(ImmutableList.of(), ImmutableList.of("a"), ImmutableMap.of("a", SortOrder.ASC_NULLS_FIRST)).maxRankingPerPartition(3).partial(false), values("a")));
}

Example 4 with OrderingScheme

use of io.trino.sql.planner.OrderingScheme in project trino by trinodb.

the class TestPushdownLimitIntoWindow method assertLimitAboveWindow.

private void assertLimitAboveWindow(String rankingFunctionName) {
    ResolvedFunction ranking = tester().getMetadata().resolveFunction(tester().getSession(), QualifiedName.of(rankingFunctionName), fromTypes());
    tester().assertThat(new PushdownLimitIntoWindow()).on(p -> {
        Symbol a = p.symbol("a");
        Symbol rowNumberSymbol = p.symbol("row_number_1");
        OrderingScheme orderingScheme = new OrderingScheme(ImmutableList.of(a), ImmutableMap.of(a, SortOrder.ASC_NULLS_FIRST));
        return p.limit(3, p.window(new WindowNode.Specification(ImmutableList.of(a), Optional.of(orderingScheme)), ImmutableMap.of(rowNumberSymbol, newWindowNodeFunction(ranking, a)), p.values(a)));
    }).matches(limit(3, topNRanking(pattern -> pattern.specification(ImmutableList.of("a"), ImmutableList.of("a"), ImmutableMap.of("a", SortOrder.ASC_NULLS_FIRST)).maxRankingPerPartition(3).partial(false), values("a"))));
}

Example 5 with OrderingScheme

use of io.trino.sql.planner.OrderingScheme in project trino by trinodb.

the class TestPushdownLimitIntoWindow method testLimitWithPreSortedInputs.

@Test
public void testLimitWithPreSortedInputs() {
    // We can push Limit with pre-sorted inputs into WindowNode if ordering scheme is satisfied
    // We don't do it currently to avoid relying on LocalProperties outside of AddExchanges/AddLocalExchanges
    ResolvedFunction ranking = tester().getMetadata().resolveFunction(tester().getSession(), QualifiedName.of("row_number"), fromTypes());
    tester().assertThat(new PushdownLimitIntoWindow()).on(p -> {
        Symbol a = p.symbol("a");
        Symbol rowNumberSymbol = p.symbol("row_number_1");
        OrderingScheme orderingScheme = new OrderingScheme(ImmutableList.of(a), ImmutableMap.of(a, SortOrder.ASC_NULLS_FIRST));
        return p.limit(3, false, ImmutableList.of(a), p.window(new WindowNode.Specification(ImmutableList.of(), Optional.of(orderingScheme)), ImmutableMap.of(rowNumberSymbol, newWindowNodeFunction(ranking, a)), p.values(a)));
    }).doesNotFire();
}