Examples with ConnectorBucketNodeMap com.facebook.presto.spi.connector.ConnectorBucketNodeMap used on opensource projects

Example 1 with ConnectorBucketNodeMap

use of com.facebook.presto.spi.connector.ConnectorBucketNodeMap in project presto by prestodb.

the class NodePartitioningManager method getConnectorBucketNodeMap.

private ConnectorBucketNodeMap getConnectorBucketNodeMap(Session session, PartitioningHandle partitioningHandle) {
    checkArgument(!(partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle));
    ConnectorId connectorId = partitioningHandle.getConnectorId().orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
    List<Node> nodes = getNodes(session, connectorId);
    ConnectorNodePartitioningProvider partitioningProvider = partitioningProviderManager.getPartitioningProvider(partitioningHandle.getConnectorId().get());
    ConnectorBucketNodeMap connectorBucketNodeMap = partitioningProvider.getBucketNodeMap(partitioningHandle.getTransactionHandle().orElse(null), session.toConnectorSession(partitioningHandle.getConnectorId().get()), partitioningHandle.getConnectorHandle(), nodes);
    checkArgument(connectorBucketNodeMap != null, "No partition map %s", partitioningHandle);
    return connectorBucketNodeMap;
}

Example 2 with ConnectorBucketNodeMap

use of com.facebook.presto.spi.connector.ConnectorBucketNodeMap in project presto by prestodb.

the class NodePartitioningManager method getBucketNodeMap.

public BucketNodeMap getBucketNodeMap(Session session, PartitioningHandle partitioningHandle, boolean preferDynamic) {
    ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle);
    NodeSelectionStrategy nodeSelectionStrategy = connectorBucketNodeMap.getNodeSelectionStrategy();
    switch(nodeSelectionStrategy) {
        case HARD_AFFINITY:
            return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap), false);
        case SOFT_AFFINITY:
            if (preferDynamic) {
                return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), connectorBucketNodeMap.getBucketCount(), getFixedMapping(connectorBucketNodeMap));
            }
            return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), getFixedMapping(connectorBucketNodeMap), true);
        case NO_PREFERENCE:
            if (preferDynamic) {
                return new DynamicBucketNodeMap(getSplitToBucket(session, partitioningHandle), connectorBucketNodeMap.getBucketCount());
            }
            return new FixedBucketNodeMap(getSplitToBucket(session, partitioningHandle), createArbitraryBucketToNode(nodeScheduler.createNodeSelector(session, partitioningHandle.getConnectorId().get()).selectRandomNodes(getMaxTasksPerStage(session)), connectorBucketNodeMap.getBucketCount()), false);
        default:
            throw new PrestoException(NODE_SELECTION_NOT_SUPPORTED, format("Unsupported node selection strategy %s", nodeSelectionStrategy));
    }
}

Example 3 with ConnectorBucketNodeMap

use of com.facebook.presto.spi.connector.ConnectorBucketNodeMap 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 4 with ConnectorBucketNodeMap

use of com.facebook.presto.spi.connector.ConnectorBucketNodeMap in project presto by prestodb.

the class NodePartitioningManager method getNodePartitioningMap.

public NodePartitionMap getNodePartitioningMap(Session session, PartitioningHandle partitioningHandle) {
    requireNonNull(session, "session is null");
    requireNonNull(partitioningHandle, "partitioningHandle is null");
    if (partitioningHandle.getConnectorHandle() instanceof SystemPartitioningHandle) {
        return ((SystemPartitioningHandle) partitioningHandle.getConnectorHandle()).getNodePartitionMap(session, nodeScheduler);
    }
    ConnectorId connectorId = partitioningHandle.getConnectorId().orElseThrow(() -> new IllegalArgumentException("No connector ID for partitioning handle: " + partitioningHandle));
    ConnectorBucketNodeMap connectorBucketNodeMap = getConnectorBucketNodeMap(session, partitioningHandle);
    // safety check for crazy partitioning
    checkArgument(connectorBucketNodeMap.getBucketCount() < 1_000_000, "Too many buckets in partitioning: %s", connectorBucketNodeMap.getBucketCount());
    List<InternalNode> bucketToNode;
    NodeSelectionStrategy nodeSelectionStrategy = connectorBucketNodeMap.getNodeSelectionStrategy();
    boolean cacheable;
    switch(nodeSelectionStrategy) {
        case HARD_AFFINITY:
            bucketToNode = getFixedMapping(connectorBucketNodeMap);
            cacheable = false;
            break;
        case SOFT_AFFINITY:
            bucketToNode = getFixedMapping(connectorBucketNodeMap);
            cacheable = true;
            break;
        case NO_PREFERENCE:
            bucketToNode = createArbitraryBucketToNode(nodeScheduler.createNodeSelector(session, connectorId).selectRandomNodes(getMaxTasksPerStage(session)), connectorBucketNodeMap.getBucketCount());
            cacheable = false;
            break;
        default:
            throw new PrestoException(NODE_SELECTION_NOT_SUPPORTED, format("Unsupported node selection strategy %s", nodeSelectionStrategy));
    }
    int[] bucketToPartition = new int[connectorBucketNodeMap.getBucketCount()];
    BiMap<InternalNode, Integer> nodeToPartition = HashBiMap.create();
    int nextPartitionId = 0;
    for (int bucket = 0; bucket < bucketToNode.size(); bucket++) {
        InternalNode node = bucketToNode.get(bucket);
        Integer partitionId = nodeToPartition.get(node);
        if (partitionId == null) {
            partitionId = nextPartitionId++;
            nodeToPartition.put(node, partitionId);
        }
        bucketToPartition[bucket] = partitionId;
    }
    List<InternalNode> partitionToNode = IntStream.range(0, nodeToPartition.size()).mapToObj(partitionId -> nodeToPartition.inverse().get(partitionId)).collect(toImmutableList());
    return new NodePartitionMap(partitionToNode, bucketToPartition, getSplitToBucket(session, partitioningHandle), cacheable);
}