Examples with SingleDistributionTaskSource io.trino.execution.scheduler.StageTaskSourceFactory.SingleDistributionTaskSource used on opensource projects

Example 1 with SingleDistributionTaskSource

use of io.trino.execution.scheduler.StageTaskSourceFactory.SingleDistributionTaskSource in project trino by trinodb.

the class TestStageTaskSourceFactory method testSingleDistributionTaskSource.

@Test
public void testSingleDistributionTaskSource() {
    ListMultimap<PlanNodeId, ExchangeSourceHandle> sources = ImmutableListMultimap.<PlanNodeId, ExchangeSourceHandle>builder().put(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 123)).put(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 321)).put(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 222)).build();
    TaskSource taskSource = new SingleDistributionTaskSource(sources, DataSize.of(4, GIGABYTE));
    assertFalse(taskSource.isFinished());
    List<TaskDescriptor> tasks = taskSource.getMoreTasks();
    assertThat(tasks).hasSize(1);
    assertTrue(taskSource.isFinished());
    TaskDescriptor task = tasks.get(0);
    assertThat(task.getNodeRequirements().getCatalogName()).isEmpty();
    assertThat(task.getNodeRequirements().getAddresses()).isEmpty();
    assertEquals(task.getPartitionId(), 0);
    assertEquals(task.getExchangeSourceHandles(), sources);
    assertEquals(task.getSplits(), ImmutableListMultimap.of());
}