APIFramework { private static AlgebricksAbsolutePartitionConstraint getComputationLocations(Map<String, NodeControllerInfo> ncMap, int parallelismHint) { int parallelism = parallelismHint <= 0 ? -2 * ncMap.size() : parallelismHint; int numNodes = ncMap.size(); int numNodesWithOneMorePartition = parallelism % numNodes; int perNodeParallelismMin = parallelism / numNodes; int perNodeParallelismMax = parallelism / numNodes + 1; List<String> allNodes = new ArrayList<>(); Set<String> selectedNodesWithOneMorePartition = new HashSet<>(); ncMap.forEach((key, value) -> allNodes.add(key)); Random random = new Random(); for (int index = numNodesWithOneMorePartition; index >= 1; --index) { int pick = random.nextInt(index); selectedNodesWithOneMorePartition.add(allNodes.get(pick)); Collections.swap(allNodes, pick, index - 1); } List<String> locations = new ArrayList<>(); ncMap.forEach((nodeId, value) -> { int availableCores = value.getNumAvailableCores(); int nodeParallelism = selectedNodesWithOneMorePartition.contains(nodeId) ? perNodeParallelismMax : perNodeParallelismMin; int coresToUse = nodeParallelism >= 0 && nodeParallelism < availableCores ? nodeParallelism : availableCores; for (int count = 0; count < coresToUse; ++count) { locations.add(nodeId); } }); return new AlgebricksAbsolutePartitionConstraint(locations.toArray(new String[0])); } APIFramework(ILangCompilationProvider compilationProvider); Pair<IReturningStatement, Integer> reWriteQuery(List<FunctionDecl> declaredFunctions,
            MetadataProvider metadataProvider, IReturningStatement q, SessionOutput output, boolean inlineUdfs,
            Collection<VarIdentifier> externalVars, IWarningCollector warningCollector); JobSpecification compileQuery(IClusterInfoCollector clusterInfoCollector, MetadataProvider metadataProvider,
            Query query, int varCounter, String outputDatasetName, SessionOutput output,
            ICompiledDmlStatement statement, Map<VarIdentifier, IAObject> externalVars, IResponsePrinter printer,
            IWarningCollector warningCollector); void executeJobArray(IHyracksClientConnection hcc, JobSpecification[] specs, PrintWriter out); void executeJobArray(IHyracksClientConnection hcc, Job[] jobs, PrintWriter out); ExecutionPlans getExecutionPlans(); static AlgebricksAbsolutePartitionConstraint getJobLocations(JobSpecification spec,
            INodeJobTracker jobTracker, AlgebricksAbsolutePartitionConstraint clusterLocations); static final String PREFIX_INTERNAL_PARAMETERS; }
@Test public void testGetComputationLocations() throws AlgebricksException { Map<String, NodeControllerInfo> map = new HashMap<>(); NodeControllerInfo nc1Info = mock(NodeControllerInfo.class); when(nc1Info.getNumAvailableCores()).thenReturn(4); NodeControllerInfo nc2Info = mock(NodeControllerInfo.class); when(nc2Info.getNumAvailableCores()).thenReturn(4); String nc1 = "nc1"; String nc2 = "nc2"; map.put(nc1, nc1Info); map.put(nc2, nc2Info); APIFramework apiFramework = new APIFramework(mock(ILangCompilationProvider.class)); AlgebricksAbsolutePartitionConstraint loc = (AlgebricksAbsolutePartitionConstraint) PA .invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, 5); int nc1Count = 0, nc2Count = 0; String[] partitions = loc.getLocations(); for (String partition : partitions) { if (partition.equals(nc1)) { ++nc1Count; } if (partition.equals(nc2)) { ++nc2Count; } } Assert.assertTrue(nc1Count > 0); Assert.assertTrue(nc2Count > 0); Assert.assertTrue(Math.abs(nc1Count - nc2Count) == 1); Assert.assertTrue(partitions.length == 5); loc = (AlgebricksAbsolutePartitionConstraint) PA.invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, 8); nc1Count = 0; nc2Count = 0; partitions = loc.getLocations(); for (String partition : partitions) { if (partition.equals(nc1)) { ++nc1Count; } if (partition.equals(nc2)) { ++nc2Count; } } Assert.assertTrue(nc1Count > 0); Assert.assertTrue(nc2Count > 0); Assert.assertTrue(Math.abs(nc1Count - nc2Count) == 0); Assert.assertTrue(partitions.length == 8); loc = (AlgebricksAbsolutePartitionConstraint) PA.invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, 1); Assert.assertTrue(loc.getLocations().length == 1); loc = (AlgebricksAbsolutePartitionConstraint) PA.invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, -100); Assert.assertTrue(loc.getLocations().length == 8); loc = (AlgebricksAbsolutePartitionConstraint) PA.invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, -1); Assert.assertTrue(loc.getLocations().length == 8); loc = (AlgebricksAbsolutePartitionConstraint) PA.invokeMethod(apiFramework, "getComputationLocations(java.util.Map,int)", map, 0); Assert.assertTrue(loc.getLocations().length == 8); }