Consider stage balancing only in UniformNodeSelector

core/trino-main/src/main/java/io/trino/execution/scheduler/NodeSchedulerConfig.java

@@ -31,12 +31,18 @@ public enum NodeSchedulerPolicy
         UNIFORM, TOPOLOGY
     }
 
+    public enum SplitsBalancingPolicy
+    {
+        NODE, STAGE
+    }
+
     private int minCandidates = 10;
     private boolean includeCoordinator = true;
     private int maxSplitsPerNode = 100;
     private int maxPendingSplitsPerTask = 10;
     private NodeSchedulerPolicy nodeSchedulerPolicy = NodeSchedulerPolicy.UNIFORM;
     private boolean optimizedLocalScheduling = true;
+    private SplitsBalancingPolicy splitsBalancingPolicy = SplitsBalancingPolicy.STAGE;
     private int maxUnacknowledgedSplitsPerTask = 500;
 
     @NotNull
@@ -132,6 +138,20 @@ public NodeSchedulerConfig setMaxUnacknowledgedSplitsPerTask(int maxUnacknowledg
         return this;
     }
 
+    @NotNull
+    public SplitsBalancingPolicy getSplitsBalancingPolicy()
+    {
+        return splitsBalancingPolicy;
+    }
+
+    @Config("node-scheduler.splits-balancing-policy")
+    @ConfigDescription("Strategy for balancing new splits on worker nodes")
+    public NodeSchedulerConfig setSplitsBalancingPolicy(SplitsBalancingPolicy splitsBalancingPolicy)
+    {
+        this.splitsBalancingPolicy = splitsBalancingPolicy;
+        return this;
+    }
+
     public boolean getOptimizedLocalScheduling()
     {
         return optimizedLocalScheduling;

core/trino-main/src/main/java/io/trino/execution/scheduler/UniformNodeSelector.java

@@ -16,6 +16,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.HashMultimap;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.SetMultimap;
@@ -24,13 +25,16 @@
 import io.trino.execution.NodeTaskMap;
 import io.trino.execution.RemoteTask;
 import io.trino.execution.resourcegroups.IndexedPriorityQueue;
+import io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy;
 import io.trino.metadata.InternalNode;
 import io.trino.metadata.InternalNodeManager;
 import io.trino.metadata.Split;
 import io.trino.spi.HostAddress;
 import io.trino.spi.SplitWeight;
 import io.trino.spi.TrinoException;
 
+import javax.annotation.Nullable;
+
 import java.net.InetAddress;
 import java.net.UnknownHostException;
 import java.util.Collection;
@@ -68,6 +72,7 @@ public class UniformNodeSelector
     private final long maxSplitsWeightPerNode;
     private final long maxPendingSplitsWeightPerTask;
     private final int maxUnacknowledgedSplitsPerTask;
+    private final SplitsBalancingPolicy splitsBalancingPolicy;
     private final boolean optimizedLocalScheduling;
 
     public UniformNodeSelector(
@@ -79,6 +84,7 @@ public UniformNodeSelector(
             long maxSplitsWeightPerNode,
             long maxPendingSplitsWeightPerTask,
             int maxUnacknowledgedSplitsPerTask,
+            SplitsBalancingPolicy splitsBalancingPolicy,
             boolean optimizedLocalScheduling)
     {
         this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
@@ -90,6 +96,7 @@ public UniformNodeSelector(
         this.maxPendingSplitsWeightPerTask = maxPendingSplitsWeightPerTask;
         this.maxUnacknowledgedSplitsPerTask = maxUnacknowledgedSplitsPerTask;
         checkArgument(maxUnacknowledgedSplitsPerTask > 0, "maxUnacknowledgedSplitsPerTask must be > 0, found: %s", maxUnacknowledgedSplitsPerTask);
+        this.splitsBalancingPolicy = requireNonNull(splitsBalancingPolicy, "splitsBalancingPolicy is null");
         this.optimizedLocalScheduling = optimizedLocalScheduling;
     }
 
@@ -171,21 +178,12 @@ public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTas
                 throw new TrinoException(NO_NODES_AVAILABLE, "No nodes available to run query");
             }
 
-            InternalNode chosenNode = null;
-            long minWeight = Long.MAX_VALUE;
-
-            for (InternalNode node : candidateNodes) {
-                long totalSplitsWeight = assignmentStats.getTotalSplitsWeight(node);
-                if (totalSplitsWeight < minWeight && totalSplitsWeight < maxSplitsWeightPerNode && assignmentStats.getUnacknowledgedSplitCountForStage(node) < maxUnacknowledgedSplitsPerTask) {
-                    chosenNode = node;
-                    minWeight = totalSplitsWeight;
-                }
-            }
+            InternalNode chosenNode = chooseNodeForSplit(assignmentStats, candidateNodes);
             if (chosenNode == null) {
-                // minWeight is guaranteed to be MAX_VALUE at this line
+                long minWeight = Long.MAX_VALUE;
                 for (InternalNode node : candidateNodes) {
                     long queuedWeight = assignmentStats.getQueuedSplitsWeightForStage(node);
-                    if (queuedWeight < minWeight && queuedWeight < maxPendingSplitsWeightPerTask && assignmentStats.getUnacknowledgedSplitCountForStage(node) < maxUnacknowledgedSplitsPerTask) {
+                    if (queuedWeight <= minWeight && queuedWeight < maxPendingSplitsWeightPerTask && assignmentStats.getUnacknowledgedSplitCountForStage(node) < maxUnacknowledgedSplitsPerTask) {
                         chosenNode = node;
                         minWeight = queuedWeight;
                     }
@@ -226,6 +224,50 @@ public SplitPlacementResult computeAssignments(Set<Split> splits, List<RemoteTas
         return selectDistributionNodes(nodeMap.get().get(), nodeTaskMap, maxSplitsWeightPerNode, maxPendingSplitsWeightPerTask, maxUnacknowledgedSplitsPerTask, splits, existingTasks, bucketNodeMap);
     }
 
+    @Nullable
+    private InternalNode chooseNodeForSplit(NodeAssignmentStats assignmentStats, List<InternalNode> candidateNodes)
+    {
+        InternalNode chosenNode = null;
+        long minWeight = Long.MAX_VALUE;
+
+        List<InternalNode> freeNodes = getFreeNodesForStage(assignmentStats, candidateNodes);
+        switch (splitsBalancingPolicy) {
+            case STAGE:
+                for (InternalNode node : freeNodes) {
+                    long queuedWeight = assignmentStats.getQueuedSplitsWeightForStage(node);
+                    if (queuedWeight <= minWeight) {
+                        chosenNode = node;
+                        minWeight = queuedWeight;
+                    }
+                }
+                break;
+            case NODE:
+                for (InternalNode node : freeNodes) {
+                    long totalSplitsWeight = assignmentStats.getTotalSplitsWeight(node);
+                    if (totalSplitsWeight <= minWeight) {
+                        chosenNode = node;
+                        minWeight = totalSplitsWeight;
+                    }
+                }
+                break;
+            default:
+                throw new UnsupportedOperationException("Unsupported split balancing policy " + splitsBalancingPolicy);
+        }
+
+        return chosenNode;
+    }
+
+    private List<InternalNode> getFreeNodesForStage(NodeAssignmentStats assignmentStats, List<InternalNode> nodes)
+    {
+        ImmutableList.Builder<InternalNode> freeNodes = ImmutableList.builder();
+        for (InternalNode node : nodes) {
+            if (assignmentStats.getTotalSplitsWeight(node) < maxSplitsWeightPerNode && assignmentStats.getUnacknowledgedSplitCountForStage(node) < maxUnacknowledgedSplitsPerTask) {
+                freeNodes.add(node);
+            }
+        }
+        return freeNodes.build();
+    }
+
     /**
      * The method tries to make the distribution of splits more uniform. All nodes are arranged into a maxHeap and a minHeap
      * based on the number of splits that are assigned to them. Splits are redistributed, one at a time, from a maxNode to a

core/trino-main/src/main/java/io/trino/execution/scheduler/UniformNodeSelectorFactory.java

@@ -23,6 +23,7 @@
 import io.trino.Session;
 import io.trino.connector.CatalogName;
 import io.trino.execution.NodeTaskMap;
+import io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy;
 import io.trino.metadata.InternalNode;
 import io.trino.metadata.InternalNodeManager;
 import io.trino.spi.HostAddress;
@@ -59,6 +60,7 @@ public class UniformNodeSelectorFactory
     private final boolean includeCoordinator;
     private final long maxSplitsWeightPerNode;
     private final long maxPendingSplitsWeightPerTask;
+    private final SplitsBalancingPolicy splitsBalancingPolicy;
     private final boolean optimizedLocalScheduling;
     private final NodeTaskMap nodeTaskMap;
     private final Duration nodeMapMemoizationDuration;
@@ -86,6 +88,7 @@ public UniformNodeSelectorFactory(
         this.nodeManager = nodeManager;
         this.minCandidates = config.getMinCandidates();
         this.includeCoordinator = config.isIncludeCoordinator();
+        this.splitsBalancingPolicy = config.getSplitsBalancingPolicy();
         this.optimizedLocalScheduling = config.getOptimizedLocalScheduling();
         this.nodeTaskMap = requireNonNull(nodeTaskMap, "nodeTaskMap is null");
         int maxSplitsPerNode = config.getMaxSplitsPerNode();
@@ -122,6 +125,7 @@ public NodeSelector createNodeSelector(Session session, Optional<CatalogName> ca
                 maxSplitsWeightPerNode,
                 maxPendingSplitsWeightPerTask,
                 getMaxUnacknowledgedSplitsPerTask(session),
+                splitsBalancingPolicy,
                 optimizedLocalScheduling);
     }
 

core/trino-main/src/test/java/io/trino/execution/TestNodeSchedulerConfig.java

@@ -23,6 +23,7 @@
 import static io.airlift.configuration.testing.ConfigAssertions.assertRecordedDefaults;
 import static io.airlift.configuration.testing.ConfigAssertions.recordDefaults;
 import static io.trino.execution.scheduler.NodeSchedulerConfig.NodeSchedulerPolicy.UNIFORM;
+import static io.trino.execution.scheduler.NodeSchedulerConfig.SplitsBalancingPolicy.NODE;
 
 public class TestNodeSchedulerConfig
 {
@@ -36,6 +37,7 @@ public void testDefaults()
                 .setMaxPendingSplitsPerTask(10)
                 .setMaxUnacknowledgedSplitsPerTask(500)
                 .setIncludeCoordinator(true)
+                .setSplitsBalancingPolicy(NodeSchedulerConfig.SplitsBalancingPolicy.STAGE)
                 .setOptimizedLocalScheduling(true));
     }
 
@@ -49,6 +51,7 @@ public void testExplicitPropertyMappings()
                 .put("node-scheduler.max-pending-splits-per-task", "11")
                 .put("node-scheduler.max-splits-per-node", "101")
                 .put("node-scheduler.max-unacknowledged-splits-per-task", "501")
+                .put("node-scheduler.splits-balancing-policy", "node")
                 .put("node-scheduler.optimized-local-scheduling", "false")
                 .build();
 
@@ -59,6 +62,7 @@ public void testExplicitPropertyMappings()
                 .setMaxPendingSplitsPerTask(11)
                 .setMaxUnacknowledgedSplitsPerTask(501)
                 .setMinCandidates(11)
+                .setSplitsBalancingPolicy(NODE)
                 .setOptimizedLocalScheduling(false);
 
         assertFullMapping(properties, expected);