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

Example 1 with ArbitraryDistributionTaskSource

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

the class TestStageTaskSourceFactory method testArbitraryDistributionTaskSource.

@Test
public void testArbitraryDistributionTaskSource() {
    ExchangeManager splittingExchangeManager = new TestingExchangeManager(true);
    ExchangeManager nonSplittingExchangeManager = new TestingExchangeManager(false);
    TaskSource taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(), ImmutableListMultimap.of(), ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    assertFalse(taskSource.isFinished());
    List<TaskDescriptor> tasks = taskSource.getMoreTasks();
    assertThat(tasks).isEmpty();
    assertTrue(taskSource.isFinished());
    TestingExchangeSourceHandle sourceHandle1 = new TestingExchangeSourceHandle(0, 1);
    TestingExchangeSourceHandle sourceHandle2 = new TestingExchangeSourceHandle(0, 2);
    TestingExchangeSourceHandle sourceHandle3 = new TestingExchangeSourceHandle(0, 3);
    TestingExchangeSourceHandle sourceHandle4 = new TestingExchangeSourceHandle(0, 4);
    TestingExchangeSourceHandle sourceHandle123 = new TestingExchangeSourceHandle(0, 123);
    TestingExchangeSourceHandle sourceHandle321 = new TestingExchangeSourceHandle(0, 321);
    Multimap<PlanNodeId, ExchangeSourceHandle> nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle3);
    Exchange exchange = splittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle3, exchange)), nonReplicatedSources, ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertTrue(taskSource.isFinished());
    assertThat(tasks).hasSize(1);
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 3)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle123);
    exchange = nonSplittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle123, exchange)), nonReplicatedSources, ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 123)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle123, PLAN_NODE_2, sourceHandle321);
    exchange = nonSplittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle123, exchange, sourceHandle321, exchange)), nonReplicatedSources, ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 123)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 321)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle1, PLAN_NODE_1, sourceHandle2, PLAN_NODE_2, sourceHandle4);
    exchange = splittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle1, exchange, sourceHandle2, exchange, sourceHandle4, exchange)), nonReplicatedSources, ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(0, 2)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 3)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle1, PLAN_NODE_1, sourceHandle3, PLAN_NODE_2, sourceHandle4);
    exchange = splittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle1, exchange, sourceHandle3, exchange, sourceHandle4, exchange)), nonReplicatedSources, ImmutableListMultimap.of(), DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 3)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 3)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(3, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    // with replicated sources
    nonReplicatedSources = ImmutableListMultimap.of(PLAN_NODE_1, sourceHandle1, PLAN_NODE_1, sourceHandle2, PLAN_NODE_1, sourceHandle4);
    Multimap<PlanNodeId, ExchangeSourceHandle> replicatedSources = ImmutableListMultimap.of(PLAN_NODE_2, sourceHandle321);
    exchange = splittingExchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    taskSource = new ArbitraryDistributionTaskSource(new IdentityHashMap<>(ImmutableMap.of(sourceHandle1, exchange, sourceHandle2, exchange, sourceHandle4, exchange, sourceHandle321, exchange)), nonReplicatedSources, replicatedSources, DataSize.of(3, BYTE), DataSize.of(4, GIGABYTE));
    tasks = taskSource.getMoreTasks();
    assertEquals(tasks, ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_1, new TestingExchangeSourceHandle(0, 2), PLAN_NODE_2, new TestingExchangeSourceHandle(0, 321)), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(1, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 3), PLAN_NODE_2, sourceHandle321), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE))), new TaskDescriptor(2, ImmutableListMultimap.of(), ImmutableListMultimap.of(PLAN_NODE_1, new TestingExchangeSourceHandle(0, 1), PLAN_NODE_2, sourceHandle321), new NodeRequirements(Optional.empty(), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
}