Examples with LocalPartitionGenerator com.facebook.presto.operator.exchange.LocalPartitionGenerator used on opensource projects

Example 1 with LocalPartitionGenerator

use of com.facebook.presto.operator.exchange.LocalPartitionGenerator in project presto by prestodb.

the class TestOptimizedPartitionedOutputOperator method createOptimizedPartitionedOutputOperator.

private OptimizedPartitionedOutputOperator createOptimizedPartitionedOutputOperator(List<Type> types, boolean replicateAllRows) {
    TestingPartitionedOutputBuffer outputBuffer = createPartitionedOutputBuffer();
    PartitionFunction partitionFunction = new LocalPartitionGenerator(new PrecomputedHashGenerator(0), PARTITION_COUNT);
    if (replicateAllRows) {
        List<Type> replicatedTypes = updateBlockTypesWithHashBlockAndNullBlock(types, false, true);
        return createOptimizedPartitionedOutputOperator(replicatedTypes, ImmutableList.of(0), partitionFunction, outputBuffer, OptionalInt.of(replicatedTypes.size() - 1), MAX_MEMORY);
    } else {
        return createOptimizedPartitionedOutputOperator(types, ImmutableList.of(0), partitionFunction, outputBuffer, OptionalInt.empty(), MAX_MEMORY);
    }
}

Example 2 with LocalPartitionGenerator

use of com.facebook.presto.operator.exchange.LocalPartitionGenerator in project presto by prestodb.

the class TestOptimizedPartitionedOutputOperator method testPartitioned.

private void testPartitioned(List<Type> types, List<Page> pages, DataSize maxMemory, List<Integer> partitionChannel, HashGenerator hashGenerator) {
    TestingPartitionedOutputBuffer outputBuffer = createPartitionedOutputBuffer();
    PartitionFunction partitionFunction = new LocalPartitionGenerator(hashGenerator, PARTITION_COUNT);
    OptimizedPartitionedOutputOperator operator = createOptimizedPartitionedOutputOperator(types, partitionChannel, partitionFunction, outputBuffer, OptionalInt.empty(), maxMemory);
    Map<Integer, List<Page>> expectedPageList = new HashMap<>();
    for (Page page : pages) {
        Map<Integer, List<Integer>> positionsByPartition = new HashMap<>();
        for (int i = 0; i < page.getPositionCount(); i++) {
            int partitionNumber = partitionFunction.getPartition(page, i);
            positionsByPartition.computeIfAbsent(partitionNumber, k -> new ArrayList<>()).add(i);
        }
        for (Map.Entry<Integer, List<Integer>> entry : positionsByPartition.entrySet()) {
            if (!entry.getValue().isEmpty()) {
                expectedPageList.computeIfAbsent(entry.getKey(), k -> new ArrayList<>()).add(copyPositions(page, entry.getValue()));
            }
        }
        operator.addInput(page);
    }
    operator.finish();
    Map<Integer, Page> expectedPages = Maps.transformValues(expectedPageList, outputPages -> mergePages(types, outputPages));
    Map<Integer, Page> actualPages = Maps.transformValues(outputBuffer.getPages(), outputPages -> mergePages(types, outputPages));
    assertEquals(actualPages.size(), expectedPages.size());
    assertEquals(actualPages.keySet(), expectedPages.keySet());
    for (Map.Entry<Integer, Page> entry : expectedPages.entrySet()) {
        int key = entry.getKey();
        assertPageEquals(types, actualPages.get(key), entry.getValue());
    }
}

Example 3 with LocalPartitionGenerator

use of com.facebook.presto.operator.exchange.LocalPartitionGenerator in project presto by prestodb.

the class TestPartitionedOutputOperator method createPartitionedOutputOperator.

private static PartitionedOutputOperator createPartitionedOutputOperator(boolean shouldReplicate) {
    PartitionFunction partitionFunction = new LocalPartitionGenerator(new InterpretedHashGenerator(ImmutableList.of(BIGINT), new int[] { 0 }), PARTITION_COUNT);
    OutputPartitioning outputPartitioning;
    if (shouldReplicate) {
        outputPartitioning = new OutputPartitioning(partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty()), true, OptionalInt.of(0));
    } else {
        outputPartitioning = new OutputPartitioning(partitionFunction, ImmutableList.of(0), ImmutableList.of(Optional.empty(), Optional.empty()), false, OptionalInt.empty());
    }
    PagesSerdeFactory serdeFactory = new PagesSerdeFactory(new BlockEncodingManager(), false);
    DriverContext driverContext = TestingTaskContext.builder(EXECUTOR, SCHEDULER, TEST_SESSION).setMemoryPoolSize(MAX_MEMORY).setQueryMaxTotalMemory(MAX_MEMORY).build().addPipelineContext(0, true, true, false).addDriverContext();
    OutputBuffers buffers = createInitialEmptyOutputBuffers(PARTITIONED);
    for (int partition = 0; partition < PARTITION_COUNT; partition++) {
        buffers = buffers.withBuffer(new OutputBuffers.OutputBufferId(partition), partition);
    }
    PartitionedOutputBuffer buffer = new PartitionedOutputBuffer("task-instance-id", new StateMachine<>("bufferState", SCHEDULER, OPEN, TERMINAL_BUFFER_STATES), buffers.withNoMoreBufferIds(), new DataSize(Long.MAX_VALUE, BYTE), () -> new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"), SCHEDULER);
    buffer.registerLifespanCompletionCallback(ignore -> {
    });
    PartitionedOutputOperator.PartitionedOutputFactory operatorFactory;
    if (shouldReplicate) {
        operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory(buffer, PARTITION_MAX_MEMORY);
        return (PartitionedOutputOperator) operatorFactory.createOutputOperator(0, new PlanNodeId("plan-node-0"), REPLICATION_TYPES, Function.identity(), Optional.of(outputPartitioning), serdeFactory).createOperator(driverContext);
    } else {
        operatorFactory = new PartitionedOutputOperator.PartitionedOutputFactory(buffer, PARTITION_MAX_MEMORY);
        return (PartitionedOutputOperator) operatorFactory.createOutputOperator(0, new PlanNodeId("plan-node-0"), TYPES, Function.identity(), Optional.of(outputPartitioning), serdeFactory).createOperator(driverContext);
    }
}

Example 4 with LocalPartitionGenerator

use of com.facebook.presto.operator.exchange.LocalPartitionGenerator in project presto by prestodb.

the class TestOptimizedPartitionedOutputOperator method testReplicated.

private void testReplicated(List<Type> types, List<Page> pages, DataSize maxMemory) {
    TestingPartitionedOutputBuffer outputBuffer = createPartitionedOutputBuffer();
    PartitionFunction partitionFunction = new LocalPartitionGenerator(new PrecomputedHashGenerator(0), PARTITION_COUNT);
    OptimizedPartitionedOutputOperator operator = createOptimizedPartitionedOutputOperator(types, ImmutableList.of(0), partitionFunction, outputBuffer, OptionalInt.of(types.size() - 1), maxMemory);
    for (Page page : pages) {
        operator.addInput(page);
    }
    operator.finish();
    Map<Integer, List<Page>> acutualPageLists = outputBuffer.getPages();
    assertEquals(acutualPageLists.size(), PARTITION_COUNT);
    Page expectedPage = mergePages(types, pages);
    acutualPageLists.values().forEach(pageList -> assertPageEquals(types, mergePages(types, pageList), expectedPage));
}