Use NodeAssignmentStats::getQueuedSplitsWeightForStage in task-based split scheduling

Author: spershin

presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/nodeSelection/SimpleNodeSelector.java

@@ -256,13 +256,11 @@ private ToLongFunction<InternalNode> createTaskLoadSplitWeightProvider(List<Remo
             RemoteTask remoteTask = tasksByNodeId.get(node.getNodeIdentifier());
             if (remoteTask == null) {
                 // No task for this node, return only the queued splits weight for the stage
-                return assignmentStats.getAssignedSplitsWeightForStage(node);
+                return assignmentStats.getQueuedSplitsWeightForStage(node);
             }
 
             TaskStatus taskStatus = remoteTask.getTaskStatus();
-            return taskStatus.getQueuedPartitionedSplitsWeight() +
-                    taskStatus.getRunningPartitionedSplitsWeight() +
-                    assignmentStats.getAssignedSplitsWeightForStage(node);
+            return taskStatus.getRunningPartitionedSplitsWeight() + assignmentStats.getQueuedSplitsWeightForStage(node);
         };
     }
 

presto-main-base/src/test/java/com/facebook/presto/execution/TestNodeScheduler.java

@@ -1324,7 +1324,6 @@ public void testScheduleSplitsBasedOnTaskLoad()
                 assertTrue(node.getNodeIdentifier().equals("other2") || node.getNodeIdentifier().equals("other3"));
             }
 
-            Map<RemoteTask, Multimap<PlanNodeId, Split>> splitsForTasks = new HashMap<>();
             PlanNodeId planNodeId = new PlanNodeId("sourceId");
             for (InternalNode node : assignments.keySet()) {
                 Multimap<PlanNodeId, Split> splits = ArrayListMultimap.create();