cryptoe commented on code in PR #13353:
URL: https://github.com/apache/druid/pull/13353#discussion_r1034291614


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/kernel/controller/ControllerStageTracker.java:
##########
@@ -165,10 +188,88 @@ ClusterByPartitions getResultPartitionBoundaries()
     }
   }
 
+
+  /**
+   * Get workers which need to be sent partition boundaries
+   *
+   * @return
+   */
+  IntSet getWorkersToSendParitionBoundaries()
+  {
+    if (!getStageDefinition().doesShuffle()) {
+      throw new ISE("Result partition information is not relevant to this 
stage because it does not shuffle");
+    }
+    IntAVLTreeSet workers = new IntAVLTreeSet();
+    for (Integer worker : workerToPhase.keySet()) {
+      if 
(WorkerStagePhase.PRESHUFFLE_WAITING_FOR_RESULT_PARTITION_BOUNDARIES.equals(workerToPhase.get(worker)))
 {
+        workers.add(worker);
+      }
+    }
+    return workers;
+  }
+
+  /**
+   * Indicates that the work order for worker has been sent. Transistions the 
state to {@link WorkerStagePhase#READING_INPUT}
+   * if no more work orders need to be sent.
+   *
+   * @param worker
+   */
+  void workOrderSentForWorker(int worker)
+  {
+
+    workerToPhase.compute(worker, (wk, state) -> {
+      if (state == null) {
+        throw new ISE("Worker[%d] not found for stage[%s]", wk, 
stageDef.getStageNumber());
+      }
+      if (!WorkerStagePhase.READING_INPUT.canTransitionFrom(state)) {
+        throw new ISE(
+            "Worker[%d] cannot transistion from state[%s] to state[%s] while 
sending work order",
+            worker,
+            state,
+            WorkerStagePhase.READING_INPUT
+        );
+      }
+      return WorkerStagePhase.READING_INPUT;
+    });
+    if (phase != ControllerStagePhase.READING_INPUT) {
+      if (workOrdersNeedToBeSent()) {

Review Comment:
   Updated the method name and the java doc to `allWorkOrdersSent`
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to