Examples with ScheduledSplit com.facebook.presto.execution.ScheduledSplit used on opensource projects

Example 21 with ScheduledSplit

use of com.facebook.presto.execution.ScheduledSplit in project presto by prestodb.

the class TestPrestoSparkSourceDistributionSplitAssigner method testAssignSplitsToPartitionWithRandomSplitSizes.

@Test
public void testAssignSplitsToPartitionWithRandomSplitSizes() {
    DataSize maxSplitDataSizePerPartition = new DataSize(2048, BYTE);
    int initialPartitionCount = 3;
    int minSparkInputPartitionCountForAutoTune = 2;
    int maxSparkInputPartitionCountForAutoTune = 5;
    int maxSplitSizeInBytes = 2048;
    AtomicInteger sequenceId = new AtomicInteger();
    for (int i = 0; i < 3; ++i) {
        List<Long> splitSizes = new ArrayList<>(1000);
        for (int j = 0; j < 1000; j++) {
            splitSizes.add(ThreadLocalRandom.current().nextLong((long) (maxSplitSizeInBytes * 1.2)));
        }
        PrestoSparkSplitAssigner assigner = new PrestoSparkSourceDistributionSplitAssigner(new PlanNodeId("test"), createSplitSource(splitSizes), 333, maxSplitDataSizePerPartition.toBytes(), initialPartitionCount, true, minSparkInputPartitionCountForAutoTune, maxSparkInputPartitionCountForAutoTune);
        HashMultimap<Integer, ScheduledSplit> actualAssignment = HashMultimap.create();
        while (true) {
            Optional<SetMultimap<Integer, ScheduledSplit>> assignment = assigner.getNextBatch();
            if (!assignment.isPresent()) {
                break;
            }
            actualAssignment.putAll(assignment.get());
        }
        long expectedSizeInBytes = splitSizes.stream().mapToLong(Long::longValue).sum();
        long actualTotalSizeInBytes = actualAssignment.values().stream().mapToLong(split -> split.getSplit().getConnectorSplit().getSplitSizeInBytes().orElseThrow(() -> new IllegalArgumentException("split size is expected to be present"))).sum();
        // check if all splits got assigned
        assertEquals(expectedSizeInBytes, actualTotalSizeInBytes);
    }
}