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

Example 86 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project components by Talend.

the class JmsOutputPTransformRuntime method expand.

@Override
public PDone expand(PCollection<Object> objectPCollection) {
    // TODO remove this method from PCollection<Object> to PCollection<IndexedRecord>, as the incoming type always PCollection<IndexedRecord>
    PCollection<IndexedRecord> indexedCollection = objectPCollection.apply("ExtractIndexedRecord", ParDo.of(new DoFn<Object, IndexedRecord>() {

        IndexedRecordConverter converter;

        @DoFn.ProcessElement
        public void processElement(ProcessContext c) throws Exception {
            if (c.element() == null) {
                return;
            }
            if (converter == null) {
                converter = new AvroRegistry().createIndexedRecordConverter(c.element().getClass());
            }
            c.output((IndexedRecord) converter.convertToAvro(c.element()));
        }
    }));
    indexedCollection.setCoder(LazyAvroCoder.of());
    PCollection<String> jmsCollection = indexedCollection.apply("IndexedRecordToJmsRecord", ParDo.of(new DoFn<IndexedRecord, String>() {

        @DoFn.ProcessElement
        public void processElement(ProcessContext c) throws Exception {
            c.output(c.element().get(0).toString());
        }
    }));
    datastoreRuntime = new JmsDatastoreRuntime();
    datastoreRuntime.initialize(null, properties.datasetRef.getReference().getDatastoreProperties());
    if (messageType.equals(JmsMessageType.QUEUE)) {
        return jmsCollection.apply(JmsIO.write().withConnectionFactory(datastoreRuntime.getConnectionFactory()).withQueue(properties.datasetRef.getReference().queueTopicName.getValue()));
    } else if (messageType.equals(JmsMessageType.TOPIC)) {
        return jmsCollection.apply(JmsIO.write().withConnectionFactory(datastoreRuntime.getConnectionFactory()).withTopic(properties.datasetRef.getReference().queueTopicName.getValue()));
    } else {
        throw new TalendRuntimeException(CommonErrorCodes.UNEXPECTED_ARGUMENT);
    }
}

Example 87 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project component-runtime by Talend.

the class TalendIOTest method processorMulti.

@Test
public void processorMulti() {
    final PCollection<SampleLength> out = pipeline.apply(Create.of(new Sample("a"), new Sample("bb")).withCoder(JsonbCoder.of(Sample.class, PLUGIN))).apply(UUID.randomUUID().toString(), ParDo.of(new DoFn<Sample, JsonObject>() {

        @ProcessElement
        public void toData(final ProcessContext sample) {
            sample.output(JSONB.fromJson(JSONB.toJson(sample.element()), JsonObject.class));
        }
    })).setCoder(JsonpJsonObjectCoder.of(PLUGIN)).apply(new ViewsMappingTransform(emptyMap(), PLUGIN)).apply(TalendFn.asFn(new BaseTestProcessor() {

        @Override
        public void onNext(final InputFactory input, final OutputFactory factory) {
            factory.create(Branches.DEFAULT_BRANCH).emit(new SampleLength(JSONB.fromJson(input.read(Branches.DEFAULT_BRANCH).toString(), Sample.class).data.length()));
        }
    })).apply(ParDo.of(new DoFn<JsonObject, SampleLength>() {

        @ProcessElement
        public void onElement(final ProcessContext ctx) {
            ctx.output(JSONB.fromJson(ctx.element().getJsonArray("__default__").getJsonObject(0).toString(), SampleLength.class));
        }
    }));
    PAssert.that(out.apply(UUID.randomUUID().toString(), ParDo.of(new DoFn<SampleLength, Integer>() {

        @ProcessElement
        public void toInt(final ProcessContext pc) {
            pc.output(pc.element().len);
        }
    }))).containsInAnyOrder(1, 2);
    assertEquals(PipelineResult.State.DONE, pipeline.run().getState());
}

Example 88 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project component-runtime by Talend.

the class InMemoryQueueIOTest method output.

@Test(timeout = 60000)
public void output() {
    final Collection<JsonObject> objects = new ArrayList<>();
    try (final LoopState state = LoopState.newTracker(null)) {
        pipeline.apply(Create.of(IntStream.range(0, 5).mapToObj(RowStruct::new).collect(toList()))).setCoder(SerializableCoder.of(RowStruct.class)).apply(ParDo.of(new DoFn<RowStruct, JsonObject>() {

            @ProcessElement
            public void onElement(final ProcessContext context) {
                final JsonObject object = ComponentManager.instance().getJsonpBuilderFactory().createObjectBuilder().add("id", context.element().id).build();
                context.output(object);
            }
        })).setCoder(JsonpJsonObjectCoder.of(null)).apply(InMemoryQueueIO.to(state));
        pipeline.run().waitUntilFinish();
        JsonObject next;
        do {
            next = state.next();
            if (next != null) {
                objects.add(next);
            }
        } while (next != null);
    }
    assertEquals(5, objects.size());
    assertEquals(IntStream.range(0, 5).boxed().collect(toSet()), objects.stream().mapToInt(o -> o.getInt("id")).boxed().collect(toSet()));
}

Example 89 with DoFn

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

the class PTransformMatchersTest method parDoWithFnTypeWithMatchingType.

@Test
public void parDoWithFnTypeWithMatchingType() {
    DoFn<Object, Object> fn = new DoFn<Object, Object>() {

        @ProcessElement
        public void process(ProcessContext ctxt) {
        }
    };
    AppliedPTransform<?, ?, ?> parDoSingle = getAppliedTransform(ParDo.of(fn));
    AppliedPTransform<?, ?, ?> parDoMulti = getAppliedTransform(ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty()));
    PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(fn.getClass());
    assertThat(matcher.matches(parDoSingle), is(true));
    assertThat(matcher.matches(parDoMulti), is(true));
}

Example 90 with DoFn

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

the class DoFnOperatorTest method testLateDroppingForStatefulFn.

@Test
public void testLateDroppingForStatefulFn() throws Exception {
    WindowingStrategy<Object, IntervalWindow> windowingStrategy = WindowingStrategy.of(FixedWindows.of(Duration.millis(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, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())), windowingStrategy, new HashMap<>(), /* side-input mapping */
    Collections.emptyList(), /* side inputs */
    FlinkPipelineOptions.defaults(), keyCoder, /* key coder */
    keySelector, DoFnSchemaInformation.create(), Collections.emptyMap());
    OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness = new KeyedOneInputStreamOperatorTestHarness<>(doFnOperator, keySelector, new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of(), FlinkPipelineOptions.defaults()));
    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();
}