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

Example 21 with PlanFragment

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

the class TestPhasedExecutionSchedule method testStageWithBroadcastAndPartitionedJoin.

@Test
public void testStageWithBroadcastAndPartitionedJoin() {
    PlanFragment broadcastBuildFragment = createTableScanPlanFragment("broadcast_build");
    PlanFragment partitionedBuildFragment = createTableScanPlanFragment("partitioned_build");
    PlanFragment probeFragment = createTableScanPlanFragment("probe");
    PlanFragment joinFragment = createBroadcastAndPartitionedJoinPlanFragment("join", broadcastBuildFragment, partitionedBuildFragment, probeFragment);
    TestingStageExecution broadcastBuildStage = new TestingStageExecution(broadcastBuildFragment);
    TestingStageExecution partitionedBuildStage = new TestingStageExecution(partitionedBuildFragment);
    TestingStageExecution probeStage = new TestingStageExecution(probeFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(broadcastBuildStage, partitionedBuildStage, probeStage, joinStage), dynamicFilterService);
    // join stage should start immediately because partitioned join forces that
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactlyInAnyOrder(new FragmentsEdge(broadcastBuildFragment.getId(), probeFragment.getId()), new FragmentsEdge(partitionedBuildFragment.getId(), probeFragment.getId()), new FragmentsEdge(broadcastBuildFragment.getId(), joinFragment.getId()));
    assertThat(getActiveFragments(schedule)).containsExactly(partitionedBuildFragment.getId(), broadcastBuildFragment.getId(), joinFragment.getId());
    // completing single build dependency shouldn't cause probe stage to start
    broadcastBuildStage.setState(FLUSHING);
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).containsExactly(partitionedBuildFragment.getId(), joinFragment.getId());
    // completing all build dependencies should cause probe stage to start
    partitionedBuildStage.setState(FLUSHING);
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).containsExactly(joinFragment.getId(), probeFragment.getId());
}

Example 22 with PlanFragment

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

the class SqlQueryExecution method planDistribution.

private void planDistribution(PlanRoot plan) {
    // if query was canceled, skip creating scheduler
    if (stateMachine.isDone()) {
        return;
    }
    // record output field
    PlanFragment rootFragment = plan.getRoot().getFragment();
    stateMachine.setColumns(((OutputNode) rootFragment.getRoot()).getColumnNames(), rootFragment.getTypes());
    // build the stage execution objects (this doesn't schedule execution)
    SqlQueryScheduler scheduler = new SqlQueryScheduler(stateMachine, plan.getRoot(), nodePartitioningManager, nodeScheduler, nodeAllocatorService, partitionMemoryEstimator, remoteTaskFactory, plan.isSummarizeTaskInfos(), scheduleSplitBatchSize, queryExecutor, schedulerExecutor, failureDetector, nodeTaskMap, executionPolicy, schedulerStats, dynamicFilterService, tableExecuteContextManager, plannerContext.getMetadata(), splitSourceFactory, coordinatorTaskManager, exchangeManagerRegistry, taskSourceFactory, taskDescriptorStorage);
    queryScheduler.set(scheduler);
    // directly since the callback may have already fired
    if (stateMachine.isDone()) {
        scheduler.abort();
        queryScheduler.set(null);
    }
}

Example 23 with PlanFragment

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

the class TestStageStateMachine method createValuesPlan.

private static PlanFragment createValuesPlan() {
    Symbol symbol = new Symbol("column");
    PlanNodeId valuesNodeId = new PlanNodeId("plan");
    PlanFragment planFragment = new PlanFragment(new PlanFragmentId("plan"), new ValuesNode(valuesNodeId, ImmutableList.of(symbol), ImmutableList.of(new Row(ImmutableList.of(new StringLiteral("foo"))))), ImmutableMap.of(symbol, VARCHAR), SOURCE_DISTRIBUTION, ImmutableList.of(valuesNodeId), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
    return planFragment;
}

Example 24 with PlanFragment

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

the class TestSourcePartitionedScheduler method testScheduleSplitsBlock.

@Test
public void testScheduleSplitsBlock() {
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(createFixedSplitSource(80, TestingSplit::createRemoteSplit), stage, nodeManager, nodeTaskMap, 1, STAGE);
    // schedule first 60 splits, which will cause the scheduler to block
    for (int i = 0; i <= 60; i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        assertFalse(scheduleResult.isFinished());
        // blocks at 20 per node
        assertEquals(scheduleResult.getBlocked().isDone(), i != 60);
        // 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);
    }
    // todo rewrite MockRemoteTask to fire a tate transition when splits are cleared, and then validate blocked future completes
    // drop the 20 splits from one node
    ((MockRemoteTask) stage.getAllTasks().get(0)).clearSplits();
    // schedule remaining 20 splits
    for (int i = 0; i < 20; i++) {
        ScheduleResult scheduleResult = scheduler.schedule();
        // finishes when last split is fetched
        if (i == 19) {
            assertEffectivelyFinished(scheduleResult, scheduler);
        } else {
            assertFalse(scheduleResult.isFinished());
        }
        // does not block again
        assertTrue(scheduleResult.getBlocked().isDone());
        // no additional tasks will be created
        assertEquals(scheduleResult.getNewTasks().size(), 0);
        assertEquals(stage.getAllTasks().size(), 3);
        // we dropped 20 splits so start at 40 and count to 60
        assertPartitionedSplitCount(stage, min(i + 41, 60));
    }
    for (RemoteTask remoteTask : stage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 20);
    }
    stage.abort();
}

Example 25 with PlanFragment

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

the class TestSourcePartitionedScheduler method testScheduleSlowSplitSource.

@Test
public void testScheduleSlowSplitSource() {
    QueuedSplitSource queuedSplitSource = new QueuedSplitSource(TestingSplit::createRemoteSplit);
    PlanFragment plan = createFragment();
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    StageExecution stage = createStageExecution(plan, nodeTaskMap);
    StageScheduler scheduler = getSourcePartitionedScheduler(queuedSplitSource, stage, nodeManager, nodeTaskMap, 1, STAGE);
    // schedule with no splits - will block
    ScheduleResult scheduleResult = scheduler.schedule();
    assertFalse(scheduleResult.isFinished());
    assertFalse(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 0);
    assertEquals(stage.getAllTasks().size(), 0);
    queuedSplitSource.addSplits(1);
    assertTrue(scheduleResult.getBlocked().isDone());
}