Examples with ExchangeContext io.trino.spi.exchange.ExchangeContext used on opensource projects

Example 1 with ExchangeContext

use of io.trino.spi.exchange.ExchangeContext in project trino by trinodb.

the class AbstractTestExchangeManager method testLargePages.

@Test
public void testLargePages() throws Exception {
    Exchange exchange = exchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 3);
    ExchangeSinkHandle sinkHandle0 = exchange.addSink(0);
    ExchangeSinkHandle sinkHandle1 = exchange.addSink(1);
    ExchangeSinkHandle sinkHandle2 = exchange.addSink(2);
    exchange.noMoreSinks();
    ExchangeSinkInstanceHandle sinkInstanceHandle = exchange.instantiateSink(sinkHandle0, 0);
    writeData(sinkInstanceHandle, new ImmutableListMultimap.Builder<Integer, String>().putAll(0, ImmutableList.of(SMALL_PAGE)).putAll(1, ImmutableList.of(MAX_PAGE, MEDIUM_PAGE)).putAll(2, ImmutableList.of()).build(), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle1, 0);
    writeData(sinkInstanceHandle, new ImmutableListMultimap.Builder<Integer, String>().putAll(0, ImmutableList.of(MEDIUM_PAGE)).putAll(1, ImmutableList.of(LARGE_PAGE)).putAll(2, ImmutableList.of(SMALL_PAGE)).build(), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle2, 0);
    writeData(sinkInstanceHandle, new ImmutableListMultimap.Builder<Integer, String>().putAll(0, ImmutableList.of(LARGE_PAGE, MAX_PAGE)).putAll(1, ImmutableList.of(SMALL_PAGE)).putAll(2, ImmutableList.of(MAX_PAGE, LARGE_PAGE, MEDIUM_PAGE)).build(), true);
    exchange.sinkFinished(sinkInstanceHandle);
    List<ExchangeSourceHandle> partitionHandles = exchange.getSourceHandles().get();
    assertThat(partitionHandles).hasSize(3);
    Map<Integer, ExchangeSourceHandle> partitions = partitionHandles.stream().collect(toImmutableMap(ExchangeSourceHandle::getPartitionId, Function.identity()));
    assertThat(readData(partitions.get(0))).containsExactlyInAnyOrder(SMALL_PAGE, MEDIUM_PAGE, LARGE_PAGE, MAX_PAGE);
    assertThat(readData(partitions.get(1))).containsExactlyInAnyOrder(SMALL_PAGE, MEDIUM_PAGE, LARGE_PAGE, MAX_PAGE);
    assertThat(readData(partitions.get(2))).containsExactlyInAnyOrder(SMALL_PAGE, MEDIUM_PAGE, LARGE_PAGE, MAX_PAGE);
    exchange.close();
}

Example 2 with ExchangeContext

use of io.trino.spi.exchange.ExchangeContext in project trino by trinodb.

the class AbstractTestExchangeManager method testHappyPath.

@Test
public void testHappyPath() throws Exception {
    Exchange exchange = exchangeManager.createExchange(new ExchangeContext(new QueryId("query"), createRandomExchangeId()), 2);
    ExchangeSinkHandle sinkHandle0 = exchange.addSink(0);
    ExchangeSinkHandle sinkHandle1 = exchange.addSink(1);
    ExchangeSinkHandle sinkHandle2 = exchange.addSink(2);
    exchange.noMoreSinks();
    ExchangeSinkInstanceHandle sinkInstanceHandle = exchange.instantiateSink(sinkHandle0, 0);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "0-0-0", 1, "0-1-0", 0, "0-0-1", 1, "0-1-1"), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle0, 1);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "0-0-0", 1, "0-1-0", 0, "0-0-1", 1, "0-1-1"), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle0, 2);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "failed", 1, "another failed"), false);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle1, 0);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "1-0-0", 1, "1-1-0", 0, "1-0-1", 1, "1-1-1"), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle1, 1);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "1-0-0", 1, "1-1-0", 0, "1-0-1", 1, "1-1-1"), true);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle1, 2);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "more failed", 1, "another failed"), false);
    exchange.sinkFinished(sinkInstanceHandle);
    sinkInstanceHandle = exchange.instantiateSink(sinkHandle2, 2);
    writeData(sinkInstanceHandle, ImmutableListMultimap.of(0, "2-0-0", 1, "2-1-0"), true);
    exchange.sinkFinished(sinkInstanceHandle);
    List<ExchangeSourceHandle> partitionHandles = exchange.getSourceHandles().get();
    assertThat(partitionHandles).hasSize(2);
    Map<Integer, ExchangeSourceHandle> partitions = partitionHandles.stream().collect(toImmutableMap(ExchangeSourceHandle::getPartitionId, Function.identity()));
    assertThat(readData(partitions.get(0))).containsExactlyInAnyOrder("0-0-0", "0-0-1", "1-0-0", "1-0-1", "2-0-0");
    assertThat(readData(partitions.get(1))).containsExactlyInAnyOrder("0-1-0", "0-1-1", "1-1-0", "1-1-1", "2-1-0");
    exchange.close();
}

Example 3 with ExchangeContext

use of io.trino.spi.exchange.ExchangeContext 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)))));
}