MSQ: Fix two issues with phase transitions. (#17053)

1) ControllerQueryKernel: Update readyToReadResults to acknowledge that sorting stages can
   go directly from READING_INPUT to RESULTS_READY.

2) WorkerStageKernel: Ignore RESULTS_COMPLETE if work is already finished, which can happen
   if the transition to FINISHED comes early due to a downstream LIMIT.

(cherry picked from commit 654e0b444b)
This commit is contained in:
Gian Merlino 2024-09-13 15:59:41 -07:00 committed by Abhishek Agarwal
parent 479fdea065
commit 75ac9051cd
2 changed files with 9 additions and 4 deletions

View File

@ -692,11 +692,10 @@ public class ControllerQueryKernel
{
if (stageOutputChannelModes.get(stageId) == OutputChannelMode.MEMORY) {
if (getStageDefinition(stageId).doesSortDuringShuffle()) {
// Stages that sort during shuffle go through a READING_INPUT phase followed by a POST_READING phase
// (once all input is read). These stages start producing output once POST_READING starts.
return newPhase == ControllerStagePhase.POST_READING;
// Sorting stages start producing output when they finish reading their input.
return newPhase.isDoneReadingInput();
} else {
// Can read results immediately.
// Non-sorting stages start producing output immediately.
return newPhase == ControllerStagePhase.NEW;
}
} else {

View File

@ -183,6 +183,12 @@ public class WorkerStageKernel
throw new NullPointerException("resultObject must not be null");
}
if (phase.isTerminal()) {
// Ignore RESULTS_COMPLETE if work is already finished. This can happen if we transition to FINISHED early
// due to a downstream stage including a limit.
return;
}
transitionTo(WorkerStagePhase.RESULTS_COMPLETE);
this.resultObject = resultObject;
}