Examples with Number2ObjectPair com.linkedin.pinot.common.utils.Pairs.Number2ObjectPair used on opensource projects

Example 1 with Number2ObjectPair

use of com.linkedin.pinot.common.utils.Pairs.Number2ObjectPair in project pinot by linkedin.

the class BalanceNumSegmentAssignmentStrategy method getAssignedInstances.

@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, String helixClusterName, SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
    String serverTenantName;
    String tableName;
    if ("realtime".equalsIgnoreCase(segmentMetadata.getIndexType())) {
        tableName = TableNameBuilder.REALTIME_TABLE_NAME_BUILDER.forTable(segmentMetadata.getTableName());
        serverTenantName = ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(tenantName);
    } else {
        tableName = TableNameBuilder.OFFLINE_TABLE_NAME_BUILDER.forTable(segmentMetadata.getTableName());
        serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);
    }
    List<String> selectedInstances = new ArrayList<String>();
    Map<String, Integer> currentNumSegmentsPerInstanceMap = new HashMap<String, Integer>();
    List<String> allTaggedInstances = HelixHelper.getEnabledInstancesWithTag(helixAdmin, helixClusterName, serverTenantName);
    for (String instance : allTaggedInstances) {
        currentNumSegmentsPerInstanceMap.put(instance, 0);
    }
    // Count number of segments assigned to each instance
    IdealState idealState = helixAdmin.getResourceIdealState(helixClusterName, tableName);
    if (idealState != null) {
        for (String partitionName : idealState.getPartitionSet()) {
            Map<String, String> instanceToStateMap = idealState.getInstanceStateMap(partitionName);
            if (instanceToStateMap != null) {
                for (String instanceName : instanceToStateMap.keySet()) {
                    if (currentNumSegmentsPerInstanceMap.containsKey(instanceName)) {
                        currentNumSegmentsPerInstanceMap.put(instanceName, currentNumSegmentsPerInstanceMap.get(instanceName) + 1);
                    }
                // else, ignore. Do not add servers, that are not tagged, to the map
                // By this approach, new segments will not be allotted to the server if tags changed
                }
            }
        }
    }
    // Select up to numReplicas instances with the fewest segments assigned
    PriorityQueue<Number2ObjectPair<String>> priorityQueue = new PriorityQueue<Number2ObjectPair<String>>(numReplicas, Pairs.getDescendingnumber2ObjectPairComparator());
    for (String key : currentNumSegmentsPerInstanceMap.keySet()) {
        priorityQueue.add(new Number2ObjectPair<String>(currentNumSegmentsPerInstanceMap.get(key), key));
        if (priorityQueue.size() > numReplicas) {
            priorityQueue.poll();
        }
    }
    while (!priorityQueue.isEmpty()) {
        selectedInstances.add(priorityQueue.poll().getB());
    }
    LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : " + segmentMetadata.getTableName() + ", selected instances: " + Arrays.toString(selectedInstances.toArray()));
    return selectedInstances;
}