Examples with PartitionFunction com.facebook.presto.operator.PartitionFunction used on opensource projects

Example 1 with PartitionFunction

use of com.facebook.presto.operator.PartitionFunction in project presto by prestodb.

the class LocalExecutionPlanner method plan.

public LocalExecutionPlan plan(Session session, PlanNode plan, Map<Symbol, Type> types, PartitioningScheme partitioningScheme, OutputBuffer outputBuffer) {
    List<Symbol> outputLayout = partitioningScheme.getOutputLayout();
    if (partitioningScheme.getPartitioning().getHandle().equals(FIXED_BROADCAST_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(SINGLE_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(COORDINATOR_DISTRIBUTION)) {
        return plan(session, plan, outputLayout, types, new TaskOutputFactory(outputBuffer));
    }
    // We can convert the symbols directly into channels, because the root must be a sink and therefore the layout is fixed
    List<Integer> partitionChannels;
    List<Optional<NullableValue>> partitionConstants;
    List<Type> partitionChannelTypes;
    if (partitioningScheme.getHashColumn().isPresent()) {
        partitionChannels = ImmutableList.of(outputLayout.indexOf(partitioningScheme.getHashColumn().get()));
        partitionConstants = ImmutableList.of(Optional.empty());
        partitionChannelTypes = ImmutableList.of(BIGINT);
    } else {
        partitionChannels = partitioningScheme.getPartitioning().getArguments().stream().map(ArgumentBinding::getColumn).map(outputLayout::indexOf).collect(toImmutableList());
        partitionConstants = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument.isConstant()) {
                return Optional.of(argument.getConstant());
            }
            return Optional.<NullableValue>empty();
        }).collect(toImmutableList());
        partitionChannelTypes = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument.isConstant()) {
                return argument.getConstant().getType();
            }
            return types.get(argument.getColumn());
        }).collect(toImmutableList());
    }
    PartitionFunction partitionFunction = nodePartitioningManager.getPartitionFunction(session, partitioningScheme, partitionChannelTypes);
    OptionalInt nullChannel = OptionalInt.empty();
    Set<Symbol> partitioningColumns = partitioningScheme.getPartitioning().getColumns();
    // partitioningColumns expected to have one column in the normal case, and zero columns when partitioning on a constant
    checkArgument(!partitioningScheme.isReplicateNulls() || partitioningColumns.size() <= 1);
    if (partitioningScheme.isReplicateNulls() && partitioningColumns.size() == 1) {
        nullChannel = OptionalInt.of(outputLayout.indexOf(getOnlyElement(partitioningColumns)));
    }
    return plan(session, plan, outputLayout, types, new PartitionedOutputFactory(partitionFunction, partitionChannels, partitionConstants, nullChannel, outputBuffer, maxPagePartitioningBufferSize));
}

Example 2 with PartitionFunction

use of com.facebook.presto.operator.PartitionFunction 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 3 with PartitionFunction

use of com.facebook.presto.operator.PartitionFunction 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 4 with PartitionFunction

use of com.facebook.presto.operator.PartitionFunction in project presto by prestodb.

the class TestLocalExchange method testCreatePartitionFunction.

@Test
public void testCreatePartitionFunction() {
    int partitionCount = 10;
    PartitioningProviderManager partitioningProviderManager = new PartitioningProviderManager();
    partitioningProviderManager.addPartitioningProvider(new ConnectorId("prism"), new ConnectorNodePartitioningProvider() {

        @Override
        public ConnectorBucketNodeMap getBucketNodeMap(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List<Node> sortedNodes) {
            return createBucketNodeMap(Stream.generate(() -> sortedNodes).flatMap(List::stream).limit(10).collect(toImmutableList()), SOFT_AFFINITY);
        }

        @Override
        public ToIntFunction<ConnectorSplit> getSplitBucketFunction(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle) {
            return null;
        }

        @Override
        public BucketFunction getBucketFunction(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle, List<Type> partitionChannelTypes, int bucketCount) {
            return (Page page, int position) -> partitionCount;
        }

        @Override
        public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle) {
            return 10;
        }
    });
    PartitioningHandle partitioningHandle = new PartitioningHandle(Optional.of(new ConnectorId("prism")), Optional.of(new ConnectorTransactionHandle() {

        @Override
        public int hashCode() {
            return super.hashCode();
        }

        @Override
        public boolean equals(Object obj) {
            return super.equals(obj);
        }
    }), new ConnectorPartitioningHandle() {

        @Override
        public boolean isSingleNode() {
            return false;
        }

        @Override
        public boolean isCoordinatorOnly() {
            return false;
        }
    });
    PartitionFunction partitionFunction = createPartitionFunction(partitioningProviderManager, session, partitioningHandle, 600, ImmutableList.of(), false);
    assertEquals(partitionFunction.getPartitionCount(), partitionCount);
}

Example 5 with PartitionFunction

use of com.facebook.presto.operator.PartitionFunction in project presto by prestodb.

the class LocalExecutionPlanner method createOutputPartitioning.

private Optional<OutputPartitioning> createOutputPartitioning(TaskContext taskContext, PartitioningScheme partitioningScheme) {
    if (partitioningScheme.getPartitioning().getHandle().equals(FIXED_BROADCAST_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(FIXED_ARBITRARY_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(SCALED_WRITER_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(SINGLE_DISTRIBUTION) || partitioningScheme.getPartitioning().getHandle().equals(COORDINATOR_DISTRIBUTION)) {
        return Optional.empty();
    }
    List<VariableReferenceExpression> outputLayout = partitioningScheme.getOutputLayout();
    // We can convert the variables directly into channels, because the root must be a sink and therefore the layout is fixed
    List<Integer> partitionChannels;
    List<Optional<ConstantExpression>> partitionConstants;
    List<Type> partitionChannelTypes;
    if (partitioningScheme.getHashColumn().isPresent()) {
        partitionChannels = ImmutableList.of(outputLayout.indexOf(partitioningScheme.getHashColumn().get()));
        partitionConstants = ImmutableList.of(Optional.empty());
        partitionChannelTypes = ImmutableList.of(BIGINT);
    } else {
        checkArgument(partitioningScheme.getPartitioning().getArguments().stream().allMatch(argument -> argument instanceof ConstantExpression || argument instanceof VariableReferenceExpression), format("Expect all partitioning arguments to be either ConstantExpression or VariableReferenceExpression, but get %s", partitioningScheme.getPartitioning().getArguments()));
        partitionChannels = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument instanceof ConstantExpression) {
                return -1;
            }
            return outputLayout.indexOf(argument);
        }).collect(toImmutableList());
        partitionConstants = partitioningScheme.getPartitioning().getArguments().stream().map(argument -> {
            if (argument instanceof ConstantExpression) {
                return Optional.of((ConstantExpression) argument);
            }
            return Optional.<ConstantExpression>empty();
        }).collect(toImmutableList());
        partitionChannelTypes = partitioningScheme.getPartitioning().getArguments().stream().map(RowExpression::getType).collect(toImmutableList());
    }
    PartitionFunction partitionFunction = nodePartitioningManager.getPartitionFunction(taskContext.getSession(), partitioningScheme, partitionChannelTypes);
    OptionalInt nullChannel = OptionalInt.empty();
    Set<VariableReferenceExpression> partitioningColumns = partitioningScheme.getPartitioning().getVariableReferences();
    // partitioningColumns expected to have one column in the normal case, and zero columns when partitioning on a constant
    checkArgument(!partitioningScheme.isReplicateNullsAndAny() || partitioningColumns.size() <= 1);
    if (partitioningScheme.isReplicateNullsAndAny() && partitioningColumns.size() == 1) {
        nullChannel = OptionalInt.of(outputLayout.indexOf(getOnlyElement(partitioningColumns)));
    }
    return Optional.of(new OutputPartitioning(partitionFunction, partitionChannels, partitionConstants, partitioningScheme.isReplicateNullsAndAny(), nullChannel));
}