Search in sources :

Example 6 with TypeSafeNodeFunction

use of org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction in project beam by apache.

the class IntrinsicMapTaskExecutorFactory method createOutputReceiversTransform.

/**
 * Returns a function which can convert {@link InstructionOutput}s into {@link OutputReceiver}s.
 */
static Function<Node, Node> createOutputReceiversTransform(final String stageName, final CounterFactory counterFactory) {
    return new TypeSafeNodeFunction<InstructionOutputNode>(InstructionOutputNode.class) {

        @Override
        public Node typedApply(InstructionOutputNode input) {
            InstructionOutput cloudOutput = input.getInstructionOutput();
            OutputReceiver outputReceiver = new OutputReceiver();
            Coder<?> coder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec()));
            @SuppressWarnings("unchecked") ElementCounter outputCounter = new DataflowOutputCounter(cloudOutput.getName(), new ElementByteSizeObservableCoder<>(coder), counterFactory, NameContext.create(stageName, cloudOutput.getOriginalName(), cloudOutput.getSystemName(), cloudOutput.getName()));
            outputReceiver.addOutputCounter(outputCounter);
            return OutputReceiverNode.create(outputReceiver, coder, input.getPcollectionId());
        }
    };
}
Also used : InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) InstructionOutput(com.google.api.services.dataflow.model.InstructionOutput) OutputReceiver(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver) TypeSafeNodeFunction(org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction) ElementCounter(org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter)

Example 7 with TypeSafeNodeFunction

use of org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction in project beam by apache.

the class LengthPrefixUnknownCoders method andReplaceForParallelInstructionNode.

/**
 * Replace unknown coders on the given {@link ParallelInstructionNode} with {@link
 * org.apache.beam.sdk.coders.LengthPrefixCoder LengthPrefixCoder&lt;T&gt;} where {@code T} is a
 * {@link org.apache.beam.sdk.coders.ByteArrayCoder}.
 */
private static Function<Node, Node> andReplaceForParallelInstructionNode() {
    return new TypeSafeNodeFunction<ParallelInstructionNode>(ParallelInstructionNode.class) {

        @Override
        public Node typedApply(ParallelInstructionNode input) {
            ParallelInstruction instruction = input.getParallelInstruction();
            Nodes.ExecutionLocation location = input.getExecutionLocation();
            try {
                instruction = forParallelInstruction(instruction, true);
            } catch (Exception e) {
                throw new RuntimeException(String.format("Failed to replace unknown coder with " + "LengthPrefixCoder for : {%s}", input.getParallelInstruction()), e);
            }
            return ParallelInstructionNode.create(instruction, location);
        }
    };
}
Also used : ParallelInstruction(com.google.api.services.dataflow.model.ParallelInstruction) ParallelInstructionNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode) TypeSafeNodeFunction(org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction) IOException(java.io.IOException)

Example 8 with TypeSafeNodeFunction

use of org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOutputReceiversTransform.

/**
 * Returns a function which can convert {@link InstructionOutput}s into {@link OutputReceiver}s.
 */
static Function<Node, Node> createOutputReceiversTransform(final String stageName, final CounterFactory counterFactory) {
    return new TypeSafeNodeFunction<InstructionOutputNode>(InstructionOutputNode.class) {

        @Override
        public Node typedApply(InstructionOutputNode input) {
            InstructionOutput cloudOutput = input.getInstructionOutput();
            OutputReceiver outputReceiver = new OutputReceiver();
            Coder<?> coder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec()));
            @SuppressWarnings("unchecked") ElementCounter outputCounter = new DataflowOutputCounter(cloudOutput.getName(), new ElementByteSizeObservableCoder<>(coder), counterFactory, NameContext.create(stageName, cloudOutput.getOriginalName(), cloudOutput.getSystemName(), cloudOutput.getName()));
            outputReceiver.addOutputCounter(outputCounter);
            return OutputReceiverNode.create(outputReceiver, coder, input.getPcollectionId());
        }
    };
}
Also used : InstructionOutputNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode) InstructionOutput(com.google.api.services.dataflow.model.InstructionOutput) OutputReceiver(org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver) TypeSafeNodeFunction(org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction) ElementCounter(org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter)

Example 9 with TypeSafeNodeFunction

use of org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForRegisterFnNodes.

private Function<Node, Node> createOperationTransformForRegisterFnNodes(final IdGenerator idGenerator, final InstructionRequestHandler instructionRequestHandler, final StateDelegator beamFnStateDelegator, final String stageName, final DataflowExecutionContext<?> executionContext) {
    return new TypeSafeNodeFunction<RegisterRequestNode>(RegisterRequestNode.class) {

        @Override
        public Node typedApply(RegisterRequestNode input) {
            ImmutableMap.Builder<String, DataflowOperationContext> ptransformIdToOperationContextBuilder = ImmutableMap.builder();
            ImmutableMap.Builder<String, DataflowStepContext> ptransformIdToStepContext = ImmutableMap.builder();
            for (Map.Entry<String, NameContext> entry : input.getPTransformIdToPartialNameContextMap().entrySet()) {
                NameContext fullNameContext = NameContext.create(stageName, entry.getValue().originalName(), entry.getValue().systemName(), entry.getValue().userName());
                DataflowOperationContext operationContext = executionContext.createOperationContext(fullNameContext);
                ptransformIdToOperationContextBuilder.put(entry.getKey(), operationContext);
                ptransformIdToStepContext.put(entry.getKey(), executionContext.getStepContext(operationContext));
            }
            ImmutableMap.Builder<String, NameContext> pcollectionIdToNameContext = ImmutableMap.builder();
            for (Map.Entry<String, NameContext> entry : input.getPCollectionToPartialNameContextMap().entrySet()) {
                pcollectionIdToNameContext.put(entry.getKey(), NameContext.create(stageName, entry.getValue().originalName(), entry.getValue().systemName(), entry.getValue().userName()));
            }
            ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts = ptransformIdToOperationContextBuilder.build();
            ImmutableMap<String, SideInputReader> ptransformIdToSideInputReaders = buildPTransformIdToSideInputReadersMap(executionContext, input, ptransformIdToOperationContexts);
            ImmutableTable<String, String, PCollectionView<?>> ptransformIdToSideInputIdToPCollectionView = buildPTransformIdToSideInputIdToPCollectionView(input);
            return OperationNode.create(new RegisterAndProcessBundleOperation(idGenerator, instructionRequestHandler, beamFnStateDelegator, input.getRegisterRequest(), ptransformIdToOperationContexts, ptransformIdToStepContext.build(), ptransformIdToSideInputReaders, ptransformIdToSideInputIdToPCollectionView, pcollectionIdToNameContext.build(), // TODO: Set NameContext properly for these operations.
            executionContext.createOperationContext(NameContext.create(stageName, stageName, stageName, stageName))));
        }
    };
}
Also used : NameContext(org.apache.beam.runners.dataflow.worker.counters.NameContext) SideInputReader(org.apache.beam.runners.core.SideInputReader) DataflowStepContext(org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext) RegisterAndProcessBundleOperation(org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) RegisterRequestNode(org.apache.beam.runners.dataflow.worker.graph.Nodes.RegisterRequestNode) PCollectionView(org.apache.beam.sdk.values.PCollectionView) TypeSafeNodeFunction(org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) Map(java.util.Map) HashMap(java.util.HashMap)

Aggregations

TypeSafeNodeFunction (org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction)9 OutputReceiver (org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver)5 NameContext (org.apache.beam.runners.dataflow.worker.counters.NameContext)4 ParallelInstruction (com.google.api.services.dataflow.model.ParallelInstruction)3 OutputReceiverNode (org.apache.beam.runners.dataflow.worker.graph.Nodes.OutputReceiverNode)3 ParallelInstructionNode (org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode)3 InstructionOutput (com.google.api.services.dataflow.model.InstructionOutput)2 HashMap (java.util.HashMap)2 Map (java.util.Map)2 SideInputReader (org.apache.beam.runners.core.SideInputReader)2 ProcessRemoteBundleOperation (org.apache.beam.runners.dataflow.worker.fn.control.ProcessRemoteBundleOperation)2 RegisterAndProcessBundleOperation (org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation)2 InstructionOutputNode (org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode)2 ElementCounter (org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter)2 Coder (org.apache.beam.sdk.coders.Coder)2 KvCoder (org.apache.beam.sdk.coders.KvCoder)2 WindowedValueCoder (org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder)2 PCollectionView (org.apache.beam.sdk.values.PCollectionView)2 ImmutableMap (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap)2 IOException (java.io.IOException)1