Java Code Examples for org.apache.beam.sdk.transforms.windowing.WindowFn#windowCoder()

public StatefulDoFnRunner(
    DoFnRunner<InputT, OutputT> doFnRunner,
    Coder<InputT> inputCoder,
    StepContext stepContext,
    WindowingStrategy<?, ?> windowingStrategy,
    CleanupTimer<InputT> cleanupTimer,
    StateCleaner<W> stateCleaner,
    boolean requiresTimeSortedInput) {
  this.doFnRunner = doFnRunner;
  this.stepContext = stepContext;
  this.windowingStrategy = windowingStrategy;
  this.cleanupTimer = cleanupTimer;
  this.stateCleaner = stateCleaner;
  this.requiresTimeSortedInput = requiresTimeSortedInput;
  WindowFn<?, ?> windowFn = windowingStrategy.getWindowFn();
  @SuppressWarnings("unchecked")
  Coder<BoundedWindow> untypedCoder = (Coder<BoundedWindow>) windowFn.windowCoder();
  this.windowCoder = untypedCoder;

  this.sortBufferTag =
      StateTags.makeSystemTagInternal(
          StateTags.bag(SORT_BUFFER_STATE, WindowedValue.getFullCoder(inputCoder, windowCoder)));

  rejectMergingWindowFn(windowFn);
}
 

Iterable<WindowedValue<T>> getValues(PCollection<T> pcollection) {
  if (windowedValues == null) {
    WindowFn<?, ?> windowFn = pcollection.getWindowingStrategy().getWindowFn();
    Coder<? extends BoundedWindow> windowCoder = windowFn.windowCoder();
    final WindowedValue.WindowedValueCoder<T> windowedValueCoder;
    if (windowFn instanceof GlobalWindows) {
      windowedValueCoder = WindowedValue.ValueOnlyWindowedValueCoder.of(pcollection.getCoder());
    } else {
      windowedValueCoder =
          WindowedValue.FullWindowedValueCoder.of(pcollection.getCoder(), windowCoder);
    }
    JavaRDDLike<byte[], ?> bytesRDD = rdd.map(CoderHelpers.toByteFunction(windowedValueCoder));
    List<byte[]> clientBytes = bytesRDD.collect();
    windowedValues =
        clientBytes.stream()
            .map(bytes -> CoderHelpers.fromByteArray(bytes, windowedValueCoder))
            .collect(Collectors.toList());
  }
  return windowedValues;
}
 

public TimeInternalsCleanupTimer(
    TimerInternals timerInternals, WindowingStrategy<?, ?> windowingStrategy) {
  this.windowingStrategy = windowingStrategy;
  WindowFn<?, ?> windowFn = windowingStrategy.getWindowFn();
  windowCoder = (Coder<BoundedWindow>) windowFn.windowCoder();
  this.timerInternals = timerInternals;
}
 

public TriggerStateMachineContextFactory(
    WindowFn<?, W> windowFn, StateInternals stateInternals, ActiveWindowSet<W> activeWindows) {
  // Future triggers may be able to exploit the active window to state address window mapping.
  this.windowFn = windowFn;
  this.stateInternals = stateInternals;
  this.windowCoder = windowFn.windowCoder();
}
 

/**
 * Tests that a users call to set a timer gets properly dispatched to the timer internals. From
 * there on, it is the duty of the runner & step context to set it in whatever way is right for
 * that runner.
 */
@Test
public void testTimerSet() {
  WindowFn<?, ?> windowFn = new GlobalWindows();
  DoFnWithTimers<GlobalWindow> fn = new DoFnWithTimers(windowFn.windowCoder());
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  // Setting the timer needs the current time, as it is set relative
  Instant currentTime = new Instant(42);
  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(currentTime);

  runner.processElement(WindowedValue.valueInGlobalWindow("anyValue"));

  verify(mockTimerInternals)
      .setTimer(
          StateNamespaces.window(new GlobalWindows().windowCoder(), GlobalWindow.INSTANCE),
          TimerDeclaration.PREFIX + DoFnWithTimers.TIMER_ID,
          "",
          currentTime.plus(DoFnWithTimers.TIMER_OFFSET),
          currentTime.plus(DoFnWithTimers.TIMER_OFFSET),
          TimeDomain.EVENT_TIME);
}
 

private <SideWindowT extends BoundedWindow> Windmill.GlobalDataRequest buildGlobalDataRequest(
    PCollectionView<?> view, BoundedWindow mainWindow) {
  @SuppressWarnings("unchecked")
  WindowingStrategy<?, SideWindowT> sideWindowStrategy =
      (WindowingStrategy<?, SideWindowT>) view.getWindowingStrategyInternal();

  WindowFn<?, SideWindowT> sideWindowFn = sideWindowStrategy.getWindowFn();

  Coder<SideWindowT> sideInputWindowCoder = sideWindowFn.windowCoder();

  SideWindowT sideInputWindow =
      (SideWindowT) view.getWindowMappingFn().getSideInputWindow(mainWindow);

  ByteString.Output windowStream = ByteString.newOutput();
  try {
    sideInputWindowCoder.encode(sideInputWindow, windowStream, Coder.Context.OUTER);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }

  return Windmill.GlobalDataRequest.newBuilder()
      .setDataId(
          Windmill.GlobalDataId.newBuilder()
              .setTag(view.getTagInternal().getId())
              .setVersion(windowStream.toByteString())
              .build())
      .setExistenceWatermarkDeadline(
          WindmillTimeUtils.harnessToWindmillTimestamp(
              sideWindowStrategy.getTrigger().getWatermarkThatGuaranteesFiring(sideInputWindow)))
      .build();
}
 

/**
 * Tests that {@link SimpleDoFnRunner#onTimer} properly dispatches to the underlying {@link DoFn}.
 */
@Test
public void testOnTimerCalled() {
  WindowFn<?, GlobalWindow> windowFn = new GlobalWindows();
  DoFnWithTimers<GlobalWindow> fn = new DoFnWithTimers(windowFn.windowCoder());
  DoFnRunner<String, String> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          null,
          null,
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(windowFn),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  Instant currentTime = new Instant(42);
  Duration offset = Duration.millis(37);

  // Mocking is not easily compatible with annotation analysis, so we manually record
  // the method call.
  runner.onTimer(
      TimerDeclaration.PREFIX + DoFnWithTimers.TIMER_ID,
      "",
      null,
      GlobalWindow.INSTANCE,
      currentTime.plus(offset),
      currentTime.plus(offset),
      TimeDomain.EVENT_TIME);

  assertThat(
      fn.onTimerInvocations,
      contains(
          TimerData.of(
              DoFnWithTimers.TIMER_ID,
              "",
              StateNamespaces.window(windowFn.windowCoder(), GlobalWindow.INSTANCE),
              currentTime.plus(offset),
              currentTime.plus(offset),
              TimeDomain.EVENT_TIME)));
}