From c40b4c136b956f4aead3eacfbb8e95fc853844d5 Mon Sep 17 00:00:00 2001
From: Sergey Pershin <spershin@hotmail.com>
Date: Mon, 22 Sep 2025 15:18:14 -0700
Subject: [PATCH] Use NodeAssignmentStats::getQueuedSplitsWeightForStage in
 task-based split scheduling

---
 .../scheduler/nodeSelection/SimpleNodeSelector.java         | 6 ++----
 .../com/facebook/presto/execution/TestNodeScheduler.java    | 1 -
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/nodeSelection/SimpleNodeSelector.java b/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/nodeSelection/SimpleNodeSelector.java
index 3f9526b064d98..3345ce52c5be0 100644
--- a/presto-main-base/src/main/java/com/facebook/presto/execution/scheduler/nodeSelection/SimpleNodeSelector.java
+++ b/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);
         };
     }
 
diff --git a/presto-main-base/src/test/java/com/facebook/presto/execution/TestNodeScheduler.java b/presto-main-base/src/test/java/com/facebook/presto/execution/TestNodeScheduler.java
index 0925f42733a77..281997a3246b9 100644
--- a/presto-main-base/src/test/java/com/facebook/presto/execution/TestNodeScheduler.java
+++ b/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();