Examples with Split io.trino.metadata.Split used on opensource projects

Example 16 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testSplitCount.

@Test
public void testSplitCount() {
    setUpNodes();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    InternalNode chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0);
    TaskId taskId1 = new TaskId(new StageId("test", 1), 1, 0);
    RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()), new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId1));
    TaskId taskId2 = new TaskId(new StageId("test", 1), 2, 0);
    RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask(taskId2, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId2));
    nodeTaskMap.addTask(chosenNode, remoteTask1);
    nodeTaskMap.addTask(chosenNode, remoteTask2);
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode).getCount(), 3);
    remoteTask1.abort();
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode).getCount(), 1);
    remoteTask2.abort();
    assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), PartitionedSplitsInfo.forZeroSplits());
}

Example 17 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testScheduleLocal.

@Test
public void testScheduleLocal() {
    setUpNodes();
    Split split = new Split(CONNECTOR_ID, new TestSplitLocallyAccessible(), Lifespan.taskWide());
    Set<Split> splits = ImmutableSet.of(split);
    Map.Entry<InternalNode, Split> assignment = getOnlyElement(nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments().entries());
    assertEquals(assignment.getKey().getHostAndPort(), split.getAddresses().get(0));
    assertEquals(assignment.getValue(), split);
}

Example 18 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testOptimizedLocalScheduling.

@Test
public void testOptimizedLocalScheduling() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Set<Split> splits = new LinkedHashSet<>();
    // 20 splits with node1 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < 20; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    }
    // computeAssignments just returns a mapping of nodes with splits to be assigned, it does not assign splits
    Multimap<InternalNode, Split> assignments1 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that all 20 splits are being assigned to node1 as optimized-local-scheduling is enabled
    assertEquals(assignments1.size(), 20);
    assertEquals(assignments1.keySet().size(), 2);
    assertTrue(assignments1.keySet().contains(node1));
    assertTrue(assignments1.keySet().contains(node2));
    // 19 splits with node2 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < 19; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromString("10.0.0.1:12")), Lifespan.taskWide()));
    }
    Multimap<InternalNode, Split> assignments2 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that all 39 splits are being assigned (20 splits assigned to node1 and 19 splits assigned to node2)
    assertEquals(assignments2.size(), 39);
    assertEquals(assignments2.keySet().size(), 2);
    assertTrue(assignments2.keySet().contains(node1));
    assertTrue(assignments2.keySet().contains(node2));
    long node1Splits = assignments2.values().stream().map(Split::getConnectorSplit).filter(TestSplitLocal.class::isInstance).count();
    assertEquals(node1Splits, 20);
    long node2Splits = assignments2.values().stream().map(Split::getConnectorSplit).filter(TestSplitRemote.class::isInstance).count();
    assertEquals(node2Splits, 19);
    // 1 split with node1 as local node
    splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    // 1 split with node2 as local node
    splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromString("10.0.0.1:12")), Lifespan.taskWide()));
    // splits now contains 41 splits : 21 with node1 as local node and 20 with node2 as local node
    Multimap<InternalNode, Split> assignments3 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that only 40 splits are being assigned as there is a single task
    assertEquals(assignments3.size(), 40);
    assertEquals(assignments3.keySet().size(), 2);
    assertTrue(assignments3.keySet().contains(node1));
    assertTrue(assignments3.keySet().contains(node2));
    // The first 20 splits have node1 as local, the next 19 have node2 as local, the 40th split has node1 as local and the 41st has node2 as local
    // If optimized-local-scheduling is disabled, the 41st split will be unassigned (the last slot in node2 will be taken up by the 40th split with node1 as local)
    // optimized-local-scheduling ensures that all splits that can be assigned locally will be assigned first
    node1Splits = assignments3.values().stream().map(Split::getConnectorSplit).filter(TestSplitLocal.class::isInstance).count();
    assertEquals(node1Splits, 20);
    node2Splits = assignments3.values().stream().map(Split::getConnectorSplit).filter(TestSplitRemote.class::isInstance).count();
    assertEquals(node2Splits, 20);
}

Example 19 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testRedistributeSplit.

@Test
public void testRedistributeSplit() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    Set<Split> splitsAssignedToNode1 = new LinkedHashSet<>();
    // Node1 to be assigned 12 splits out of which 6 are local to it
    for (int i = 0; i < 6; i++) {
        splitsAssignedToNode1.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
        splitsAssignedToNode1.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    for (Split split : splitsAssignedToNode1) {
        assignment.put(node1, split);
    }
    Set<Split> splitsAssignedToNode2 = new LinkedHashSet<>();
    // Node2 to be assigned 10 splits
    for (int i = 0; i < 10; i++) {
        splitsAssignedToNode2.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    for (Split split : splitsAssignedToNode2) {
        assignment.put(node2, split);
    }
    assertEquals(assignment.get(node1).size(), 12);
    assertEquals(assignment.get(node2).size(), 10);
    ImmutableSetMultimap.Builder<InetAddress, InternalNode> nodesByHost = ImmutableSetMultimap.builder();
    try {
        nodesByHost.put(node1.getInternalAddress(), node1);
        nodesByHost.put(node2.getInternalAddress(), node2);
    } catch (UnknownHostException e) {
        System.out.println("Could not convert the address");
    }
    // Redistribute 1 split from Node 1 to Node 2
    UniformNodeSelector.redistributeSplit(assignment, node1, node2, nodesByHost.build());
    assertEquals(assignment.get(node1).size(), 11);
    assertEquals(assignment.get(node2).size(), 11);
    Set<Split> redistributedSplit = Sets.difference(new HashSet<>(assignment.get(node2)), splitsAssignedToNode2);
    assertEquals(redistributedSplit.size(), 1);
    // Assert that the redistributed split is not a local split in Node 1. This test ensures that redistributeSingleSplit() prioritizes the transfer of a non-local split
    assertTrue(redistributedSplit.iterator().next().getConnectorSplit() instanceof TestSplitRemote);
}

Example 20 with Split

use of io.trino.metadata.Split in project trino by trinodb.

the class TestNodeScheduler method testMaxUnacknowledgedSplitsPerTask.

@Test
public void testMaxUnacknowledgedSplitsPerTask() {
    int maxUnacknowledgedSplitsPerTask = 5;
    nodeSelector = nodeScheduler.createNodeSelector(sessionWithMaxUnacknowledgedSplitsPerTask(maxUnacknowledgedSplitsPerTask), Optional.of(CONNECTOR_ID));
    setUpNodes();
    ImmutableList.Builder<Split> initialSplits = ImmutableList.builder();
    for (int i = 0; i < maxUnacknowledgedSplitsPerTask; i++) {
        initialSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    List<InternalNode> nodes = new ArrayList<>();
    List<MockRemoteTaskFactory.MockRemoteTask> tasks = new ArrayList<>();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int counter = 1;
    for (InternalNode node : nodeManager.getActiveConnectorNodes(CONNECTOR_ID)) {
        // Max out number of unacknowledged splits on each task
        TaskId taskId = new TaskId(new StageId("test", 1), counter, 0);
        counter++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        nodeTaskMap.addTask(node, remoteTask);
        remoteTask.setMaxUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask);
        remoteTask.setUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask);
        nodes.add(node);
        tasks.add(remoteTask);
    }
    // One split per node
    Set<Split> splits = new HashSet<>();
    for (int i = 0; i < nodes.size(); i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    SplitPlacementResult splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    // No splits should have been placed, max unacknowledged was already reached
    assertEquals(splitPlacements.getAssignments().size(), 0);
    // Unblock one task
    MockRemoteTaskFactory.MockRemoteTask taskOne = tasks.get(0);
    taskOne.finishSplits(1);
    taskOne.setUnacknowledgedSplits(taskOne.getUnacknowledgedPartitionedSplitCount() - 1);
    assertTrue(splitPlacements.getBlocked().isDone());
    // Attempt to schedule again, only the node with the unblocked task should be chosen
    splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    assertEquals(splitPlacements.getAssignments().size(), 1);
    assertTrue(splitPlacements.getAssignments().keySet().contains(nodes.get(0)));
    // Make the first node appear to have no splits, unacknowledged splits alone should force the splits to be spread across nodes
    taskOne.clearSplits();
    // Give all tasks with room for 1 unacknowledged split
    tasks.forEach(task -> task.setUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask - 1));
    splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    // One split placed on each node
    assertEquals(splitPlacements.getAssignments().size(), nodes.size());
    assertTrue(splitPlacements.getAssignments().keySet().containsAll(nodes));
}