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

Example 1 with PlanBuilder

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

the class TestLocalDynamicFilterConsumer method testDynamicFilterPruning.

@Test
public void testDynamicFilterPruning() throws Exception {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), dummyMetadata(), getQueryRunner().getDefaultSession());
    Symbol left1 = planBuilder.symbol("left1", BIGINT);
    Symbol left2 = planBuilder.symbol("left2", INTEGER);
    Symbol left3 = planBuilder.symbol("left3", SMALLINT);
    Symbol right1 = planBuilder.symbol("right1", BIGINT);
    Symbol right2 = planBuilder.symbol("right2", INTEGER);
    Symbol right3 = planBuilder.symbol("right3", SMALLINT);
    DynamicFilterId filter1 = new DynamicFilterId("filter1");
    DynamicFilterId filter2 = new DynamicFilterId("filter2");
    DynamicFilterId filter3 = new DynamicFilterId("filter3");
    JoinNode joinNode = planBuilder.join(INNER, planBuilder.values(left1, left2, left3), planBuilder.values(right1, right2, right3), ImmutableList.of(new EquiJoinClause(left1, right1), new EquiJoinClause(left2, right2), new EquiJoinClause(left3, right3)), ImmutableList.of(), ImmutableList.of(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(filter1, right1, filter2, right2, filter3, right3));
    LocalDynamicFilterConsumer consumer = LocalDynamicFilterConsumer.create(joinNode, ImmutableList.of(BIGINT, INTEGER, SMALLINT), 1, ImmutableSet.of(filter1, filter3));
    assertEquals(consumer.getBuildChannels(), ImmutableMap.of(filter1, 0, filter3, 2));
    // make sure domain types got propagated correctly
    consumer.getTupleDomainConsumer().accept(TupleDomain.none());
    assertEquals(consumer.getDynamicFilterDomains().get(), ImmutableMap.of(filter1, Domain.none(BIGINT), filter3, Domain.none(SMALLINT)));
}

Example 2 with PlanBuilder

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

the class TestSchedulingOrderVisitor method testIndexJoinOrder.

@Test
public void testIndexJoinOrder() {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), dummyMetadata(), TEST_SESSION);
    TableScanNode a = planBuilder.tableScan(emptyList(), emptyMap());
    TableScanNode b = planBuilder.tableScan(emptyList(), emptyMap());
    List<PlanNodeId> order = scheduleOrder(planBuilder.indexJoin(IndexJoinNode.Type.INNER, a, b));
    assertEquals(order, ImmutableList.of(b.getId(), a.getId()));
}

Example 3 with PlanBuilder

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

the class TestRemoveAggregationInSemiJoin method semiJoinWithDistinctAsFilteringSource.

private static PlanNode semiJoinWithDistinctAsFilteringSource(PlanBuilder p) {
    Symbol leftKey = p.symbol("leftKey");
    Symbol rightKey = p.symbol("rightKey");
    return p.semiJoin(leftKey, rightKey, p.symbol("match"), Optional.empty(), Optional.empty(), p.values(leftKey), p.aggregation(builder -> builder.singleGroupingSet(rightKey).source(p.values(rightKey))));
}

Example 4 with PlanBuilder

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

the class TestRemoveAggregationInSemiJoin method semiJoinWithAggregationAsFilteringSource.

private static PlanNode semiJoinWithAggregationAsFilteringSource(PlanBuilder p) {
    Symbol leftKey = p.symbol("leftKey");
    Symbol rightKey = p.symbol("rightKey");
    return p.semiJoin(leftKey, rightKey, p.symbol("match"), Optional.empty(), Optional.empty(), p.values(leftKey), p.aggregation(builder -> builder.globalGrouping().addAggregation(rightKey, expression("count(rightValue)"), ImmutableList.of(BIGINT)).source(p.values(p.symbol("rightValue")))));
}

Example 5 with PlanBuilder

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

the class TestDetermineJoinDistributionType method testGetSourceTablesSizeInBytes.

@Test
public void testGetSourceTablesSizeInBytes() {
    PlanBuilder planBuilder = new PlanBuilder(new PlanNodeIdAllocator(), tester.getMetadata(), tester.getSession());
    Symbol symbol = planBuilder.symbol("col");
    Symbol sourceSymbol1 = planBuilder.symbol("source1");
    Symbol sourceSymbol2 = planBuilder.symbol("soruce2");
    // missing source stats
    assertEquals(getSourceTablesSizeInBytes(planBuilder.values(symbol), noLookup(), node -> PlanNodeStatsEstimate.unknown(), planBuilder.getTypes()), NaN);
    // two source plan nodes
    PlanNodeStatsEstimate sourceStatsEstimate1 = PlanNodeStatsEstimate.builder().setOutputRowCount(10).build();
    PlanNodeStatsEstimate sourceStatsEstimate2 = PlanNodeStatsEstimate.builder().setOutputRowCount(20).build();
    assertEquals(getSourceTablesSizeInBytes(planBuilder.union(ImmutableListMultimap.<Symbol, Symbol>builder().put(symbol, sourceSymbol1).put(symbol, sourceSymbol2).build(), ImmutableList.of(planBuilder.tableScan(ImmutableList.of(sourceSymbol1), ImmutableMap.of(sourceSymbol1, new TestingColumnHandle("col"))), planBuilder.values(new PlanNodeId("valuesNode"), sourceSymbol2))), noLookup(), node -> {
        if (node instanceof TableScanNode) {
            return sourceStatsEstimate1;
        }
        if (node instanceof ValuesNode) {
            return sourceStatsEstimate2;
        }
        return PlanNodeStatsEstimate.unknown();
    }, planBuilder.getTypes()), 270.0);
    // join node
    assertEquals(getSourceTablesSizeInBytes(planBuilder.join(INNER, planBuilder.values(sourceSymbol1), planBuilder.values(sourceSymbol2)), noLookup(), node -> sourceStatsEstimate1, planBuilder.getTypes()), NaN);
}