Examples with RegisterAndProcessBundleOperation org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation used on opensource projects

Example 1 with RegisterAndProcessBundleOperation

use of org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForGrpcPortNodes.

private Function<Node, Node> createOperationTransformForGrpcPortNodes(final Network<Node, Edge> network, final FnDataService beamFnDataService, final OperationContext context) {
    return new TypeSafeNodeFunction<RemoteGrpcPortNode>(RemoteGrpcPortNode.class) {

        @Override
        public Node typedApply(RemoteGrpcPortNode input) {
            RegisterAndProcessBundleOperation registerFnOperation = (RegisterAndProcessBundleOperation) Iterables.getOnlyElement(Iterables.filter(network.adjacentNodes(input), OperationNode.class)).getOperation();
            // The coder comes from the one and only adjacent output node
            Coder<?> coder = Iterables.getOnlyElement(Iterables.filter(network.adjacentNodes(input), OutputReceiverNode.class)).getCoder();
            // We figure out whether we are outputting some where if the output node is a
            // successor.
            Iterable<OutputReceiverNode> outputReceiverNodes = Iterables.filter(network.successors(input), OutputReceiverNode.class);
            Operation operation;
            if (outputReceiverNodes.iterator().hasNext()) {
                OutputReceiver[] outputReceivers = new OutputReceiver[] { Iterables.getOnlyElement(outputReceiverNodes).getOutputReceiver() };
                operation = new RemoteGrpcPortReadOperation<>(beamFnDataService, input.getPrimitiveTransformId(), registerFnOperation::getProcessBundleInstructionId, (Coder) coder, outputReceivers, context);
            } else {
                operation = new RemoteGrpcPortWriteOperation<>(beamFnDataService, input.getPrimitiveTransformId(), registerFnOperation::getProcessBundleInstructionId, (Coder) coder, context);
            }
            return OperationNode.create(operation);
        }
    };
}

Example 2 with RegisterAndProcessBundleOperation

use of org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation 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))));
        }
    };
}