Search in sources :

Example 1 with AutoKVWrapper

use of org.talend.sdk.component.runtime.beam.transform.AutoKVWrapper in project component-runtime by Talend.

the class BeamExecutor method run.

@Override
public void run() {
    try {
        final Map<String, Mapper> mappers = delegate.getLevels().values().stream().flatMap(Collection::stream).filter(Job.Component::isSource).collect(toMap(Job.Component::getId, e -> delegate.getManager().findMapper(e.getNode().getFamily(), e.getNode().getComponent(), e.getNode().getVersion(), e.getNode().getConfiguration()).orElseThrow(() -> new IllegalStateException("No mapper found for: " + e.getNode()))));
        final Map<String, Processor> processors = delegate.getLevels().values().stream().flatMap(Collection::stream).filter(component -> !component.isSource()).collect(toMap(Job.Component::getId, e -> delegate.getManager().findProcessor(e.getNode().getFamily(), e.getNode().getComponent(), e.getNode().getVersion(), e.getNode().getConfiguration()).orElseThrow(() -> new IllegalStateException("No processor found for:" + e.getNode()))));
        final Pipeline pipeline = Pipeline.create(createPipelineOptions());
        final Map<String, PCollection<JsonObject>> pCollections = new HashMap<>();
        delegate.getLevels().values().stream().flatMap(Collection::stream).forEach(component -> {
            if (component.isSource()) {
                final Mapper mapper = mappers.get(component.getId());
                pCollections.put(component.getId(), pipeline.apply(toName("TalendIO", component), TalendIO.read(mapper)).apply(toName("RecordNormalizer", component), RecordNormalizer.of(mapper.plugin())));
            } else {
                final Processor processor = processors.get(component.getId());
                final List<Job.Edge> joins = getEdges(delegate.getEdges(), component, e -> e.getTo().getNode());
                final Map<String, PCollection<KV<String, JsonObject>>> inputs = joins.stream().collect(toMap(e -> e.getTo().getBranch(), e -> {
                    final PCollection<JsonObject> pc = pCollections.get(e.getFrom().getNode().getId());
                    final PCollection<JsonObject> filteredInput = pc.apply(toName("RecordBranchFilter", component, e), RecordBranchFilter.of(processor.plugin(), e.getFrom().getBranch()));
                    final PCollection<JsonObject> mappedInput;
                    if (e.getFrom().getBranch().equals(e.getTo().getBranch())) {
                        mappedInput = filteredInput;
                    } else {
                        mappedInput = filteredInput.apply(toName("RecordBranchMapper", component, e), RecordBranchMapper.of(processor.plugin(), e.getFrom().getBranch(), e.getTo().getBranch()));
                    }
                    return mappedInput.apply(toName("RecordBranchUnwrapper", component, e), RecordBranchUnwrapper.of(processor.plugin(), e.getTo().getBranch())).apply(toName("AutoKVWrapper", component, e), AutoKVWrapper.of(processor.plugin(), delegate.getKeyProvider(component.getId()), component.getId(), e.getFrom().getBranch()));
                }));
                KeyedPCollectionTuple<String> join = null;
                for (final Map.Entry<String, PCollection<KV<String, JsonObject>>> entry : inputs.entrySet()) {
                    final TupleTag<JsonObject> branch = new TupleTag<>(entry.getKey());
                    join = join == null ? KeyedPCollectionTuple.of(branch, entry.getValue()) : join.and(branch, entry.getValue());
                }
                final PCollection<JsonObject> preparedInput = join.apply(toName("CoGroupByKey", component), CoGroupByKey.create()).apply(toName("CoGroupByKeyResultMappingTransform", component), new CoGroupByKeyResultMappingTransform<>(processor.plugin(), true));
                if (getEdges(delegate.getEdges(), component, e -> e.getFrom().getNode()).isEmpty()) {
                    final PTransform<PCollection<JsonObject>, PDone> write = TalendIO.write(processor);
                    preparedInput.apply(toName("Output", component), write);
                } else {
                    final PTransform<PCollection<JsonObject>, PCollection<JsonObject>> process = TalendFn.asFn(processor);
                    pCollections.put(component.getId(), preparedInput.apply(toName("Processor", component), process));
                }
            }
        });
        final PipelineResult result = pipeline.run();
        // the wait until finish don't wait for the job to complete on the direct runner
        result.waitUntilFinish();
        while (PipelineResult.State.RUNNING.equals(result.getState())) {
            try {
                Thread.sleep(100L);
            } catch (final InterruptedException e) {
                throw new IllegalStateException("the job was aborted", e);
            }
        }
    } finally {
        delegate.getLevels().values().stream().flatMap(Collection::stream).map(Job.Component::getId).forEach(JobImpl.LocalSequenceHolder::clean);
    }
}
Also used : TalendIO(org.talend.sdk.component.runtime.beam.TalendIO) KV(org.apache.beam.sdk.values.KV) PipelineResult(org.apache.beam.sdk.PipelineResult) RecordBranchFilter(org.talend.sdk.component.runtime.beam.transform.RecordBranchFilter) HashMap(java.util.HashMap) PipelineOptionsFactory(org.apache.beam.sdk.options.PipelineOptionsFactory) Function(java.util.function.Function) PTransform(org.apache.beam.sdk.transforms.PTransform) RecordBranchMapper(org.talend.sdk.component.runtime.beam.transform.RecordBranchMapper) Collectors.toMap(java.util.stream.Collectors.toMap) TupleTag(org.apache.beam.sdk.values.TupleTag) Map(java.util.Map) RecordNormalizer(org.talend.sdk.component.runtime.beam.transform.RecordNormalizer) Pipeline(org.apache.beam.sdk.Pipeline) KeyedPCollectionTuple(org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) JsonObject(javax.json.JsonObject) PDone(org.apache.beam.sdk.values.PDone) Collection(java.util.Collection) PCollection(org.apache.beam.sdk.values.PCollection) Processor(org.talend.sdk.component.runtime.output.Processor) RecordBranchUnwrapper(org.talend.sdk.component.runtime.beam.transform.RecordBranchUnwrapper) CoGroupByKey(org.apache.beam.sdk.transforms.join.CoGroupByKey) AutoKVWrapper(org.talend.sdk.component.runtime.beam.transform.AutoKVWrapper) Collectors.toList(java.util.stream.Collectors.toList) List(java.util.List) Mapper(org.talend.sdk.component.runtime.input.Mapper) CoGroupByKeyResultMappingTransform(org.talend.sdk.component.runtime.beam.transform.CoGroupByKeyResultMappingTransform) Job(org.talend.sdk.component.runtime.manager.chain.Job) JobImpl(org.talend.sdk.component.runtime.manager.chain.internal.JobImpl) AllArgsConstructor(lombok.AllArgsConstructor) TalendFn(org.talend.sdk.component.runtime.beam.TalendFn) Processor(org.talend.sdk.component.runtime.output.Processor) HashMap(java.util.HashMap) JsonObject(javax.json.JsonObject) TupleTag(org.apache.beam.sdk.values.TupleTag) RecordBranchMapper(org.talend.sdk.component.runtime.beam.transform.RecordBranchMapper) Mapper(org.talend.sdk.component.runtime.input.Mapper) Job(org.talend.sdk.component.runtime.manager.chain.Job) PipelineResult(org.apache.beam.sdk.PipelineResult) Pipeline(org.apache.beam.sdk.Pipeline) PCollection(org.apache.beam.sdk.values.PCollection) PDone(org.apache.beam.sdk.values.PDone) Collection(java.util.Collection) PCollection(org.apache.beam.sdk.values.PCollection) HashMap(java.util.HashMap) Collectors.toMap(java.util.stream.Collectors.toMap) Map(java.util.Map)

Aggregations

Collection (java.util.Collection)1 HashMap (java.util.HashMap)1 List (java.util.List)1 Map (java.util.Map)1 Function (java.util.function.Function)1 Collectors.toList (java.util.stream.Collectors.toList)1 Collectors.toMap (java.util.stream.Collectors.toMap)1 JsonObject (javax.json.JsonObject)1 AllArgsConstructor (lombok.AllArgsConstructor)1 Pipeline (org.apache.beam.sdk.Pipeline)1 PipelineResult (org.apache.beam.sdk.PipelineResult)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 PipelineOptionsFactory (org.apache.beam.sdk.options.PipelineOptionsFactory)1 PTransform (org.apache.beam.sdk.transforms.PTransform)1 CoGroupByKey (org.apache.beam.sdk.transforms.join.CoGroupByKey)1 KeyedPCollectionTuple (org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple)1 KV (org.apache.beam.sdk.values.KV)1 PCollection (org.apache.beam.sdk.values.PCollection)1 PDone (org.apache.beam.sdk.values.PDone)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1