/** * Clean up {@link ExecutionStage}-related state and re-create {@link StageActivator}s etc. from the {@link ExecutionPlan}/ * * @param executionPlan whose {@link ExecutionStage}s will be executed * @param optimizationContext contains additional optimization info for the {@code executionPlan} */ public void prepare(ExecutionPlan executionPlan, OptimizationContext optimizationContext) { this.allStages.clear(); this.activatedStageActivators.clear(); this.suspendedStages.clear(); // Remove obsolete StageActivators (after re-optimization). this.allStages.addAll(executionPlan.getStages()); new ArrayList<>(this.pendingStageActivators.keySet()).stream() .filter(stage -> !this.allStages.contains(stage)) .forEach(this.pendingStageActivators::remove); // Create StageActivators for all ExecutionStages. for (ExecutionStage stage : this.allStages) { // Avoid re-activating already executed ExecutionStages. if (this.completedStages.contains(stage)) continue; final StageActivator activator = this.getOrCreateActivator( stage, () -> this.determineInitialOptimizationContext(stage, optimizationContext) ); this.tryToActivate(activator); } }