Java Code Examples for org.apache.beam.sdk.transforms.DoFnSchemaInformation#create()

The following examples show how to use org.apache.beam.sdk.transforms.DoFnSchemaInformation#create() . 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: GroupByKeyAndWindowDoFnTransform.java    From incubator-nemo with Apache License 2.0 6 votes vote down vote up
/**
 * GroupByKey constructor.
 *
 * @param outputCoders      output coders
 * @param mainOutputTag     main output tag
 * @param windowingStrategy windowing strategy
 * @param options           pipeline options
 * @param reduceFn          reduce function
 * @param displayData       display data.
 */
public GroupByKeyAndWindowDoFnTransform(final Map<TupleTag<?>, Coder<?>> outputCoders,
                                        final TupleTag<KV<K, Iterable<InputT>>> mainOutputTag,
                                        final WindowingStrategy<?, ?> windowingStrategy,
                                        final PipelineOptions options,
                                        final SystemReduceFn reduceFn,
                                        final DisplayData displayData) {
  super(null, /* doFn */
    null, /* inputCoder */
    outputCoders,
    mainOutputTag,
    Collections.emptyList(),  /*  GBK does not have additional outputs */
    windowingStrategy,
    Collections.emptyMap(), /*  GBK does not have additional side inputs */
    options,
    displayData,
    DoFnSchemaInformation.create(),
    Collections.emptyMap());
  this.keyToValues = new HashMap<>();
  this.reduceFn = reduceFn;
  this.prevOutputWatermark = new Watermark(Long.MIN_VALUE);
  this.keyAndWatermarkHoldMap = new HashMap<>();
}
 
Example 2
Source File: FlinkPipelineOptionsTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test(expected = Exception.class)
public void parDoBaseClassPipelineOptionsNullTest() {
  TupleTag<String> mainTag = new TupleTag<>("main-output");
  Coder<WindowedValue<String>> coder = WindowedValue.getValueOnlyCoder(StringUtf8Coder.of());
  new DoFnOperator<>(
      new TestDoFn(),
      "stepName",
      coder,
      Collections.emptyMap(),
      mainTag,
      Collections.emptyList(),
      new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, coder),
      WindowingStrategy.globalDefault(),
      new HashMap<>(),
      Collections.emptyList(),
      null,
      null, /* key coder */
      null /* key selector */,
      DoFnSchemaInformation.create(),
      Collections.emptyMap());
}
 
Example 3
Source File: SimplePushbackSideInputDoFnRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
private DoFnRunner<KV<String, Integer>, Integer> getDoFnRunner(
    DoFn<KV<String, Integer>, Integer> fn) {
  return new SimpleDoFnRunner<>(
      null,
      fn,
      NullSideInputReader.empty(),
      null,
      null,
      Collections.emptyList(),
      mockStepContext,
      null,
      Collections.emptyMap(),
      WINDOWING_STRATEGY,
      DoFnSchemaInformation.create(),
      Collections.emptyMap());
}
 
Example 4
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@Test
public void testProcessElementExceptionsWrappedAsUserCodeException() {
  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.processElement(WindowedValue.valueInGlobalWindow("anyValue"));
}
 
Example 5
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 6
Source File: DoFnTransformTest.java    From incubator-nemo with Apache License 2.0 5 votes vote down vote up
@Test
@SuppressWarnings("unchecked")
public void testSingleOutput() {

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

  final DoFnTransform<String, String> doFnTransform =
    new DoFnTransform<>(
      new IdentityDoFn<>(),
      NULL_INPUT_CODER,
      NULL_OUTPUT_CODERS,
      outputTag,
      Collections.emptyList(),
      WindowingStrategy.globalDefault(),
      PipelineOptionsFactory.as(NemoPipelineOptions.class),
      DisplayData.none(),
      DoFnSchemaInformation.create(),
      Collections.emptyMap());

  final Transform.Context context = mock(Transform.Context.class);
  final OutputCollector<WindowedValue<String>> oc = new TestOutputCollector<>();
  doFnTransform.prepare(context, oc);

  doFnTransform.onData(WindowedValue.valueInGlobalWindow("Hello"));

  assertEquals(((TestOutputCollector<String>) oc).outputs.get(0), WindowedValue.valueInGlobalWindow("Hello"));

  doFnTransform.close();
}
 
Example 7
Source File: DoFnOperatorTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
@SuppressWarnings("unchecked")
public void testSingleOutput() throws Exception {

  Coder<WindowedValue<String>> coder = WindowedValue.getValueOnlyCoder(StringUtf8Coder.of());

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

  DoFnOperator<String, String> doFnOperator =
      new DoFnOperator<>(
          new IdentityDoFn<>(),
          "stepName",
          coder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, coder),
          WindowingStrategy.globalDefault(),
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          PipelineOptionsFactory.as(FlinkPipelineOptions.class),
          null,
          null,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
      new OneInputStreamOperatorTestHarness<>(doFnOperator);

  testHarness.open();

  testHarness.processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow("Hello")));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(WindowedValue.valueInGlobalWindow("Hello")));

  testHarness.close();
}
 
Example 8
Source File: SimpleDoFnRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * Demonstrates that attempting to output an element with a timestamp before the current one
 * always succeeds when {@link DoFn#getAllowedTimestampSkew()} is equal to {@link Long#MAX_VALUE}
 * milliseconds.
 */
@Test
public void testInfiniteSkew() {
  SkewingDoFn fn = new SkewingDoFn(Duration.millis(Long.MAX_VALUE));
  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();
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(Duration.millis(1L), new Instant(0)));
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(
          Duration.millis(1L), BoundedWindow.TIMESTAMP_MIN_VALUE.plus(Duration.millis(1))));
  runner.processElement(
      WindowedValue.timestampedValueInGlobalWindow(
          // This is the maximum amount a timestamp in beam can move (from the maximum timestamp
          // to the minimum timestamp).
          Duration.millis(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis())
              .minus(Duration.millis(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis())),
          BoundedWindow.TIMESTAMP_MAX_VALUE));
}
 
Example 9
Source File: DoFnOperatorTest.java    From beam with Apache License 2.0 5 votes vote down vote up
private static DoFnOperator getOperatorForCleanupInspection() {
  FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
  options.setParallelism(4);

  TupleTag<String> outputTag = new TupleTag<>("main-output");
  WindowedValue.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
      WindowedValue.getValueOnlyCoder(StringUtf8Coder.of());
  IdentityDoFn<String> doFn =
      new IdentityDoFn<String>() {
        @FinishBundle
        public void finishBundle(FinishBundleContext context) {
          context.output(
              "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
        }
      };

  DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(
          outputTag,
          WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));

  return new DoFnOperator<>(
      doFn,
      "stepName",
      windowedValueCoder,
      Collections.emptyMap(),
      outputTag,
      Collections.emptyList(),
      outputManagerFactory,
      WindowingStrategy.globalDefault(),
      new HashMap<>(), /* side-input mapping */
      Collections.emptyList(), /* side inputs */
      options,
      null,
      null,
      DoFnSchemaInformation.create(),
      Collections.emptyMap());
}
 
Example 10
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 11
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 12
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 13
Source File: FlinkDoFnFunctionTest.java    From beam with Apache License 2.0 5 votes vote down vote up
@Test
public void testAccumulatorRegistrationOnOperatorClose() throws Exception {
  FlinkDoFnFunction doFnFunction =
      new TestDoFnFunction(
          "step",
          WindowingStrategy.globalDefault(),
          Collections.emptyMap(),
          PipelineOptionsFactory.create(),
          Collections.emptyMap(),
          new TupleTag<>(),
          null,
          Collections.emptyMap(),
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  doFnFunction.open(new Configuration());

  String metricContainerFieldName = "metricContainer";
  FlinkMetricContainer monitoredContainer =
      Mockito.spy(
          (FlinkMetricContainer)
              Whitebox.getInternalState(doFnFunction, metricContainerFieldName));
  Whitebox.setInternalState(doFnFunction, metricContainerFieldName, monitoredContainer);

  doFnFunction.close();
  Mockito.verify(monitoredContainer).registerMetricsForPipelineResult();
}
 
Example 14
Source File: DoFnOperatorTest.java    From beam with Apache License 2.0 4 votes vote down vote up
@Test
public void testBundleProcessingExceptionIsFatalDuringCheckpointing() throws Exception {
  FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
  options.setMaxBundleSize(10L);
  options.setCheckpointingInterval(1L);

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

  StringUtf8Coder coder = StringUtf8Coder.of();
  WindowedValue.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
      WindowedValue.getValueOnlyCoder(coder);

  DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(
          outputTag,
          WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));

  @SuppressWarnings("unchecked")
  DoFnOperator doFnOperator =
      new DoFnOperator<>(
          new IdentityDoFn() {
            @FinishBundle
            public void finishBundle() {
              throw new RuntimeException("something went wrong here");
            }
          },
          "stepName",
          windowedValueCoder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          outputManagerFactory,
          WindowingStrategy.globalDefault(),
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          options,
          null,
          null,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  @SuppressWarnings("unchecked")
  OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
      new OneInputStreamOperatorTestHarness<>(doFnOperator);

  testHarness.open();

  // start a bundle
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.valueInGlobalWindow("regular element")));

  // Make sure we throw Error, not a regular Exception.
  // A regular exception would just cause the checkpoint to fail.
  assertThrows(Error.class, () -> testHarness.snapshot(0, 0));
}
 
Example 15
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)));
}
 
Example 16
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 17
Source File: DoFnTransformTest.java    From incubator-nemo with Apache License 2.0 4 votes vote down vote up
@Test
@SuppressWarnings("unchecked")
public void testTimeBundle() {

  final long maxBundleTimeMills = 1000L;
  final TupleTag<String> outputTag = new TupleTag<>("main-output");
  final NemoPipelineOptions pipelineOptions = PipelineOptionsFactory.as(NemoPipelineOptions.class);
  pipelineOptions.setMaxBundleSize(10000000L);
  pipelineOptions.setMaxBundleTimeMills(maxBundleTimeMills);

  final List<Integer> bundleOutput = new ArrayList<>();

  final DoFnTransform<String, String> doFnTransform =
    new DoFnTransform<>(
      new BundleTestDoFn(bundleOutput),
      NULL_INPUT_CODER,
      NULL_OUTPUT_CODERS,
      outputTag,
      Collections.emptyList(),
      WindowingStrategy.globalDefault(),
      pipelineOptions,
      DisplayData.none(),
      DoFnSchemaInformation.create(),
      Collections.emptyMap());

  final Transform.Context context = mock(Transform.Context.class);
  final OutputCollector<WindowedValue<String>> oc = new TestOutputCollector<>();

  long startTime = System.currentTimeMillis();
  doFnTransform.prepare(context, oc);

  int count = 0;
  while (bundleOutput.isEmpty()) {
    doFnTransform.onData(WindowedValue.valueInGlobalWindow("a"));
    count += 1;
    try {
      Thread.sleep(10);
    } catch (InterruptedException e) {
      e.printStackTrace();
      throw new RuntimeException(e);
    }
  }

  long endTime = System.currentTimeMillis();
  assertEquals(count, (int) bundleOutput.get(0));
  assertTrue(endTime - startTime >= maxBundleTimeMills);

  doFnTransform.close();
}
 
Example 18
Source File: SimpleParDoFnTest.java    From beam with Apache License 2.0 4 votes vote down vote up
/**
 * Set up and execute a basic {@link ParDoFn} to validate reported counter values.
 *
 * @param inputData Input elements to process. For each element X, the DoFn will output a string
 *     repeated X times.
 * @return Delta counter updates extracted after execution.
 * @throws Exception
 */
private List<CounterUpdate> executeParDoFnCounterTest(int... inputData) throws Exception {
  class RepeaterDoFn extends DoFn<Integer, String> {
    /** Takes as input the number of times to output a message. */
    @ProcessElement
    public void processElement(ProcessContext c) {
      int numTimes = c.element();
      for (int i = 0; i < numTimes; i++) {
        c.output(String.format("I will repeat this message %d times", numTimes));
      }
    }
  }

  DoFn<Integer, String> fn = new RepeaterDoFn();
  DoFnInfo<?, ?> fnInfo =
      DoFnInfo.forFn(
          fn,
          WindowingStrategy.globalDefault(),
          null /* side input views */,
          null /* input coder */,
          MAIN_OUTPUT,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  ParDoFn parDoFn =
      new SimpleParDoFn<>(
          options,
          DoFnInstanceManagers.singleInstance(fnInfo),
          new EmptySideInputReader(),
          MAIN_OUTPUT,
          ImmutableMap.of(MAIN_OUTPUT, 0),
          stepContext,
          operationContext,
          DoFnSchemaInformation.create(),
          Collections.emptyMap(),
          SimpleDoFnRunnerFactory.INSTANCE);

  parDoFn.startBundle(new TestReceiver());
  for (int input : inputData) {
    parDoFn.processElement(WindowedValue.valueInGlobalWindow(input));
  }

  return operationContext
      .counterSet()
      .extractUpdates(true, DataflowCounterUpdateExtractor.INSTANCE);
}
 
Example 19
Source File: DoFnOperatorTest.java    From beam with Apache License 2.0 4 votes vote down vote up
@Test(expected = IllegalStateException.class)
public void testFailOnRequiresStableInputAndDisabledCheckpointing() {
  TupleTag<String> outputTag = new TupleTag<>("main-output");

  StringUtf8Coder keyCoder = StringUtf8Coder.of();
  KvToByteBufferKeySelector keySelector = new KvToByteBufferKeySelector<>(keyCoder);
  KvCoder<String, String> kvCoder = KvCoder.of(keyCoder, StringUtf8Coder.of());
  WindowedValue.ValueOnlyWindowedValueCoder<KV<String, String>> windowedValueCoder =
      WindowedValue.getValueOnlyCoder(kvCoder);

  DoFn<String, String> doFn =
      new DoFn<String, String>() {
        @ProcessElement
        // Use RequiresStableInput to force buffering elements
        @RequiresStableInput
        public void processElement(ProcessContext context) {
          context.output(context.element());
        }
      };

  DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(
          outputTag,
          WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));

  FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
  // should make the DoFnOperator creation fail
  options.setCheckpointingInterval(-1L);
  new DoFnOperator(
      doFn,
      "stepName",
      windowedValueCoder,
      Collections.emptyMap(),
      outputTag,
      Collections.emptyList(),
      outputManagerFactory,
      WindowingStrategy.globalDefault(),
      new HashMap<>(), /* side-input mapping */
      Collections.emptyList(), /* side inputs */
      options,
      keyCoder,
      keySelector,
      DoFnSchemaInformation.create(),
      Collections.emptyMap());
}
 
Example 20
Source File: FlinkPipelineOptionsTest.java    From beam with Apache License 2.0 4 votes vote down vote up
/** Tests that PipelineOptions are present after serialization. */
@Test
public void parDoBaseClassPipelineOptionsSerializationTest() throws Exception {

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

  Coder<WindowedValue<String>> coder = WindowedValue.getValueOnlyCoder(StringUtf8Coder.of());
  DoFnOperator<String, String> doFnOperator =
      new DoFnOperator<>(
          new TestDoFn(),
          "stepName",
          coder,
          Collections.emptyMap(),
          mainTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, coder),
          WindowingStrategy.globalDefault(),
          new HashMap<>(),
          Collections.emptyList(),
          options,
          null, /* key coder */
          null /* key selector */,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  final byte[] serialized = SerializationUtils.serialize(doFnOperator);

  @SuppressWarnings("unchecked")
  DoFnOperator<Object, Object> deserialized = SerializationUtils.deserialize(serialized);

  TypeInformation<WindowedValue<Object>> typeInformation =
      TypeInformation.of(new TypeHint<WindowedValue<Object>>() {});

  OneInputStreamOperatorTestHarness<WindowedValue<Object>, WindowedValue<Object>> testHarness =
      new OneInputStreamOperatorTestHarness<>(
          deserialized, typeInformation.createSerializer(new ExecutionConfig()));
  testHarness.open();

  // execute once to access options
  testHarness.processElement(
      new StreamRecord<>(
          WindowedValue.of(
              new Object(), Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)));

  testHarness.close();
}