Examples with InferableFunction org.apache.beam.sdk.transforms.InferableFunction used on opensource projects

Example 1 with InferableFunction

use of org.apache.beam.sdk.transforms.InferableFunction in project beam by apache.

the class FlinkSavepointTest method createStreamingJob.

private static PCollection createStreamingJob(Pipeline pipeline, boolean restored, boolean isPortablePipeline) {
    final PCollection<KV<String, Long>> key;
    if (isPortablePipeline) {
        key = pipeline.apply("ImpulseStage", Impulse.create()).apply("KvMapperStage", MapElements.via(new InferableFunction<byte[], KV<String, Void>>() {

            @Override
            public KV<String, Void> apply(byte[] input) throws Exception {
                // https://jira.apache.org/jira/browse/BEAM-7144
                return KV.of("key", null);
            }
        })).apply("TimerStage", ParDo.of(new DoFn<KV<String, Void>, KV<String, Long>>() {

            @StateId("nextInteger")
            private final StateSpec<ValueState<Long>> valueStateSpec = StateSpecs.value();

            @TimerId("timer")
            private final TimerSpec timer = TimerSpecs.timer(TimeDomain.EVENT_TIME);

            @ProcessElement
            public void processElement(ProcessContext context, @TimerId("timer") Timer timer) {
                timer.set(new Instant(0));
            }

            @OnTimer("timer")
            public void onTimer(OnTimerContext context, @StateId("nextInteger") ValueState<Long> nextInteger, @TimerId("timer") Timer timer) {
                Long current = nextInteger.read();
                current = current != null ? current : 0L;
                context.output(KV.of("key", current));
                LOG.debug("triggering timer {}", current);
                nextInteger.write(current + 1);
                // Trigger timer again and continue to hold back the watermark
                timer.withOutputTimestamp(new Instant(0)).set(context.fireTimestamp());
            }
        }));
    } else {
        key = pipeline.apply("IdGeneratorStage", GenerateSequence.from(0)).apply("KvMapperStage", ParDo.of(new DoFn<Long, KV<String, Long>>() {

            @ProcessElement
            public void processElement(ProcessContext context) {
                context.output(KV.of("key", context.element()));
            }
        }));
    }
    if (restored) {
        return key.apply("VerificationStage", ParDo.of(new DoFn<KV<String, Long>, String>() {

            @StateId("valueState")
            private final StateSpec<ValueState<Integer>> valueStateSpec = StateSpecs.value();

            @StateId("bagState")
            private final StateSpec<BagState<Integer>> bagStateSpec = StateSpecs.bag();

            @ProcessElement
            public void processElement(ProcessContext context, @StateId("valueState") ValueState<Integer> intValueState, @StateId("bagState") BagState<Integer> intBagState) {
                assertThat(intValueState.read(), Matchers.is(42));
                assertThat(intBagState.read(), IsIterableContaining.hasItems(40, 1, 1));
                oneShotLatch.countDown();
            }
        }));
    } else {
        return key.apply("VerificationStage", ParDo.of(new DoFn<KV<String, Long>, String>() {

            @StateId("valueState")
            private final StateSpec<ValueState<Integer>> valueStateSpec = StateSpecs.value();

            @StateId("bagState")
            private final StateSpec<BagState<Integer>> bagStateSpec = StateSpecs.bag();

            @ProcessElement
            public void processElement(ProcessContext context, @StateId("valueState") ValueState<Integer> intValueState, @StateId("bagState") BagState<Integer> intBagState) {
                long value = Objects.requireNonNull(context.element().getValue());
                LOG.debug("value: {} timestamp: {}", value, context.timestamp().getMillis());
                if (value == 0L) {
                    intValueState.write(42);
                    intBagState.add(40);
                    intBagState.add(1);
                    intBagState.add(1);
                } else if (value >= 1) {
                    oneShotLatch.countDown();
                }
            }
        }));
    }
}