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

Example 26 with PlanFragment

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBatched.

@Test
public void testScheduleSplitsBatched() {
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(60, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 7, STAGE);
    for (int i = 0; i <= (60 / 7); i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // finishes when last split is fetched
        if (i == (60 / 7)) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // never blocks
        assertTrue(scheduleResult.getBlocked().isDone());
        // first three splits create new tasks
        assertEquals(scheduleResult.getNewTasks().size(), i == 0 ? 3 : 0);
        assertEquals(stage.getAllTasks().size(), 3);
        assertPartitionedSplitCount(stage, min((i + 1) * 7, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    stage.abort();
}

Example 27 with PlanFragment

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

the class TestSourcePartitionedScheduler method testWorkerBalancedSplitAssignment.

@Test
public void testWorkerBalancedSplitAssignment() {
    // use private node manager so we can add a node later
    InMemoryNodeManager nodeManager = new InMemoryNodeManager();
    nodeManager.addNode(CONNECTOR_ID, new InternalNode("other1", URI.create("http://127.0.0.1:11"), NodeVersion.UNKNOWN, false), new InternalNode("other2", URI.create("http://127.0.0.1:12"), NodeVersion.UNKNOWN, false), new InternalNode("other3", URI.create("http://127.0.0.1:13"), NodeVersion.UNKNOWN, false));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    // Schedule 15 splits - there are 3 nodes, each node should get 5 splits
    PlanFragment firstPlan = createFragment();
    StageExecution firstStage = createStageExecution(firstPlan, nodeTaskMap);
    StageScheduler firstScheduler = getSourcePartitionedScheduler(createFixedSplitSource(15, TestingSplit::createRemoteSplit), firstStage, nodeManager, nodeTaskMap, 200, NODE);
    ScheduleResult scheduleResult = firstScheduler.schedule();
    assertEffectivelyFinished(scheduleResult, firstScheduler);
    assertTrue(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 3);
    assertEquals(firstStage.getAllTasks().size(), 3);
    for (RemoteTask remoteTask : firstStage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 5);
    }
    // Add new node
    InternalNode additionalNode = new InternalNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, additionalNode);
    // Schedule 5 splits in another query. Since the new node does not have any splits, all 5 splits are assigned to the new node
    PlanFragment secondPlan = createFragment();
    StageExecution secondStage = createStageExecution(secondPlan, nodeTaskMap);
    StageScheduler secondScheduler = getSourcePartitionedScheduler(createFixedSplitSource(5, TestingSplit::createRemoteSplit), secondStage, nodeManager, nodeTaskMap, 200, NODE);
    scheduleResult = secondScheduler.schedule();
    assertEffectivelyFinished(scheduleResult, secondScheduler);
    assertTrue(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 1);
    assertEquals(secondStage.getAllTasks().size(), 1);
    RemoteTask task = secondStage.getAllTasks().get(0);
    assertEquals(task.getPartitionedSplitsInfo().getCount(), 5);
    firstStage.abort();
    secondStage.abort();
}

Example 28 with PlanFragment

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

the class TestSourcePartitionedScheduler method testScheduleSplitsOneAtATime.

@Test
public void testScheduleSplitsOneAtATime() {
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(60, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 1, STAGE);
    for (int i = 0; i < 60; i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // only finishes when last split is fetched
        if (i == 59) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // never blocks
        assertTrue(scheduleResult.getBlocked().isDone());
        // first three splits create new tasks
        assertEquals(scheduleResult.getNewTasks().size(), i < 3 ? 1 : 0);
        assertEquals(stage.getAllTasks().size(), i < 3 ? i + 1 : 3);
        assertPartitionedSplitCount(stage, min(i + 1, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    stage.abort();
}

Example 29 with PlanFragment

use of io.trino.sql.planner.PlanFragment 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);
}

Example 30 with PlanFragment

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

the class TestFaultTolerantStageScheduler method createPlanFragment.

private PlanFragment createPlanFragment() {
    Symbol probeColumnSymbol = new Symbol("probe_column");
    Symbol buildColumnSymbol = new Symbol("build_column");
    TableScanNode tableScan = new TableScanNode(TABLE_SCAN_NODE_ID, TEST_TABLE_HANDLE, ImmutableList.of(probeColumnSymbol), ImmutableMap.of(probeColumnSymbol, new TestingColumnHandle("column")), TupleDomain.none(), Optional.empty(), false, Optional.empty());
    RemoteSourceNode remoteSource = new RemoteSourceNode(new PlanNodeId("remote_source_id"), ImmutableList.of(SOURCE_FRAGMENT_ID_1, SOURCE_FRAGMENT_ID_2), ImmutableList.of(buildColumnSymbol), Optional.empty(), REPLICATE, TASK);
    return new PlanFragment(FRAGMENT_ID, new JoinNode(new PlanNodeId("join_id"), INNER, tableScan, remoteSource, ImmutableList.of(), tableScan.getOutputSymbols(), remoteSource.getOutputSymbols(), false, Optional.empty(), Optional.empty(), Optional.empty(), Optional.of(REPLICATED), Optional.empty(), ImmutableMap.of(), Optional.empty()), ImmutableMap.of(probeColumnSymbol, VARCHAR, buildColumnSymbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(TABLE_SCAN_NODE_ID), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(probeColumnSymbol, buildColumnSymbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
}