How to use

getStage

method

in

org.qcri.rheem.core.plan.executionplan.ExecutionTask

public void markAsTerminalTask(ExecutionTask executionTask) {
  Validate.isTrue(executionTask.getStage() == this);
  this.terminalTasks.add(executionTask);
}

public void markAsStartTask(ExecutionTask executionTask) {
  Validate.isTrue(executionTask.getStage() == this);
  this.startTasks.add(executionTask);
}

/**
 * Tests for inter-stage instances.
 *
 * @return whether this instance connects {@link ExecutionTask}s of different {@link ExecutionStage}s.
 */
public boolean isBetweenStages() {
  if (this.producer == null || this.consumers.isEmpty()) {
    return false;
  }
  final ExecutionStage producerStage = this.producer.getStage();
  if (producerStage == null) {
    return false;
  }
  for (ExecutionTask consumer : this.consumers) {
    if (consumer.getStage() != null && !producerStage.equals(consumer.getStage())) {
      return true;
    }
  }
  return false;
}

/**
 * Inspect whether {@link ExecutionTask} is the {@link LoopHeadOperator} of the {@link #loopSubplan}. If so,
 * promote its {@link ExecutionStage} as the loop head.
 *
 * @param task to be checked
 */
public void update(ExecutionTask task) {
  if (this.headStageCache == null && this.isLoopHead(task)) {
    this.headStageCache = task.getStage();
  }
}

/**
 * Scrap any consumer {@link ExecutionTask}s and sibling {@link Channel}s that are not within the given
 * {@link ExecutionStage}s.
 *
 * @return whether consumer {@link ExecutionTask}s have been removed
 */
public boolean retain(Set<ExecutionStage> retainableStages) {
  boolean isConsumersRemoved = this.consumers.removeIf(consumer -> !retainableStages.contains(consumer.getStage()));
  this.removeSiblingsWhere((sibling) -> !retainableStages.contains(sibling.getProducer().getStage()));
  return isConsumersRemoved;
}

/**
 * Collects all {@link ExecutionTask}s of this instance.
 */
public Set<ExecutionTask> getAllTasks() {
  final Queue<ExecutionTask> nextTasks = new LinkedList<>(this.startTasks);
  final Set<ExecutionTask> allTasks = new HashSet<>();
  while (!nextTasks.isEmpty()) {
    final ExecutionTask task = nextTasks.poll();
    assert task.getStage() == this;
    if (allTasks.add(task) && !this.terminalTasks.contains(task)) {
      Arrays.stream(task.getOutputChannels())
          .flatMap(channel -> channel.getConsumers().stream())
          .filter(consumer -> consumer.getStage() == this)
          .forEach(nextTasks::add);
    }
  }
  return allTasks;
}

/**
 * @return all {@link Channel}s of this instance that connect from other {@link ExecutionStage}s
 */
public Collection<Channel> getInboundChannels() {
  return this.getAllTasks().stream()
      .flatMap(task ->
          Arrays.stream(task.getInputChannels()).filter(
              channel -> channel.getProducer().getStage() != this
          )
      ).collect(Collectors.toList());
}

/**
 * Checks whether the given {@link Channel} is inside of a {@link ExecutionStageLoop}.
 *
 * @param channel the said {@link Channel}
 * @return whether the {@link Channel} is in a {@link ExecutionStageLoop}
 */
private static boolean checkIfIsInLoopChannel(Channel channel) {
  final ExecutionStageLoop producerLoop = channel.getProducer().getStage().getLoop();
  return producerLoop != null && channel.getConsumers().stream().anyMatch(
      consumer -> consumer.getStage().getLoop() == producerLoop
  );
}

@Override
public void execute(final ExecutionStage stage, OptimizationContext optimizationContext, ExecutionState executionState) {
  Queue<ExecutionTask> scheduledTasks = new LinkedList<>(stage.getStartTasks());
  Set<ExecutionTask> executedTasks = new HashSet<>();
  while (!scheduledTasks.isEmpty()) {
    final ExecutionTask task = scheduledTasks.poll();
    if (executedTasks.contains(task)) continue;
    this.execute(task, optimizationContext, executionState);
    executedTasks.add(task);
    Arrays.stream(task.getOutputChannels())
        .flatMap(channel -> channel.getConsumers().stream())
        .filter(consumer -> consumer.getStage() == stage)
        .forEach(scheduledTasks::add);
  }
}

/**
 * Implements various sanity checks. Problems are logged.
 */
public boolean isSane() {
  // 1. Check if every ExecutionTask is assigned an ExecutionStage.
  final Set<ExecutionTask> allTasks = this.collectAllTasks();
  boolean isAllTasksAssigned = allTasks.stream().allMatch(task -> task.getStage() != null);
  if (!isAllTasksAssigned) {
    this.logger.error("There are tasks without stages.");
  }
  final Set<Channel> allChannels = allTasks.stream()
      .flatMap(task -> Stream.concat(Arrays.stream(task.getInputChannels()), Arrays.stream(task.getOutputChannels())))
      .collect(Collectors.toSet());
  boolean isAllChannelsOriginal = allChannels.stream()
      .allMatch(channel -> !channel.isCopy());
  if (!isAllChannelsOriginal) {
    this.logger.error("There are channels that are copies.");
  }
  boolean isAllSiblingsConsistent = true;
  for (Channel channel : allChannels) {
    for (Channel sibling : channel.getSiblings()) {
      if (!allChannels.contains(sibling)) {
        this.logger.error("A sibling of {}, namely {}, seems to be invalid.", channel, sibling);
        isAllSiblingsConsistent = false;
      }
    }
  }
  return isAllTasksAssigned && isAllChannelsOriginal && isAllSiblingsConsistent;
}

private void toExtensiveStringAux(ExecutionTask task, Set<ExecutionTask> seenTasks, StringBuilder sb, String indent) {
  if (!seenTasks.add(task)) {
    return;
  }
  for (Channel channel : task.getOutputChannels()) {
    for (ExecutionTask consumer : channel.getConsumers()) {
      if (consumer.getStage() == this) {
        sb.append(indent)
            .append("    ")
            .append(this.prettyPrint(task))
            .append(" => ")
            .append(this.prettyPrint(channel))
            .append(" => ")
            .append(this.prettyPrint(consumer)).append('\n');
        this.toExtensiveStringAux(consumer, seenTasks, sb, indent);
      } else {
        sb.append(indent)
            .append("Out ")
            .append(this.prettyPrint(task))
            .append(" => ")
            .append(this.prettyPrint(channel)).append('\n');
      }
    }
  }
}

/**
 * Starts building an {@link InterimStage} starting from the given {@link ExecutionTask} unless there is one already.
 * If a {@link PlatformExecution} is provided, the {@link InterimStage} will be associated with it.
 */
private void createStageFor(ExecutionTask task, PlatformExecution platformExecution) {
  assert task.getStage() == null : String.format("%s has already stage %s.", task, task.getStage());
  // See if there is already an InterimStage.
  if (this.assignedInterimStages.containsKey(task)) {
    return;
  }
  // Create a new PlatformExecution if none.
  if (platformExecution == null) {
    Platform platform = task.getOperator().getPlatform();
    platformExecution = new PlatformExecution(platform);
  }
  // Create the InterimStage and expand it.
  InterimStage initialStage = new InterimStageImpl(platformExecution);
  this.addStage(initialStage);
  this.assignTaskAndExpand(task, initialStage);
}

/**
 * Determine the consuming {@link InputSlot}s of the given {@link Channel} that lie within a {@link RheemPlan} and
 * have not been executed yet.
 * We follow non-RheemPlan {@link ExecutionOperator}s because they should merely forward data.
 */
private Collection<InputSlot<?>> findRheemPlanInputSlotFor(Channel channel, Set<ExecutionStage> executedStages) {
  Collection<InputSlot<?>> result = new LinkedList<>();
  for (ExecutionTask consumerTask : channel.getConsumers()) {
    if (executedStages.contains(consumerTask.getStage())) continue;
    if (!consumerTask.getOperator().isAuxiliary()) {
      result.add(consumerTask.getInputSlotFor(channel));
    } else {
      for (Channel consumerOutputChannel : consumerTask.getOutputChannels()) {
        result.addAll(this.findRheemPlanInputSlotFor(consumerOutputChannel, executedStages));
      }
    }
  }
  return result;
}

private void activateSuccessorTasks(ExecutionTask task, Collection<ChannelInstance> outputChannelInstances) {
  for (ChannelInstance outputChannelInstance : outputChannelInstances) {
    if (outputChannelInstance == null) continue; // Partial results possible (cf. LoopHeadOperator).
    final Channel channel = outputChannelInstance.getChannel();
    for (ExecutionTask consumer : channel.getConsumers()) {
      // Stay within ExecutionStage.
      if (consumer.getStage() != task.getStage() || consumer.getOperator().isLoopHead()) continue;
      // Get or create the TaskActivator.
      final TaskActivator consumerActivator = this.stagedActivators.computeIfAbsent(
          consumer, (key) -> new TaskActivator(key, this.fetchOperatorContext(key), this.executionState)
      );
      // Register the outputChannelInstance.
      consumerActivator.accept(outputChannelInstance);
      // Schedule the consumerActivator if it isReady.
      if (consumerActivator.isReady()) {
        this.stagedActivators.remove(consumer);
        this.readyActivators.add(consumerActivator);
      }
    }
  }
}

/**
 * The given instance should build upon the open {@link Channel}s of this instance. Then, this instance will be
 * expanded with the content of the given instance.
 *
 * @param expansion extends this instance, but they are not overlapping
 */
public void expand(ExecutionPlan expansion) {
  for (Channel openChannel : expansion.getOpenInputChannels()) {
    openChannel.mergeIntoOriginal();
    final Channel original = openChannel.getOriginal();
    final ExecutionStage producerStage = original.getProducer().getStage();
    assert producerStage != null : String.format("No stage found for %s.", original.getProducer());
    for (ExecutionTask consumer : original.getConsumers()) {
      final ExecutionStage consumerStage = consumer.getStage();
      assert consumerStage != null : String.format("No stage found for %s.", consumer);
      // Equal stages possible on "partially open" Channels.
      if (producerStage != consumerStage) {
        producerStage.addSuccessor(consumerStage);
      }
    }
  }
}

jsonThatOp.put("via", prettyPrint(channel));
perOutputThatList.add(jsonThatOp);
if (consumer.getStage() == this)
  this.toJsonMapAux(consumer, seenTasks, operators);

this.createStageFor(task, producer.getStage().getPlatformExecution());

/**
 * Checks if the given {@link Channel} (inbound to {@link #stage}), is entering an {@link ExecutionStageLoop}
 * when serving the {@link #stage}.
 *
 * @param inboundChannel {@link Channel} that is inbound to {@link #stage} and that might be a {@link ExecutionStageLoop} input
 * @return whether the {@link Channel} is a {@link ExecutionStageLoop} input, i.e., it is not produced in an {@link ExecutionStageLoop}
 * while this {@link #stage} is part of a {@link ExecutionStageLoop}
 */
private boolean checkIfIsLoopInput(Channel inboundChannel) {
  // NB: This code assumes no nested loops.
  return this.stage.getLoop() != null && this.stage.getLoop() != inboundChannel.getProducer().getStage().getLoop();
}

for (ExecutionTask consumer : outboundChannel.getConsumers()) {
  if (this.getChannelInstance(outboundChannel, consumer.isFeedbackInput(outboundChannel)) != null) {
    final ExecutionStage consumerStage = consumer.getStage();

/**
 * Determine the {@link ExecutionStageLoop} the given {@link Channel} belongs to.
 *
 * @param channel the {@link Channel}
 * @return the {@link ExecutionStageLoop} or {@code null} if none
 */
private static ExecutionStageLoop getExecutionStageLoop(Channel channel) {
  final ExecutionStage producerStage = channel.getProducer().getStage();
  if (producerStage.getLoop() == null) return null;
  final OutputSlot<?> output = channel.getProducer().getOutputSlotFor(channel);
  if (output != null
      && output.getOwner().isLoopHead()
      && ((LoopHeadOperator) output.getOwner()).getFinalLoopOutputs().contains(output)) {
    return null;
  }
  return producerStage.getLoop();
}