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

Example 6 with ExchangeSourceHandle

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

the class TestStageTaskSourceFactory method testSourceDistributionTaskSource.

@Test
public void testSourceDistributionTaskSource() {
    TaskSource taskSource = createSourceDistributionTaskSource(ImmutableList.of(), ImmutableListMultimap.of(), 2, 0, 3 * STANDARD_WEIGHT, 1000);
    assertFalse(taskSource.isFinished());
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of());
    assertTrue(taskSource.isFinished());
    Split split1 = createSplit(1);
    Split split2 = createSplit(2);
    Split split3 = createSplit(3);
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(split1), ImmutableListMultimap.of(), 2, 0, 2 * STANDARD_WEIGHT, 1000);
    assertEquals(taskSource.getMoreTasks(), ImmutableList.of(new TaskDescriptor(0, ImmutableListMultimap.of(PLAN_NODE_1, split1), ImmutableListMultimap.of(), new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE)))));
    assertTrue(taskSource.isFinished());
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(split1, split2, split3), ImmutableListMultimap.of(), 3, 0, 2 * STANDARD_WEIGHT, 1000);
    List<TaskDescriptor> tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(2);
    assertThat(tasks.get(0).getSplits().values()).hasSize(2);
    assertThat(tasks.get(1).getSplits().values()).hasSize(1);
    assertThat(tasks).allMatch(taskDescriptor -> taskDescriptor.getNodeRequirements().equals(new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))));
    assertThat(tasks).allMatch(taskDescriptor -> taskDescriptor.getExchangeSourceHandles().isEmpty());
    assertThat(flattenSplits(tasks)).hasSameEntriesAs(ImmutableMultimap.of(PLAN_NODE_1, split1, PLAN_NODE_1, split2, PLAN_NODE_1, split3));
    assertTrue(taskSource.isFinished());
    ImmutableListMultimap<PlanNodeId, ExchangeSourceHandle> replicatedSources = ImmutableListMultimap.of(PLAN_NODE_2, new TestingExchangeSourceHandle(0, 1));
    taskSource = createSourceDistributionTaskSource(ImmutableList.of(split1, split2, split3), replicatedSources, 2, 0, 2 * STANDARD_WEIGHT, 1000);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(2);
    assertThat(tasks.get(0).getSplits().values()).hasSize(2);
    assertThat(tasks.get(1).getSplits().values()).hasSize(1);
    assertThat(tasks).allMatch(taskDescriptor -> taskDescriptor.getNodeRequirements().equals(new NodeRequirements(Optional.of(CATALOG), ImmutableSet.of(), DataSize.of(4, GIGABYTE))));
    assertThat(tasks).allMatch(taskDescriptor -> taskDescriptor.getExchangeSourceHandles().equals(replicatedSources));
    assertThat(flattenSplits(tasks)).hasSameEntriesAs(ImmutableMultimap.of(PLAN_NODE_1, split1, PLAN_NODE_1, split2, PLAN_NODE_1, split3));
    assertTrue(taskSource.isFinished());
    // non remotely accessible splits
    ImmutableList<Split> splits = ImmutableList.of(createSplit(1, "host1:8080", "host2:8080"), createSplit(2, "host2:8080"), createSplit(3, "host1:8080", "host3:8080"), createSplit(4, "host3:8080", "host1:8080"), createSplit(5, "host1:8080", "host2:8080"), createSplit(6, "host2:8080", "host3:8080"), createSplit(7, "host3:8080", "host4:8080"));
    taskSource = createSourceDistributionTaskSource(splits, ImmutableListMultimap.of(), 3, 0, 2 * STANDARD_WEIGHT, 1000);
    tasks = readAllTasks(taskSource);
    assertThat(tasks).hasSize(4);
    assertThat(tasks.stream()).allMatch(taskDescriptor -> taskDescriptor.getExchangeSourceHandles().isEmpty());
    assertThat(flattenSplits(tasks)).hasSameEntriesAs(Multimaps.index(splits, split -> PLAN_NODE_1));
    assertThat(tasks).allMatch(task -> task.getSplits().values().stream().allMatch(split -> {
        HostAddress requiredAddress = getOnlyElement(task.getNodeRequirements().getAddresses());
        return split.getAddresses().contains(requiredAddress);
    }));
    assertTrue(taskSource.isFinished());
}

Example 7 with ExchangeSourceHandle

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

Example 8 with ExchangeSourceHandle

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

the class TestingTaskSourceFactory method getHandlesForRemoteSources.

private static ListMultimap<PlanNodeId, ExchangeSourceHandle> getHandlesForRemoteSources(List<RemoteSourceNode> remoteSources, Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles) {
    ImmutableListMultimap.Builder<PlanNodeId, ExchangeSourceHandle> result = ImmutableListMultimap.builder();
    for (RemoteSourceNode remoteSource : remoteSources) {
        checkArgument(remoteSource.getExchangeType() == REPLICATE, "expected exchange type to be REPLICATE, got: %s", remoteSource.getExchangeType());
        for (PlanFragmentId fragmentId : remoteSource.getSourceFragmentIds()) {
            Collection<ExchangeSourceHandle> handles = requireNonNull(exchangeSourceHandles.get(fragmentId), () -> "exchange source handle is missing for fragment: " + fragmentId);
            checkArgument(handles.size() == 1, "single exchange source handle is expected, got: %s", handles);
            result.putAll(remoteSource.getId(), handles);
        }
    }
    return result.build();
}

Example 9 with ExchangeSourceHandle

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

the class StageTaskSourceFactory method getInputsForRemoteSources.

private static ListMultimap<PlanNodeId, ExchangeSourceHandle> getInputsForRemoteSources(List<RemoteSourceNode> remoteSources, Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles) {
    ImmutableListMultimap.Builder<PlanNodeId, ExchangeSourceHandle> result = ImmutableListMultimap.builder();
    for (RemoteSourceNode remoteSource : remoteSources) {
        for (PlanFragmentId fragmentId : remoteSource.getSourceFragmentIds()) {
            Collection<ExchangeSourceHandle> handles = requireNonNull(exchangeSourceHandles.get(fragmentId), () -> "exchange source handle is missing for fragment: " + fragmentId);
            if (remoteSource.getExchangeType() == GATHER || remoteSource.getExchangeType() == REPLICATE) {
                checkArgument(handles.size() <= 1, "at most 1 exchange source handle is expected, got: %s", handles);
            }
            result.putAll(remoteSource.getId(), handles);
        }
    }
    return result.build();
}

Example 10 with ExchangeSourceHandle

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

the class FileSystemExchange method createExchangeSourceHandles.

private List<ExchangeSourceHandle> createExchangeSourceHandles() {
    Multimap<Integer, FileStatus> partitionFiles = ArrayListMultimap.create();
    List<Integer> finishedTaskPartitions;
    synchronized (this) {
        finishedTaskPartitions = ImmutableList.copyOf(finishedSinks);
    }
    for (Integer taskPartition : finishedTaskPartitions) {
        URI committedAttemptPath = getCommittedAttemptPath(taskPartition);
        Map<Integer, FileStatus> partitions = getCommittedPartitions(committedAttemptPath);
        partitions.forEach(partitionFiles::put);
    }
    ImmutableList.Builder<ExchangeSourceHandle> result = ImmutableList.builder();
    for (Integer partitionId : partitionFiles.keySet()) {
        result.add(new FileSystemExchangeSourceHandle(partitionId, ImmutableList.copyOf(partitionFiles.get(partitionId)), secretKey.map(SecretKey::getEncoded)));
    }
    return result.build();
}