Search in sources :

Example 1 with DoFnInfo

use of org.apache.beam.runners.dataflow.util.DoFnInfo in project beam by apache.

the class ProcessBundleHandler method createDoFnRunner.

/**
   * Converts a {@link org.apache.beam.fn.v1.BeamFnApi.FunctionSpec} into a {@link DoFnRunner}.
   */
private <InputT, OutputT> DoFnRunner<InputT, OutputT> createDoFnRunner(BeamFnApi.FunctionSpec functionSpec, Map<String, Collection<ThrowingConsumer<WindowedValue<OutputT>>>> outputMap) {
    ByteString serializedFn;
    try {
        serializedFn = functionSpec.getData().unpack(BytesValue.class).getValue();
    } catch (InvalidProtocolBufferException e) {
        throw new IllegalArgumentException(String.format("Unable to unwrap DoFn %s", functionSpec), e);
    }
    DoFnInfo<?, ?> doFnInfo = (DoFnInfo<?, ?>) SerializableUtils.deserializeFromByteArray(serializedFn.toByteArray(), "DoFnInfo");
    checkArgument(Objects.equals(new HashSet<>(Collections2.transform(outputMap.keySet(), Long::parseLong)), doFnInfo.getOutputMap().keySet()), "Unexpected mismatch between transform output map %s and DoFnInfo output map %s.", outputMap.keySet(), doFnInfo.getOutputMap());
    ImmutableMultimap.Builder<TupleTag<?>, ThrowingConsumer<WindowedValue<OutputT>>> tagToOutput = ImmutableMultimap.builder();
    for (Map.Entry<Long, TupleTag<?>> entry : doFnInfo.getOutputMap().entrySet()) {
        tagToOutput.putAll(entry.getValue(), outputMap.get(Long.toString(entry.getKey())));
    }
    @SuppressWarnings({ "unchecked", "rawtypes" }) final Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tagBasedOutputMap = (Map) tagToOutput.build().asMap();
    OutputManager outputManager = new OutputManager() {

        Map<TupleTag<?>, Collection<ThrowingConsumer<WindowedValue<?>>>> tupleTagToOutput = tagBasedOutputMap;

        @Override
        public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
            try {
                Collection<ThrowingConsumer<WindowedValue<?>>> consumers = tupleTagToOutput.get(tag);
                if (consumers == null) {
                    /* This is a normal case, e.g., if a DoFn has output but that output is not
                 * consumed. Drop the output. */
                    return;
                }
                for (ThrowingConsumer<WindowedValue<?>> consumer : consumers) {
                    consumer.accept(output);
                }
            } catch (Throwable t) {
                throw new RuntimeException(t);
            }
        }
    };
    @SuppressWarnings({ "unchecked", "rawtypes", "deprecation" }) DoFnRunner<InputT, OutputT> runner = DoFnRunners.simpleRunner(PipelineOptionsFactory.create(), /* TODO */
    (DoFn) doFnInfo.getDoFn(), NullSideInputReader.empty(), /* TODO */
    outputManager, (TupleTag) doFnInfo.getOutputMap().get(doFnInfo.getMainOutput()), new ArrayList<>(doFnInfo.getOutputMap().values()), new FakeStepContext(), (WindowingStrategy) doFnInfo.getWindowingStrategy());
    return runner;
}
Also used : ByteString(com.google.protobuf.ByteString) TupleTag(org.apache.beam.sdk.values.TupleTag) FakeStepContext(org.apache.beam.fn.harness.fake.FakeStepContext) WindowedValue(org.apache.beam.sdk.util.WindowedValue) ImmutableMultimap(com.google.common.collect.ImmutableMultimap) HashSet(java.util.HashSet) DoFnInfo(org.apache.beam.runners.dataflow.util.DoFnInfo) InvalidProtocolBufferException(com.google.protobuf.InvalidProtocolBufferException) Collection(java.util.Collection) ThrowingConsumer(org.apache.beam.fn.harness.fn.ThrowingConsumer) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) OutputManager(org.apache.beam.runners.core.DoFnRunners.OutputManager)

Aggregations

ImmutableMap (com.google.common.collect.ImmutableMap)1 ImmutableMultimap (com.google.common.collect.ImmutableMultimap)1 ByteString (com.google.protobuf.ByteString)1 InvalidProtocolBufferException (com.google.protobuf.InvalidProtocolBufferException)1 Collection (java.util.Collection)1 HashSet (java.util.HashSet)1 Map (java.util.Map)1 FakeStepContext (org.apache.beam.fn.harness.fake.FakeStepContext)1 ThrowingConsumer (org.apache.beam.fn.harness.fn.ThrowingConsumer)1 OutputManager (org.apache.beam.runners.core.DoFnRunners.OutputManager)1 DoFnInfo (org.apache.beam.runners.dataflow.util.DoFnInfo)1 WindowedValue (org.apache.beam.sdk.util.WindowedValue)1 TupleTag (org.apache.beam.sdk.values.TupleTag)1