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

Example 11 with SymbolAllocator

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

the class TestRemoveUnsupportedDynamicFilters method removeUnsupportedDynamicFilters.

private PlanNode removeUnsupportedDynamicFilters(PlanNode root) {
    return getQueryRunner().inTransaction(session -> {
        // metadata.getCatalogHandle() registers the catalog for the transaction
        session.getCatalog().ifPresent(catalog -> metadata.getCatalogHandle(session, catalog));
        PlanNode rewrittenPlan = new RemoveUnsupportedDynamicFilters(plannerContext).optimize(root, session, builder.getTypes(), new SymbolAllocator(), new PlanNodeIdAllocator(), WarningCollector.NOOP);
        new DynamicFiltersChecker().validate(rewrittenPlan, session, plannerContext, createTestingTypeAnalyzer(plannerContext), builder.getTypes(), WarningCollector.NOOP);
        return rewrittenPlan;
    });
}

Example 12 with SymbolAllocator

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

the class RuleAssert method applyRule.

private RuleApplication applyRule() {
    SymbolAllocator symbolAllocator = new SymbolAllocator(types.allTypes());
    Memo memo = new Memo(idAllocator, plan);
    Lookup lookup = Lookup.from(planNode -> Stream.of(memo.resolve(planNode)));
    PlanNode memoRoot = memo.getNode(memo.getRootGroup());
    return inTransaction(session -> applyRule(rule, memoRoot, ruleContext(statsCalculator, costCalculator, symbolAllocator, memo, lookup, session)));
}

Example 13 with SymbolAllocator

use of io.trino.sql.planner.SymbolAllocator 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 14 with SymbolAllocator

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

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

the class TestSourcePartitionedScheduler method testDynamicFiltersUnblockedOnBlockedBuildSource.

@Test
public void testDynamicFiltersUnblockedOnBlockedBuildSource() {
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(nodeManager, new NodeSchedulerConfig().setIncludeCoordinator(false), nodeTaskMap));
    DynamicFilterService dynamicFilterService = new DynamicFilterService(metadata, functionManager, typeOperators, new DynamicFilterConfig());
    dynamicFilterService.registerQuery(QUERY_ID, TEST_SESSION, ImmutableSet.of(DYNAMIC_FILTER_ID), ImmutableSet.of(DYNAMIC_FILTER_ID), ImmutableSet.of(DYNAMIC_FILTER_ID));
    StageScheduler scheduler = newSourcePartitionedSchedulerAsStageScheduler(stage, TABLE_SCAN_NODE_ID, new ConnectorAwareSplitSource(CONNECTOR_ID, createBlockedSplitSource()), new DynamicSplitPlacementPolicy(nodeScheduler.createNodeSelector(session, Optional.of(CONNECTOR_ID)), stage::getAllTasks), 2, dynamicFilterService, new TableExecuteContextManager(), () -> true);
    SymbolAllocator symbolAllocator = new SymbolAllocator();
    Symbol symbol = symbolAllocator.newSymbol("DF_SYMBOL1", BIGINT);
    DynamicFilter dynamicFilter = dynamicFilterService.createDynamicFilter(QUERY_ID, ImmutableList.of(new DynamicFilters.Descriptor(DYNAMIC_FILTER_ID, symbol.toSymbolReference())), ImmutableMap.of(symbol, new TestingColumnHandle("probeColumnA")), symbolAllocator.getTypes());
    // make sure dynamic filtering collecting task was created immediately
    assertEquals(stage.getState(), PLANNED);
    scheduler.start();
    assertEquals(stage.getAllTasks().size(), 1);
    assertEquals(stage.getState(), SCHEDULING);
    // make sure dynamic filter is initially blocked
    assertFalse(dynamicFilter.isBlocked().isDone());
    // make sure dynamic filter is unblocked due to build side source tasks being blocked
    ScheduleResult scheduleResult = scheduler.schedule();
    assertTrue(dynamicFilter.isBlocked().isDone());
    // no new probe splits should be scheduled
    assertEquals(scheduleResult.getSplitsScheduled(), 0);
}