Java Code Examples for org.apache.beam.sdk.values.WindowingStrategy#of()

The following examples show how to use org.apache.beam.sdk.values.WindowingStrategy#of() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: GroupAlsoByWindowProperties.java    From beam with Apache License 2.0 6 votes vote down vote up
/**
 * Tests that for empty input and the given {@link WindowingStrategy}, the provided GABW
 * implementation produces no output.
 *
 * <p>The input type is deliberately left as a wildcard, since it is not relevant.
 */
public static <K, InputT, OutputT> void emptyInputEmptyOutput(
    GroupAlsoByWindowDoFnFactory<K, InputT, OutputT> gabwFactory) throws Exception {

  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));

  // This key should never actually be used, though it is eagerly passed to the
  // StateInternalsFactory so must be non-null
  @SuppressWarnings("unchecked")
  K fakeKey = (K) "this key should never be used";

  List<WindowedValue<KV<K, OutputT>>> result =
      runGABW(
          gabwFactory,
          windowingStrategy,
          fakeKey,
          Collections.<WindowedValue<InputT>>emptyList());

  assertThat(result, hasSize(0));
}
 
Example 2
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testStartBundleExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  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());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.startBundle();
}
 
Example 3
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testFinishBundleExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  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());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.finishBundle();
}
 
Example 4
Source File: StreamingKeyedWorkItemSideInputDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@SuppressWarnings("unchecked")
private StreamingKeyedWorkItemSideInputDoFnRunner<
        String, Integer, KV<String, Integer>, IntervalWindow>
    createRunner(DoFnRunners.OutputManager outputManager) throws Exception {
  CoderRegistry registry = CoderRegistry.createDefault();
  Coder<String> keyCoder = StringUtf8Coder.of();
  Coder<Integer> inputCoder = BigEndianIntegerCoder.of();

  AppliedCombineFn<String, Integer, ?, Integer> combineFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofIntegers(), registry, KvCoder.of(keyCoder, inputCoder));

  WindowingStrategy<Object, IntervalWindow> windowingStrategy = WindowingStrategy.of(WINDOW_FN);
  @SuppressWarnings("rawtypes")
  StreamingGroupAlsoByWindowViaWindowSetFn doFn =
      (StreamingGroupAlsoByWindowViaWindowSetFn)
          StreamingGroupAlsoByWindowsDoFns.create(
              windowingStrategy, key -> state, combineFn, keyCoder);

  DoFnRunner<KeyedWorkItem<String, Integer>, KV<String, Integer>> simpleDoFnRunner =
      new GroupAlsoByWindowFnRunner<>(
          PipelineOptionsFactory.create(),
          doFn.asDoFn(),
          mockSideInputReader,
          outputManager,
          mainOutputTag,
          stepContext);
  return new StreamingKeyedWorkItemSideInputDoFnRunner<
      String, Integer, KV<String, Integer>, IntervalWindow>(
      simpleDoFnRunner, keyCoder, sideInputFetcher, stepContext);
}
 
Example 5
Source File: BatchGroupAlsoByWindowFnsTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testCreateCombiningNonmerging() throws Exception {
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofLongs(),
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));

  assertThat(
      BatchGroupAlsoByWindowsDoFns.create(windowingStrategy, appliedFn),
      instanceOf(BatchGroupAlsoByWindowAndCombineFn.class));
}
 
Example 6
Source File: BatchGroupAlsoByWindowReshuffleDoFnTest.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Tests that for a simple sequence of elements on the same key, {@link
 * BatchGroupAlsoByWindowReshuffleFn} fires each element in a single pane.
 */
@Test
public void testReshuffleFiresEveryElement() throws Exception {
  GABWReshuffleDoFnFactory gabwFactory = new GABWReshuffleDoFnFactory();

  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));

  List<WindowedValue<KV<String, Iterable<String>>>> result =
      runGABW(
          gabwFactory,
          windowingStrategy,
          "key",
          WindowedValue.of(
              "v1", new Instant(1), Arrays.asList(window(0, 10)), PaneInfo.NO_FIRING),
          WindowedValue.of(
              "v2", new Instant(2), Arrays.asList(window(0, 10)), PaneInfo.NO_FIRING),
          WindowedValue.of(
              "v3", new Instant(13), Arrays.asList(window(10, 20)), PaneInfo.NO_FIRING));

  assertThat(result.size(), equalTo(3));

  WindowedValue<KV<String, Iterable<String>>> item0 = result.get(0);
  assertThat(item0.getValue().getValue(), contains("v1"));
  assertThat(item0.getTimestamp(), equalTo(new Instant(1)));
  assertThat(item0.getWindows(), contains(window(0, 10)));

  WindowedValue<KV<String, Iterable<String>>> item1 = result.get(1);
  assertThat(item1.getValue().getValue(), contains("v2"));
  assertThat(item1.getTimestamp(), equalTo(new Instant(2)));
  assertThat(item1.getWindows(), contains(window(0, 10)));

  WindowedValue<KV<String, Iterable<String>>> item2 = result.get(2);
  assertThat(item2.getValue().getValue(), contains("v3"));
  assertThat(item2.getTimestamp(), equalTo(new Instant(13)));
  assertThat(item2.getWindows(), contains(window(10, 20)));
}
 
Example 7
Source File: StreamingDataflowWorkerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
private ParallelInstruction makeDoFnInstruction(
    DoFn<?, ?> doFn, int producerIndex, Coder<?> outputCoder) {
  // The windows used in this test are actually fairly arbitrary and could not be assigned
  // by FixedWindow. However, the WindowingStrategy on the DoFnInfo is used only for
  // the window Coder, which is IntervalWindow.Coder
  WindowingStrategy<?, ?> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));

  return makeDoFnInstruction(doFn, producerIndex, outputCoder, windowingStrategy);
}
 
Example 8
Source File: StreamingSideInputFetcherTest.java    From beam with Apache License 2.0 5 votes vote down vote up
private <ReceiverT> StreamingSideInputFetcher<String, IntervalWindow> createFetcher(
    List<PCollectionView<String>> views) throws Exception {
  @SuppressWarnings({"unchecked", "rawtypes"})
  Iterable<PCollectionView<?>> typedViews = (Iterable) views;

  return new StreamingSideInputFetcher<String, IntervalWindow>(
      typedViews, StringUtf8Coder.of(), WindowingStrategy.of(WINDOW_FN), stepContext);
}
 
Example 9
Source File: BatchGroupAlsoByWindowFnsTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testCreateNoncombiningMerging() throws Exception {
  Coder<Long> inputCoder = VarLongCoder.of();
  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)));

  assertThat(
      BatchGroupAlsoByWindowsDoFns.createForIterable(
          windowingStrategy, new InMemoryStateInternalsFactory<>(), inputCoder),
      instanceOf(BatchGroupAlsoByWindowViaOutputBufferFn.class));
}
 
Example 10
Source File: WindowDoFnOperatorTest.java    From beam with Apache License 2.0 5 votes vote down vote up
private WindowDoFnOperator<Long, Long, Long> getWindowDoFnOperator() {
  WindowingStrategy<Object, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(standardMinutes(1)));

  TupleTag<KV<Long, Long>> outputTag = new TupleTag<>("main-output");

  SystemReduceFn<Long, Long, long[], Long, BoundedWindow> reduceFn =
      SystemReduceFn.combining(
          VarLongCoder.of(),
          AppliedCombineFn.withInputCoder(
              Sum.ofLongs(),
              CoderRegistry.createDefault(),
              KvCoder.of(VarLongCoder.of(), VarLongCoder.of())));

  Coder<IntervalWindow> windowCoder = windowingStrategy.getWindowFn().windowCoder();
  SingletonKeyedWorkItemCoder<Long, Long> workItemCoder =
      SingletonKeyedWorkItemCoder.of(VarLongCoder.of(), VarLongCoder.of(), windowCoder);
  FullWindowedValueCoder<SingletonKeyedWorkItem<Long, Long>> inputCoder =
      WindowedValue.getFullCoder(workItemCoder, windowCoder);
  FullWindowedValueCoder<KV<Long, Long>> outputCoder =
      WindowedValue.getFullCoder(KvCoder.of(VarLongCoder.of(), VarLongCoder.of()), windowCoder);

  return new WindowDoFnOperator<Long, Long, Long>(
      reduceFn,
      "stepName",
      (Coder) inputCoder,
      outputTag,
      emptyList(),
      new MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
      windowingStrategy,
      emptyMap(),
      emptyList(),
      PipelineOptionsFactory.as(FlinkPipelineOptions.class),
      VarLongCoder.of(),
      new WorkItemKeySelector(VarLongCoder.of()));
}
 
Example 11
Source File: GroupAlsoByWindowProperties.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Tests that the given {@link BatchGroupAlsoByWindowFn} implementation combines elements per
 * session window correctly according to the provided {@link CombineFn}.
 */
public static void combinesElementsPerSession(
    GroupAlsoByWindowDoFnFactory<String, Long, Long> gabwFactory,
    CombineFn<Long, ?, Long> combineFn)
    throws Exception {

  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)));

  List<WindowedValue<KV<String, Long>>> result =
      runGABW(
          gabwFactory,
          windowingStrategy,
          "k",
          WindowedValue.of(1L, new Instant(0), Arrays.asList(window(0, 10)), PaneInfo.NO_FIRING),
          WindowedValue.of(2L, new Instant(5), Arrays.asList(window(5, 15)), PaneInfo.NO_FIRING),
          WindowedValue.of(
              4L, new Instant(15), Arrays.asList(window(15, 25)), PaneInfo.NO_FIRING));

  assertThat(result, hasSize(2));

  TimestampedValue<KV<String, Long>> item0 = getOnlyElementInWindow(result, window(0, 15));
  assertThat(item0.getValue().getKey(), equalTo("k"));
  assertThat(item0.getValue().getValue(), equalTo(combineFn.apply(ImmutableList.of(1L, 2L))));
  assertThat(item0.getTimestamp(), equalTo(window(0, 15).maxTimestamp()));

  TimestampedValue<KV<String, Long>> item1 = getOnlyElementInWindow(result, window(15, 25));
  assertThat(item1.getValue().getKey(), equalTo("k"));
  assertThat(item1.getValue().getValue(), equalTo(combineFn.apply(ImmutableList.of(4L))));
  assertThat(item1.getTimestamp(), equalTo(window(15, 25).maxTimestamp()));
}
 
Example 12
Source File: LateDataDroppingDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testLateDataFilter() throws Exception {
  MetricsContainerImpl container = new MetricsContainerImpl("any");
  MetricsEnvironment.setCurrentContainer(container);
  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(15L));

  LateDataFilter lateDataFilter =
      new LateDataFilter(WindowingStrategy.of(WINDOW_FN), mockTimerInternals);

  Iterable<WindowedValue<Integer>> actual =
      lateDataFilter.filter(
          "a",
          ImmutableList.of(
              createDatum(13, 13L),
              createDatum(5, 5L), // late element, earlier than 4L.
              createDatum(16, 16L),
              createDatum(18, 18L)));

  Iterable<WindowedValue<Integer>> expected =
      ImmutableList.of(createDatum(13, 13L), createDatum(16, 16L), createDatum(18, 18L));
  assertThat(expected, containsInAnyOrder(Iterables.toArray(actual, WindowedValue.class)));
  long droppedValues =
      container
          .getCounter(
              MetricName.named(
                  LateDataDroppingDoFnRunner.class,
                  LateDataDroppingDoFnRunner.DROPPED_DUE_TO_LATENESS))
          .getCumulative();
  assertEquals(1, droppedValues);
  // Ensure that reiterating returns the same results and doesn't increment the counter again.
  assertThat(expected, containsInAnyOrder(Iterables.toArray(actual, WindowedValue.class)));
  droppedValues =
      container
          .getCounter(
              MetricName.named(
                  LateDataDroppingDoFnRunner.class,
                  LateDataDroppingDoFnRunner.DROPPED_DUE_TO_LATENESS))
          .getCumulative();
  assertEquals(1, droppedValues);
}
 
Example 13
Source File: BatchGroupAlsoByWindowFnsTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testCreateCombiningMerging() throws Exception {
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofLongs(),
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)));

  assertThat(
      BatchGroupAlsoByWindowsDoFns.create(windowingStrategy, appliedFn),
      instanceOf(BatchGroupAlsoByWindowAndCombineFn.class));
}
 
Example 14
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Demonstrates that attempting to output an element before the timestamp of the current element
 * plus the value of {@link DoFn#getAllowedTimestampSkew()} throws, but between that value and the
 * current timestamp succeeds.
 */
@Test
public void testSkew() {
  SkewingDoFn fn = new SkewingDoFn(Duration.standardMinutes(10L));
  DoFnRunner<Duration, Duration> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          new ListOutputManager(),
          new TupleTag<>(),
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  runner.startBundle();
  // Outputting between "now" and "now - allowed skew" succeeds.
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.standardMinutes(5L), new Instant(0)));
  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(IllegalArgumentException.class));
  thrown.expectMessage("must be no earlier");
  thrown.expectMessage(
      String.format("timestamp of the current input (%s)", new Instant(0).toString()));
  thrown.expectMessage(
      String.format(
          "the allowed skew (%s)",
          PeriodFormat.getDefault().print(Duration.standardMinutes(10L).toPeriod())));
  // Outputting before "now - allowed skew" fails.
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.standardHours(1L), new Instant(0)));
}
 
Example 15
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Demonstrates that attempting to output an element before the timestamp of the current element
 * with zero {@link DoFn#getAllowedTimestampSkew() allowed timestamp skew} throws.
 */
@Test
public void testBackwardsInTimeNoSkew() {
  SkewingDoFn fn = new SkewingDoFn(Duration.ZERO);
  DoFnRunner<Duration, Duration> runner =
      new SimpleDoFnRunner<>(
          null,
          fn,
          NullSideInputReader.empty(),
          new ListOutputManager(),
          new TupleTag<>(),
          Collections.emptyList(),
          mockStepContext,
          null,
          Collections.emptyMap(),
          WindowingStrategy.of(new GlobalWindows()),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  runner.startBundle();
  // An element output at the current timestamp is fine.
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.ZERO, new Instant(0)));
  thrown.expect(UserCodeException.class);
  thrown.expectCause(isA(IllegalArgumentException.class));
  thrown.expectMessage("must be no earlier");
  thrown.expectMessage(
      String.format("timestamp of the current input (%s)", new Instant(0).toString()));
  thrown.expectMessage(
      String.format(
          "the allowed skew (%s)", PeriodFormat.getDefault().print(Duration.ZERO.toPeriod())));
  // An element output before (current time - skew) is forbidden
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.millis(1L), new Instant(0)));
}
 
Example 16
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * 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);
}
 
Example 17
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testOnTimerExceptionsWrappedAsUserCodeException() {
  ThrowingDoFn fn = new ThrowingDoFn();
  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());

  thrown.expect(UserCodeException.class);
  thrown.expectCause(is(fn.exceptionToThrow));

  runner.onTimer(
      TimerDeclaration.PREFIX + ThrowingDoFn.TIMER_ID,
      "",
      null,
      GlobalWindow.INSTANCE,
      new Instant(0),
      new Instant(0),
      TimeDomain.EVENT_TIME);
}
 
Example 18
Source File: GroupAlsoByWindowProperties.java    From beam with Apache License 2.0 5 votes vote down vote up
/** Tests that the given GABW implementation correctly groups elements into merged sessions. */
public static void groupsElementsInMergedSessions(
    GroupAlsoByWindowDoFnFactory<String, String, Iterable<String>> gabwFactory) throws Exception {

  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)));

  List<WindowedValue<KV<String, Iterable<String>>>> result =
      runGABW(
          gabwFactory,
          windowingStrategy,
          "key",
          WindowedValue.of(
              "v1", new Instant(0), Arrays.asList(window(0, 10)), PaneInfo.NO_FIRING),
          WindowedValue.of(
              "v2", new Instant(5), Arrays.asList(window(5, 15)), PaneInfo.NO_FIRING),
          WindowedValue.of(
              "v3", new Instant(15), Arrays.asList(window(15, 25)), PaneInfo.NO_FIRING));

  assertThat(result, hasSize(2));

  TimestampedValue<KV<String, Iterable<String>>> item0 =
      getOnlyElementInWindow(result, window(0, 15));
  assertThat(item0.getValue().getValue(), containsInAnyOrder("v1", "v2"));
  assertThat(item0.getTimestamp(), equalTo(window(0, 15).maxTimestamp()));

  TimestampedValue<KV<String, Iterable<String>>> item1 =
      getOnlyElementInWindow(result, window(15, 25));
  assertThat(item1.getValue().getValue(), contains("v3"));
  assertThat(item1.getTimestamp(), equalTo(window(15, 25).maxTimestamp()));
}
 
Example 19
Source File: DoFnOperatorTest.java    From beam with Apache License 2.0 4 votes vote down vote up
@Test
public void testLateDroppingForStatefulFn() throws Exception {

  WindowingStrategy<Object, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(new Duration(10)));

  DoFn<Integer, String> fn =
      new DoFn<Integer, String>() {

        @StateId("state")
        private final StateSpec<ValueState<String>> stateSpec =
            StateSpecs.value(StringUtf8Coder.of());

        @ProcessElement
        public void processElement(ProcessContext context) {
          context.output(context.element().toString());
        }
      };

  VarIntCoder keyCoder = VarIntCoder.of();
  Coder<WindowedValue<Integer>> inputCoder =
      WindowedValue.getFullCoder(keyCoder, windowingStrategy.getWindowFn().windowCoder());
  Coder<WindowedValue<String>> outputCoder =
      WindowedValue.getFullCoder(
          StringUtf8Coder.of(), windowingStrategy.getWindowFn().windowCoder());

  KeySelector<WindowedValue<Integer>, ByteBuffer> keySelector =
      e -> FlinkKeyUtils.encodeKey(e.getValue(), keyCoder);

  TupleTag<String> outputTag = new TupleTag<>("main-output");

  DoFnOperator<Integer, String> doFnOperator =
      new DoFnOperator<>(
          fn,
          "stepName",
          inputCoder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
          windowingStrategy,
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          PipelineOptionsFactory.as(FlinkPipelineOptions.class),
          keyCoder, /* key coder */
          keySelector,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
      new KeyedOneInputStreamOperatorTestHarness<>(
          doFnOperator,
          keySelector,
          new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of()));

  testHarness.open();

  testHarness.processWatermark(0);

  IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10));

  // this should not be late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(WindowedValue.of("13", new Instant(0), window1, PaneInfo.NO_FIRING)));

  testHarness.getOutput().clear();

  testHarness.processWatermark(9);

  // this should still not be considered late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(WindowedValue.of("17", new Instant(0), window1, PaneInfo.NO_FIRING)));

  testHarness.getOutput().clear();

  testHarness.processWatermark(10);

  // this should now be considered late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());

  testHarness.close();
}
 
Example 20
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 4 votes vote down vote up
/**
 * 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)));
}