Examples with PlanBuilder io.trino.sql.planner.iterative.rule.test.PlanBuilder used on opensource projects

Example 11 with PlanBuilder

use of io.trino.sql.planner.iterative.rule.test.PlanBuilder in project trino by trinodb.

the class TestValidateLimitWithPresortedInput method validatePlan.

private void validatePlan(Function<PlanBuilder, PlanNode> planProvider) {
    LocalQueryRunner queryRunner = getQueryRunner();
    PlanBuilder builder = new PlanBuilder(idAllocator, queryRunner.getMetadata(), queryRunner.getDefaultSession());
    PlanNode planNode = planProvider.apply(builder);
    TypeProvider types = builder.getTypes();
    queryRunner.inTransaction(session -> {
        // metadata.getCatalogHandle() registers the catalog for the transaction
        session.getCatalog().ifPresent(catalog -> queryRunner.getMetadata().getCatalogHandle(session, catalog));
        TypeAnalyzer typeAnalyzer = createTestingTypeAnalyzer(queryRunner.getPlannerContext());
        new ValidateLimitWithPresortedInput().validate(planNode, session, queryRunner.getPlannerContext(), typeAnalyzer, types, WarningCollector.NOOP);
        return null;
    });
}

Example 12 with PlanBuilder

use of io.trino.sql.planner.iterative.rule.test.PlanBuilder in project trino by trinodb.

the class TestUnaliasSymbolReferences method assertOptimizedPlan.

private void assertOptimizedPlan(PlanOptimizer optimizer, PlanCreator planCreator, PlanMatchPattern pattern) {
    LocalQueryRunner queryRunner = getQueryRunner();
    queryRunner.inTransaction(session -> {
        Metadata metadata = queryRunner.getMetadata();
        session.getCatalog().ifPresent(catalog -> metadata.getCatalogHandle(session, catalog));
        PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
        PlanBuilder planBuilder = new PlanBuilder(idAllocator, metadata, session);
        SymbolAllocator symbolAllocator = new SymbolAllocator();
        PlanNode plan = planCreator.create(planBuilder, session, metadata);
        PlanNode optimized = optimizer.optimize(plan, session, planBuilder.getTypes(), symbolAllocator, idAllocator, WarningCollector.NOOP);
        Plan actual = new Plan(optimized, planBuilder.getTypes(), StatsAndCosts.empty());
        PlanAssert.assertPlan(session, queryRunner.getMetadata(), queryRunner.getFunctionManager(), queryRunner.getStatsCalculator(), actual, pattern);
        return null;
    });
}

Example 13 with PlanBuilder

use of io.trino.sql.planner.iterative.rule.test.PlanBuilder in project trino by trinodb.

the class TestRemoveUnsupportedDynamicFilters method setup.

@BeforeClass
public void setup() {
    plannerContext = getQueryRunner().getPlannerContext();
    metadata = plannerContext.getMetadata();
    builder = new PlanBuilder(new PlanNodeIdAllocator(), metadata, TEST_SESSION);
    CatalogName catalogName = getCurrentConnectorId();
    lineitemTableHandle = new TableHandle(catalogName, new TpchTableHandle("sf1", "lineitem", 1.0), TestingTransactionHandle.create());
    lineitemOrderKeySymbol = builder.symbol("LINEITEM_OK", BIGINT);
    lineitemTableScanNode = builder.tableScan(lineitemTableHandle, ImmutableList.of(lineitemOrderKeySymbol), ImmutableMap.of(lineitemOrderKeySymbol, new TpchColumnHandle("orderkey", BIGINT)));
    TableHandle ordersTableHandle = new TableHandle(catalogName, new TpchTableHandle("sf1", "orders", 1.0), TestingTransactionHandle.create());
    ordersOrderKeySymbol = builder.symbol("ORDERS_OK", BIGINT);
    ordersTableScanNode = builder.tableScan(ordersTableHandle, ImmutableList.of(ordersOrderKeySymbol), ImmutableMap.of(ordersOrderKeySymbol, new TpchColumnHandle("orderkey", BIGINT)));
}

Example 14 with PlanBuilder

use of io.trino.sql.planner.iterative.rule.test.PlanBuilder in project trino by trinodb.

the class TestBeginTableWrite method applyOptimization.

private void applyOptimization(Function<PlanBuilder, PlanNode> planProvider) {
    Metadata metadata = new MockMetadata();
    new BeginTableWrite(metadata, createTestingFunctionManager()).optimize(planProvider.apply(new PlanBuilder(new PlanNodeIdAllocator(), metadata, testSessionBuilder().build())), testSessionBuilder().build(), empty(), new SymbolAllocator(), new PlanNodeIdAllocator(), WarningCollector.NOOP);
}

Example 15 with PlanBuilder

use of io.trino.sql.planner.iterative.rule.test.PlanBuilder in project trino by trinodb.

the class TestCardinalityExtractorPlanVisitor method testAggregation.

@Test
public void testAggregation() {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), dummyMetadata(), TEST_SESSION);
    Symbol symbol = planBuilder.symbol("symbol");
    ColumnHandle columnHandle = new TestingColumnHandle("column");
    // single default aggregation
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet().source(planBuilder.values(10)))), Range.singleton(1L));
    // multiple grouping sets with default aggregation with source that produces arbitrary number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.tableScan(ImmutableList.of(symbol), ImmutableMap.of(symbol, columnHandle))))), Range.atLeast(1L));
    // multiple grouping sets with default aggregation with source that produces exact number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.values(10, symbol)))), Range.closed(1L, 10L));
    // multiple grouping sets with default aggregation with source that produces no rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.groupingSets(AggregationNode.groupingSets(ImmutableList.of(symbol), 2, ImmutableSet.of(0))).source(planBuilder.values(0, symbol)))), Range.singleton(1L));
    // single non-default aggregation with source that produces arbitrary number of rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.tableScan(ImmutableList.of(symbol), ImmutableMap.of(symbol, columnHandle))))), Range.atLeast(0L));
    // single non-default aggregation with source that produces at least single row
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.values(10, symbol)))), Range.closed(1L, 10L));
    // single non-default aggregation with source that produces no rows
    assertEquals(extractCardinality(planBuilder.aggregation(builder -> builder.singleGroupingSet(symbol).source(planBuilder.values(0, symbol)))), Range.singleton(0L));
}