Examples with FragmentsEdge io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge used on opensource projects

Example 1 with FragmentsEdge

use of io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge in project trino by trinodb.

the class TestPhasedExecutionSchedule method testBroadcastSourceJoin.

@Test
public void testBroadcastSourceJoin() {
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment joinSourceFragment = createBroadcastJoinPlanFragment("probe", buildFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution joinSourceStage = new TestingStageExecution(joinSourceFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(joinSourceStage, buildStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), joinSourceFragment.getId());
    // single dependency between build and join stages
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactlyInAnyOrder(new FragmentsEdge(buildFragment.getId(), joinSourceFragment.getId()));
    // build stage should start immediately
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId());
    // join stage should start after build stage buffer is full
    buildStage.setAnyTaskBlocked(true);
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), joinSourceFragment.getId());
}

Example 2 with FragmentsEdge

use of io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge 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 3 with FragmentsEdge

use of io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge in project trino by trinodb.

the class TestPhasedExecutionSchedule method testAggregation.

@Test
public void testAggregation() {
    PlanFragment sourceFragment = createTableScanPlanFragment("probe");
    PlanFragment aggregationFragment = createAggregationFragment("aggregation", sourceFragment);
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, REPLICATED, "join", buildFragment, aggregationFragment);
    TestingStageExecution sourceStage = new TestingStageExecution(sourceFragment);
    TestingStageExecution aggregationStage = new TestingStageExecution(aggregationFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(sourceStage, aggregationStage, buildStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId(), joinFragment.getId());
    // aggregation and source stage should start immediately, join stage should wait for build stage to complete
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactly(new FragmentsEdge(buildFragment.getId(), joinFragment.getId()));
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
}

Example 4 with FragmentsEdge

use of io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge in project trino by trinodb.

the class TestPhasedExecutionSchedule method testDependentStageAbortedBeforeStarted.

@Test
public void testDependentStageAbortedBeforeStarted() {
    PlanFragment sourceFragment = createTableScanPlanFragment("probe");
    PlanFragment aggregationFragment = createAggregationFragment("aggregation", sourceFragment);
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, REPLICATED, "join", buildFragment, aggregationFragment);
    TestingStageExecution sourceStage = new TestingStageExecution(sourceFragment);
    TestingStageExecution aggregationStage = new TestingStageExecution(aggregationFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(sourceStage, aggregationStage, buildStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId(), joinFragment.getId());
    // aggregation and source stage should start immediately, join stage should wait for build stage to complete
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactly(new FragmentsEdge(buildFragment.getId(), joinFragment.getId()));
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
    // abort non-active join stage
    joinStage.setState(ABORTED);
    // dependencies finish
    buildStage.setState(FINISHED);
    aggregationStage.setState(FINISHED);
    sourceStage.setState(FINISHED);
    // join stage already aborted. Whole schedule should be marked as finished
    schedule.schedule();
    assertThat(schedule.isFinished()).isTrue();
}

Example 5 with FragmentsEdge

use of io.trino.execution.scheduler.policy.PhasedExecutionSchedule.FragmentsEdge in project trino by trinodb.

the class TestPhasedExecutionSchedule method testPartitionedJoin.

@Test
public void testPartitionedJoin() {
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment probeFragment = createTableScanPlanFragment("probe");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, PARTITIONED, "join", buildFragment, probeFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution probeStage = new TestingStageExecution(probeFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(buildStage, probeStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), probeFragment.getId(), joinFragment.getId());
    // single dependency between build and probe stages
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactlyInAnyOrder(new FragmentsEdge(buildFragment.getId(), probeFragment.getId()));
    // build and join stage should start immediately
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), joinFragment.getId());
    // probe stage should start after build stage is completed
    ListenableFuture<Void> rescheduleFuture = schedule.getRescheduleFuture().orElseThrow();
    assertThat(rescheduleFuture).isNotDone();
    buildStage.setState(FLUSHING);
    assertThat(rescheduleFuture).isDone();
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).containsExactly(joinFragment.getId(), probeFragment.getId());
    // make sure scheduler finishes
    rescheduleFuture = schedule.getRescheduleFuture().orElseThrow();
    assertThat(rescheduleFuture).isNotDone();
    probeStage.setState(FINISHED);
    assertThat(rescheduleFuture).isNotDone();
    joinStage.setState(FINISHED);
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).isEmpty();
    assertThat(schedule.isFinished()).isTrue();
}