Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#Trigger

@Test
public void testAfterWatermarkEarlyLateTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setAfterEndOfWindow(
              RunnerApi.Trigger.AfterEndOfWindow.newBuilder()
                  .setEarlyFirings(subtrigger1)
                  .setLateFirings(subtrigger2))
          .build();
  AfterWatermarkStateMachine.AfterWatermarkEarlyAndLate machine =
      (AfterWatermarkStateMachine.AfterWatermarkEarlyAndLate)
          TriggerStateMachines.stateMachineForTrigger(trigger);

  assertThat(
      machine,
      equalTo(
          AfterWatermarkStateMachine.pastEndOfWindow()
              .withEarlyFirings(submachine1)
              .withLateFirings(submachine2)));
}
 

private RunnerApi.Trigger convertSpecific(AfterWatermarkEarlyAndLate v) {
  RunnerApi.Trigger.AfterEndOfWindow.Builder builder =
      RunnerApi.Trigger.AfterEndOfWindow.newBuilder();

  builder.setEarlyFirings(toProto(v.getEarlyTrigger()));
  if (v.getLateTrigger() != null) {
    builder.setLateFirings(toProto(v.getLateTrigger()));
  }

  return RunnerApi.Trigger.newBuilder().setAfterEndOfWindow(builder).build();
}
 

@Test
public void testAfterAllTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setAfterAll(
              RunnerApi.Trigger.AfterAll.newBuilder()
                  .addSubtriggers(subtrigger1)
                  .addSubtriggers(subtrigger2))
          .build();
  AfterAllStateMachine machine =
      (AfterAllStateMachine) TriggerStateMachines.stateMachineForTrigger(trigger);

  assertThat(machine, equalTo(AfterAllStateMachine.of(submachine1, submachine2)));
}
 

@Test
public void testAfterEachTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setAfterEach(
              RunnerApi.Trigger.AfterEach.newBuilder()
                  .addSubtriggers(subtrigger1)
                  .addSubtriggers(subtrigger2))
          .build();
  AfterEachStateMachine machine =
      (AfterEachStateMachine) TriggerStateMachines.stateMachineForTrigger(trigger);

  assertThat(machine, equalTo(AfterEachStateMachine.inOrder(submachine1, submachine2)));
}
 

public static TriggerStateMachine stateMachineForTrigger(RunnerApi.Trigger trigger) {
  switch (trigger.getTriggerCase()) {
    case AFTER_ALL:
      return AfterAllStateMachine.of(
          stateMachinesForTriggers(trigger.getAfterAll().getSubtriggersList()));
    case AFTER_ANY:
      return AfterFirstStateMachine.of(
          stateMachinesForTriggers(trigger.getAfterAny().getSubtriggersList()));
    case AFTER_END_OF_WINDOW:
      return stateMachineForAfterEndOfWindow(trigger.getAfterEndOfWindow());
    case ELEMENT_COUNT:
      return AfterPaneStateMachine.elementCountAtLeast(
          trigger.getElementCount().getElementCount());
    case AFTER_SYNCHRONIZED_PROCESSING_TIME:
      return AfterSynchronizedProcessingTimeStateMachine.ofFirstElement();
    case DEFAULT:
      return DefaultTriggerStateMachine.of();
    case NEVER:
      return NeverStateMachine.ever();
    case ALWAYS:
      return ReshuffleTriggerStateMachine.create();
    case OR_FINALLY:
      return stateMachineForTrigger(trigger.getOrFinally().getMain())
          .orFinally(stateMachineForTrigger(trigger.getOrFinally().getFinally()));
    case REPEAT:
      return RepeatedlyStateMachine.forever(
          stateMachineForTrigger(trigger.getRepeat().getSubtrigger()));
    case AFTER_EACH:
      return AfterEachStateMachine.inOrder(
          stateMachinesForTriggers(trigger.getAfterEach().getSubtriggersList()));
    case AFTER_PROCESSING_TIME:
      return stateMachineForAfterProcessingTime(trigger.getAfterProcessingTime());
    case TRIGGER_NOT_SET:
      throw new IllegalArgumentException(
          String.format("Required field 'trigger' not set on %s", trigger));
    default:
      throw new IllegalArgumentException(String.format("Unknown trigger type %s", trigger));
  }
}
 

private static List<Trigger> protosToTriggers(List<RunnerApi.Trigger> triggers) {
  List<Trigger> result = Lists.newArrayList();
  for (RunnerApi.Trigger trigger : triggers) {
    result.add(fromProto(trigger));
  }
  return result;
}
 

private RunnerApi.Trigger convertSpecific(AfterProcessingTime v) {
  RunnerApi.Trigger.AfterProcessingTime.Builder builder =
      RunnerApi.Trigger.AfterProcessingTime.newBuilder();

  for (TimestampTransform transform : v.getTimestampTransforms()) {
    builder.addTimestampTransforms(convertTimestampTransform(transform));
  }

  return RunnerApi.Trigger.newBuilder().setAfterProcessingTime(builder).build();
}
 

@Test
public void testOrFinallyTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setOrFinally(
              RunnerApi.Trigger.OrFinally.newBuilder()
                  .setMain(subtrigger1)
                  .setFinally(subtrigger2))
          .build();
  OrFinallyStateMachine machine =
      (OrFinallyStateMachine) TriggerStateMachines.stateMachineForTrigger(trigger);

  assertThat(machine, equalTo(submachine1.orFinally(submachine2)));
}
 

@Override
public void processElement(WindowedValue<KeyedWorkItem<K, V>> element) throws Exception {
  KeyedWorkItem<K, V> workItem = element.getValue();
  K key = workItem.key();

  UncommittedBundle<KV<K, Iterable<V>>> bundle =
      evaluationContext.createKeyedBundle(
          structuralKey,
          (PCollection<KV<K, Iterable<V>>>)
              Iterables.getOnlyElement(application.getOutputs().values()));
  outputBundles.add(bundle);
  CopyOnAccessInMemoryStateInternals stateInternals = stepContext.stateInternals();
  DirectTimerInternals timerInternals = stepContext.timerInternals();
  RunnerApi.Trigger runnerApiTrigger =
      TriggerTranslation.toProto(windowingStrategy.getTrigger());
  ReduceFnRunner<K, V, Iterable<V>, BoundedWindow> reduceFnRunner =
      new ReduceFnRunner<>(
          key,
          windowingStrategy,
          ExecutableTriggerStateMachine.create(
              TriggerStateMachines.stateMachineForTrigger(runnerApiTrigger)),
          stateInternals,
          timerInternals,
          new OutputWindowedValueToBundle<>(bundle),
          new UnsupportedSideInputReader(DirectGroupAlsoByWindow.class.getSimpleName()),
          reduceFn,
          options);

  // Drop any elements within expired windows
  reduceFnRunner.processElements(
      dropExpiredWindows(key, workItem.elementsIterable(), timerInternals));
  reduceFnRunner.onTimers(workItem.timersIterable());
  reduceFnRunner.persist();
}
 

private RunnerApi.Trigger convertSpecific(AfterEach v) {
  RunnerApi.Trigger.AfterEach.Builder builder = RunnerApi.Trigger.AfterEach.newBuilder();

  for (Trigger subtrigger : v.subTriggers()) {
    builder.addSubtriggers(toProto(subtrigger));
  }

  return RunnerApi.Trigger.newBuilder().setAfterEach(builder).build();
}
 

@Test
public void testNeverTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setNever(RunnerApi.Trigger.Never.getDefaultInstance())
          .build();
  NeverStateMachine machine =
      (NeverStateMachine) checkNotNull(TriggerStateMachines.stateMachineForTrigger(trigger));
  // No parameters, so if it doesn't crash, we win!
}
 

private RunnerApi.Trigger convertSpecific(AfterAll v) {
  RunnerApi.Trigger.AfterAll.Builder builder = RunnerApi.Trigger.AfterAll.newBuilder();

  for (Trigger subtrigger : v.subTriggers()) {
    builder.addSubtriggers(toProto(subtrigger));
  }

  return RunnerApi.Trigger.newBuilder().setAfterAll(builder).build();
}
 

private RunnerApi.Trigger convertSpecific(AfterFirst v) {
  RunnerApi.Trigger.AfterAny.Builder builder = RunnerApi.Trigger.AfterAny.newBuilder();

  for (Trigger subtrigger : v.subTriggers()) {
    builder.addSubtriggers(toProto(subtrigger));
  }

  return RunnerApi.Trigger.newBuilder().setAfterAny(builder).build();
}
 

@Test
public void testAfterFirstTranslation() {
  RunnerApi.Trigger trigger =
      RunnerApi.Trigger.newBuilder()
          .setAfterAny(
              RunnerApi.Trigger.AfterAny.newBuilder()
                  .addSubtriggers(subtrigger1)
                  .addSubtriggers(subtrigger2))
          .build();
  AfterFirstStateMachine machine =
      (AfterFirstStateMachine) TriggerStateMachines.stateMachineForTrigger(trigger);

  assertThat(machine, equalTo(AfterFirstStateMachine.of(submachine1, submachine2)));
}
 

public static RunnerApi.Trigger toProto(Trigger trigger) {
  return CONVERTER.convertTrigger(trigger);
}
 

private RunnerApi.Trigger convertSpecific(AfterSynchronizedProcessingTime v) {
  return RunnerApi.Trigger.newBuilder()
      .setAfterSynchronizedProcessingTime(
          RunnerApi.Trigger.AfterSynchronizedProcessingTime.getDefaultInstance())
      .build();
}
 

private RunnerApi.Trigger convertSpecific(NeverTrigger v) {
  return RunnerApi.Trigger.newBuilder()
      .setNever(RunnerApi.Trigger.Never.getDefaultInstance())
      .build();
}
 

private RunnerApi.Trigger convertSpecific(FromEndOfWindow v) {
  return RunnerApi.Trigger.newBuilder()
      .setAfterEndOfWindow(RunnerApi.Trigger.AfterEndOfWindow.newBuilder())
      .build();
}
 

private RunnerApi.Trigger convertSpecific(DefaultTrigger v) {
  return RunnerApi.Trigger.newBuilder()
      .setDefault(RunnerApi.Trigger.Default.getDefaultInstance())
      .build();
}
 

public RunnerApi.Trigger convertTrigger(Trigger trigger) {
  Method evaluationMethod = getEvaluationMethod(trigger.getClass());
  return tryConvert(evaluationMethod, trigger);
}