| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.dataflow.worker; |
| |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; |
| |
| import com.google.api.client.util.Throwables; |
| import com.google.api.services.dataflow.model.InstructionOutput; |
| import com.google.api.services.dataflow.model.MapTask; |
| import com.google.api.services.dataflow.model.MultiOutputInfo; |
| import com.google.api.services.dataflow.model.ParDoInstruction; |
| import com.google.api.services.dataflow.model.ParallelInstruction; |
| import com.google.api.services.dataflow.model.PartialGroupByKeyInstruction; |
| import com.google.api.services.dataflow.model.ReadInstruction; |
| import com.google.api.services.dataflow.model.Source; |
| import com.google.api.services.dataflow.model.WriteInstruction; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.HashMap; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.function.Function; |
| import org.apache.beam.model.pipeline.v1.Endpoints; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.runners.core.ElementByteSizeObservable; |
| import org.apache.beam.runners.core.SideInputReader; |
| import org.apache.beam.runners.dataflow.DataflowRunner; |
| import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions; |
| import org.apache.beam.runners.dataflow.util.CloudObject; |
| import org.apache.beam.runners.dataflow.util.CloudObjects; |
| import org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext; |
| import org.apache.beam.runners.dataflow.worker.counters.CounterFactory; |
| import org.apache.beam.runners.dataflow.worker.counters.CounterSet; |
| import org.apache.beam.runners.dataflow.worker.counters.NameContext; |
| import org.apache.beam.runners.dataflow.worker.fn.control.BeamFnMapTaskExecutor; |
| import org.apache.beam.runners.dataflow.worker.fn.control.ProcessRemoteBundleOperation; |
| import org.apache.beam.runners.dataflow.worker.fn.control.RegisterAndProcessBundleOperation; |
| import org.apache.beam.runners.dataflow.worker.fn.control.TimerReceiver; |
| import org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortReadOperation; |
| import org.apache.beam.runners.dataflow.worker.fn.data.RemoteGrpcPortWriteOperation; |
| import org.apache.beam.runners.dataflow.worker.graph.Edges.Edge; |
| import org.apache.beam.runners.dataflow.worker.graph.Edges.MultiOutputInfoEdge; |
| import org.apache.beam.runners.dataflow.worker.graph.Networks; |
| import org.apache.beam.runners.dataflow.worker.graph.Networks.TypeSafeNodeFunction; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.ExecutableStageNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.FetchAndFilterStreamingSideInputsNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.InstructionOutputNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.Node; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.OperationNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.OutputReceiverNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.ParallelInstructionNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.RegisterRequestNode; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.RemoteGrpcPortNode; |
| import org.apache.beam.runners.dataflow.worker.util.CloudSourceUtils; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.FlattenOperation; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.OperationContext; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoOperation; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ReadOperation; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.Receiver; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.Sink; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.WriteOperation; |
| import org.apache.beam.runners.fnexecution.GrpcFnServer; |
| import org.apache.beam.runners.fnexecution.control.InstructionRequestHandler; |
| import org.apache.beam.runners.fnexecution.control.JobBundleFactory; |
| import org.apache.beam.runners.fnexecution.control.SingleEnvironmentInstanceJobBundleFactory; |
| import org.apache.beam.runners.fnexecution.control.StageBundleFactory; |
| import org.apache.beam.runners.fnexecution.data.FnDataService; |
| import org.apache.beam.runners.fnexecution.data.GrpcDataService; |
| import org.apache.beam.runners.fnexecution.environment.StaticRemoteEnvironmentFactory; |
| import org.apache.beam.runners.fnexecution.state.GrpcStateService; |
| import org.apache.beam.runners.fnexecution.state.StateDelegator; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.fn.IdGenerator; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder; |
| import org.apache.beam.sdk.util.common.ElementByteSizeObserver; |
| import org.apache.beam.sdk.values.PCollectionView; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableTable; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.MutableNetwork; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.graph.Network; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** Creates a {@link DataflowMapTaskExecutor} from a {@link MapTask} definition. */ |
| public class BeamFnMapTaskExecutorFactory implements DataflowMapTaskExecutorFactory { |
| private static final Logger LOG = LoggerFactory.getLogger(BeamFnMapTaskExecutorFactory.class); |
| |
| public static BeamFnMapTaskExecutorFactory defaultFactory() { |
| return new BeamFnMapTaskExecutorFactory(); |
| } |
| |
| private BeamFnMapTaskExecutorFactory() {} |
| |
| /** |
| * Creates a new {@link DataflowMapTaskExecutor} from the given {@link MapTask} definition using |
| * the provided {@link ReaderFactory}. |
| */ |
| @Override |
| public DataflowMapTaskExecutor create( |
| InstructionRequestHandler instructionRequestHandler, |
| GrpcFnServer<GrpcDataService> grpcDataFnServer, |
| Endpoints.ApiServiceDescriptor dataApiServiceDescriptor, |
| GrpcFnServer<GrpcStateService> grpcStateFnServer, |
| MutableNetwork<Node, Edge> network, |
| PipelineOptions options, |
| String stageName, |
| ReaderFactory readerFactory, |
| SinkFactory sinkFactory, |
| DataflowExecutionContext<?> executionContext, |
| CounterSet counterSet, |
| IdGenerator idGenerator) { |
| |
| // TODO: remove this once we trust the code paths |
| checkArgument( |
| DataflowRunner.hasExperiment(options.as(DataflowPipelineDebugOptions.class), "beam_fn_api"), |
| "%s should only be used when beam_fn_api is enabled", |
| getClass().getSimpleName()); |
| |
| // Swap out all the InstructionOutput nodes with OutputReceiver nodes |
| Networks.replaceDirectedNetworkNodes( |
| network, createOutputReceiversTransform(stageName, counterSet)); |
| |
| if (DataflowRunner.hasExperiment( |
| options.as(DataflowPipelineDebugOptions.class), "use_executable_stage_bundle_execution")) { |
| LOG.debug("Using SingleEnvironmentInstanceJobBundleFactory"); |
| JobBundleFactory jobBundleFactory = |
| SingleEnvironmentInstanceJobBundleFactory.create( |
| StaticRemoteEnvironmentFactory.forService(instructionRequestHandler), |
| grpcDataFnServer, |
| grpcStateFnServer, |
| idGenerator); |
| // If the use_executable_stage_bundle_execution is enabled, use ExecutableStage instead. |
| Networks.replaceDirectedNetworkNodes( |
| network, |
| createOperationTransformForExecutableStageNode( |
| network, stageName, executionContext, jobBundleFactory)); |
| } else { |
| // Swap out all the RegisterFnRequest nodes with Operation nodes |
| Networks.replaceDirectedNetworkNodes( |
| network, |
| createOperationTransformForRegisterFnNodes( |
| idGenerator, |
| instructionRequestHandler, |
| grpcStateFnServer.getService(), |
| stageName, |
| executionContext)); |
| // Swap out all the RemoteGrpcPort nodes with Operation nodes, note that it is expected |
| // that the RegisterFnRequest nodes have already been replaced. |
| Networks.replaceDirectedNetworkNodes( |
| network, |
| createOperationTransformForGrpcPortNodes( |
| network, |
| grpcDataFnServer.getService(), |
| // TODO: Set NameContext properly for these operations. |
| executionContext.createOperationContext( |
| NameContext.create(stageName, stageName, stageName, stageName)))); |
| } |
| |
| // Swap out all the FetchAndFilterStreamingSideInput nodes with operation nodes |
| Networks.replaceDirectedNetworkNodes( |
| network, |
| createOperationTransformForFetchAndFilterStreamingSideInputNodes( |
| network, |
| idGenerator, |
| instructionRequestHandler, |
| grpcDataFnServer.getService(), |
| dataApiServiceDescriptor, |
| executionContext, |
| stageName)); |
| |
| // Swap out all the ParallelInstruction nodes with Operation nodes |
| Networks.replaceDirectedNetworkNodes( |
| network, |
| createOperationTransformForParallelInstructionNodes( |
| stageName, network, options, readerFactory, sinkFactory, executionContext)); |
| |
| // Collect all the operations within the network and attach all the operations as receivers |
| // to preceding output receivers. |
| List<Operation> topoSortedOperations = new ArrayList<>(); |
| for (OperationNode node : |
| Iterables.filter(Networks.topologicalOrder(network), OperationNode.class)) { |
| topoSortedOperations.add(node.getOperation()); |
| |
| for (Node predecessor : |
| Iterables.filter(network.predecessors(node), OutputReceiverNode.class)) { |
| ((OutputReceiverNode) predecessor) |
| .getOutputReceiver() |
| .addOutput((Receiver) node.getOperation()); |
| } |
| } |
| |
| if (LOG.isDebugEnabled()) { |
| LOG.info("Map task network: {}", Networks.toDot(network)); |
| } |
| |
| return BeamFnMapTaskExecutor.withSharedCounterSet( |
| topoSortedOperations, counterSet, executionContext.getExecutionStateTracker()); |
| } |
| |
| private Function<Node, Node> createOperationTransformForFetchAndFilterStreamingSideInputNodes( |
| MutableNetwork<Node, Edge> network, |
| IdGenerator idGenerator, |
| InstructionRequestHandler instructionRequestHandler, |
| FnDataService beamFnDataService, |
| Endpoints.ApiServiceDescriptor dataApiServiceDescriptor, |
| DataflowExecutionContext executionContext, |
| String stageName) { |
| return new TypeSafeNodeFunction<FetchAndFilterStreamingSideInputsNode>( |
| FetchAndFilterStreamingSideInputsNode.class) { |
| |
| @Override |
| public Node typedApply(FetchAndFilterStreamingSideInputsNode input) { |
| OutputReceiverNode output = |
| (OutputReceiverNode) Iterables.getOnlyElement(network.successors(input)); |
| DataflowOperationContext operationContext = |
| executionContext.createOperationContext( |
| NameContext.create( |
| stageName, |
| input.getNameContext().originalName(), |
| input.getNameContext().systemName(), |
| input.getNameContext().userName())); |
| return OperationNode.create( |
| new FetchAndFilterStreamingSideInputsOperation<>( |
| new OutputReceiver[] {output.getOutputReceiver()}, |
| operationContext, |
| instructionRequestHandler, |
| beamFnDataService, |
| dataApiServiceDescriptor, |
| idGenerator, |
| (Coder<WindowedValue<Object>>) output.getCoder(), |
| (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy(), |
| executionContext.getStepContext(operationContext), |
| input.getPCollectionViewsToWindowMappingFns())); |
| } |
| }; |
| } |
| |
| 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); |
| } |
| }; |
| } |
| |
| private Function<Node, Node> createOperationTransformForExecutableStageNode( |
| final Network<Node, Edge> network, |
| final String stageName, |
| final DataflowExecutionContext<?> executionContext, |
| final JobBundleFactory jobBundleFactory) { |
| return new TypeSafeNodeFunction<ExecutableStageNode>(ExecutableStageNode.class) { |
| @Override |
| public Node typedApply(ExecutableStageNode input) { |
| StageBundleFactory stageBundleFactory = |
| jobBundleFactory.forStage(input.getExecutableStage()); |
| Iterable<OutputReceiverNode> outputReceiverNodes = |
| Iterables.filter(network.successors(input), OutputReceiverNode.class); |
| |
| Map<String, OutputReceiver> outputReceiverMap = new HashMap<>(); |
| Lists.newArrayList(outputReceiverNodes).stream() |
| .forEach( |
| outputReceiverNode -> |
| outputReceiverMap.put( |
| outputReceiverNode.getPcollectionId(), |
| outputReceiverNode.getOutputReceiver())); |
| |
| DataflowOperationContext operationContextStage = |
| executionContext.createOperationContext( |
| NameContext.create(stageName, stageName, stageName, stageName)); |
| TimerReceiver timerReceiver = |
| new TimerReceiver( |
| input.getExecutableStage().getComponents(), |
| executionContext.getStepContext(operationContextStage).namespacedToUser(), |
| stageBundleFactory); |
| |
| ImmutableMap.Builder<String, DataflowOperationContext> |
| ptransformIdToOperationContextBuilder = 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); |
| } |
| |
| ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts = |
| ptransformIdToOperationContextBuilder.build(); |
| |
| ImmutableMap<String, SideInputReader> ptransformIdToSideInputReaders = |
| buildPTransformIdToSideInputReadersMap( |
| executionContext, input, ptransformIdToOperationContexts); |
| |
| Map<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> |
| ptransformIdToSideInputIdToPCollectionView = buildSideInputIdToPCollectionView(input); |
| |
| return OperationNode.create( |
| new ProcessRemoteBundleOperation( |
| input.getExecutableStage(), |
| executionContext.createOperationContext( |
| NameContext.create(stageName, stageName, stageName, stageName)), |
| stageBundleFactory, |
| outputReceiverMap, |
| timerReceiver, |
| ptransformIdToSideInputReaders, |
| ptransformIdToSideInputIdToPCollectionView)); |
| } |
| }; |
| } |
| |
| 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)))); |
| } |
| }; |
| } |
| |
| /** Returns a map from PTransform id to side input reader. */ |
| private static ImmutableMap<String, SideInputReader> buildPTransformIdToSideInputReadersMap( |
| DataflowExecutionContext executionContext, |
| RegisterRequestNode registerRequestNode, |
| ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts) { |
| |
| ImmutableMap.Builder<String, SideInputReader> ptransformIdToSideInputReaders = |
| ImmutableMap.builder(); |
| for (Map.Entry<String, Iterable<PCollectionView<?>>> ptransformIdToPCollectionView : |
| registerRequestNode.getPTransformIdToPCollectionViewMap().entrySet()) { |
| try { |
| ptransformIdToSideInputReaders.put( |
| ptransformIdToPCollectionView.getKey(), |
| executionContext.getSideInputReader( |
| // Note that the side input infos will only be populated for a batch pipeline |
| registerRequestNode |
| .getPTransformIdToSideInputInfoMap() |
| .get(ptransformIdToPCollectionView.getKey()), |
| ptransformIdToPCollectionView.getValue(), |
| ptransformIdToOperationContexts.get(ptransformIdToPCollectionView.getKey()))); |
| } catch (Exception e) { |
| throw Throwables.propagate(e); |
| } |
| } |
| return ptransformIdToSideInputReaders.build(); |
| } |
| |
| /** Returns a map from PTransform id to side input reader. */ |
| private static ImmutableMap<String, SideInputReader> buildPTransformIdToSideInputReadersMap( |
| DataflowExecutionContext executionContext, |
| ExecutableStageNode registerRequestNode, |
| ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts) { |
| |
| ImmutableMap.Builder<String, SideInputReader> ptransformIdToSideInputReaders = |
| ImmutableMap.builder(); |
| for (Map.Entry<String, Iterable<PCollectionView<?>>> ptransformIdToPCollectionView : |
| registerRequestNode.getPTransformIdToPCollectionViewMap().entrySet()) { |
| try { |
| ptransformIdToSideInputReaders.put( |
| ptransformIdToPCollectionView.getKey(), |
| executionContext.getSideInputReader( |
| // Note that the side input infos will only be populated for a batch pipeline |
| registerRequestNode |
| .getPTransformIdToSideInputInfoMap() |
| .get(ptransformIdToPCollectionView.getKey()), |
| ptransformIdToPCollectionView.getValue(), |
| ptransformIdToOperationContexts.get(ptransformIdToPCollectionView.getKey()))); |
| } catch (Exception e) { |
| throw Throwables.propagate(e); |
| } |
| } |
| return ptransformIdToSideInputReaders.build(); |
| } |
| |
| /** |
| * Returns a table where the row key is the PTransform id, the column key is the side input id, |
| * and the value is the corresponding PCollectionView. |
| */ |
| private static ImmutableTable<String, String, PCollectionView<?>> |
| buildPTransformIdToSideInputIdToPCollectionView(RegisterRequestNode registerRequestNode) { |
| ImmutableTable.Builder<String, String, PCollectionView<?>> |
| ptransformIdToSideInputIdToPCollectionViewBuilder = ImmutableTable.builder(); |
| for (Map.Entry<String, Iterable<PCollectionView<?>>> ptransformIdToPCollectionViews : |
| registerRequestNode.getPTransformIdToPCollectionViewMap().entrySet()) { |
| for (PCollectionView<?> pCollectionView : ptransformIdToPCollectionViews.getValue()) { |
| ptransformIdToSideInputIdToPCollectionViewBuilder.put( |
| ptransformIdToPCollectionViews.getKey(), |
| pCollectionView.getTagInternal().getId(), |
| pCollectionView); |
| } |
| } |
| |
| return ptransformIdToSideInputIdToPCollectionViewBuilder.build(); |
| } |
| |
| /** Returns a map where key is the SideInput id, value is PCollectionView. */ |
| private static Map<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> |
| buildSideInputIdToPCollectionView(ExecutableStageNode executableStageNode) { |
| ImmutableMap.Builder<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> |
| sideInputIdToPCollectionViewMapBuilder = ImmutableMap.builder(); |
| |
| for (Map.Entry<String, Iterable<PCollectionView<?>>> ptransformIdToPCollectionViews : |
| executableStageNode.getPTransformIdToPCollectionViewMap().entrySet()) { |
| for (PCollectionView<?> pCollectionView : ptransformIdToPCollectionViews.getValue()) { |
| sideInputIdToPCollectionViewMapBuilder.put( |
| RunnerApi.ExecutableStagePayload.SideInputId.newBuilder() |
| .setTransformId(ptransformIdToPCollectionViews.getKey()) |
| .setLocalName(pCollectionView.getTagInternal().getId()) |
| .build(), |
| pCollectionView); |
| } |
| } |
| |
| return sideInputIdToPCollectionViewMapBuilder.build(); |
| } |
| |
| /** |
| * Creates an {@link Operation} from the given {@link ParallelInstruction} definition using the |
| * provided {@link ReaderFactory}. |
| */ |
| Function<Node, Node> createOperationTransformForParallelInstructionNodes( |
| final String stageName, |
| final Network<Node, Edge> network, |
| final PipelineOptions options, |
| final ReaderFactory readerFactory, |
| final SinkFactory sinkFactory, |
| final DataflowExecutionContext<?> executionContext) { |
| |
| return new TypeSafeNodeFunction<ParallelInstructionNode>(ParallelInstructionNode.class) { |
| @Override |
| public Node typedApply(ParallelInstructionNode node) { |
| ParallelInstruction instruction = node.getParallelInstruction(); |
| NameContext nameContext = |
| NameContext.create( |
| stageName, |
| instruction.getOriginalName(), |
| instruction.getSystemName(), |
| instruction.getName()); |
| try { |
| DataflowOperationContext context = executionContext.createOperationContext(nameContext); |
| if (instruction.getRead() != null) { |
| return createReadOperation( |
| network, node, options, readerFactory, executionContext, context); |
| } else if (instruction.getWrite() != null) { |
| return createWriteOperation(node, options, sinkFactory, executionContext, context); |
| } else if (instruction.getParDo() != null) { |
| return createParDoOperation(network, node, options, executionContext, context); |
| } else if (instruction.getPartialGroupByKey() != null) { |
| return createPartialGroupByKeyOperation( |
| network, node, options, executionContext, context); |
| } else if (instruction.getFlatten() != null) { |
| return createFlattenOperation(network, node, context); |
| } else { |
| throw new IllegalArgumentException( |
| String.format("Unexpected instruction: %s", instruction)); |
| } |
| } catch (Exception e) { |
| throw new RuntimeException(e); |
| } |
| } |
| }; |
| } |
| |
| OperationNode createReadOperation( |
| Network<Node, Edge> network, |
| ParallelInstructionNode node, |
| PipelineOptions options, |
| ReaderFactory readerFactory, |
| DataflowExecutionContext<?> executionContext, |
| DataflowOperationContext operationContext) |
| throws Exception { |
| |
| ParallelInstruction instruction = node.getParallelInstruction(); |
| ReadInstruction read = instruction.getRead(); |
| |
| Source cloudSource = CloudSourceUtils.flattenBaseSpecs(read.getSource()); |
| CloudObject sourceSpec = CloudObject.fromSpec(cloudSource.getSpec()); |
| Coder<?> coder = |
| CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudSource.getCodec())); |
| NativeReader<?> reader = |
| readerFactory.create(sourceSpec, coder, options, executionContext, operationContext); |
| OutputReceiver[] receivers = getOutputReceivers(network, node); |
| return OperationNode.create(ReadOperation.create(reader, receivers, operationContext)); |
| } |
| |
| OperationNode createWriteOperation( |
| ParallelInstructionNode node, |
| PipelineOptions options, |
| SinkFactory sinkFactory, |
| DataflowExecutionContext executionContext, |
| DataflowOperationContext context) |
| throws Exception { |
| |
| ParallelInstruction instruction = node.getParallelInstruction(); |
| WriteInstruction write = instruction.getWrite(); |
| Coder<?> coder = |
| CloudObjects.coderFromCloudObject(CloudObject.fromSpec(write.getSink().getCodec())); |
| CloudObject cloudSink = CloudObject.fromSpec(write.getSink().getSpec()); |
| Sink<?> sink = sinkFactory.create(cloudSink, coder, options, executionContext, context); |
| return OperationNode.create(WriteOperation.create(sink, EMPTY_OUTPUT_RECEIVER_ARRAY, context)); |
| } |
| |
| private ParDoFnFactory parDoFnFactory = new DefaultParDoFnFactory(); |
| |
| private OperationNode createParDoOperation( |
| Network<Node, Edge> network, |
| ParallelInstructionNode node, |
| PipelineOptions options, |
| DataflowExecutionContext<?> executionContext, |
| DataflowOperationContext operationContext) |
| throws Exception { |
| |
| ParallelInstruction instruction = node.getParallelInstruction(); |
| ParDoInstruction parDo = instruction.getParDo(); |
| |
| TupleTag<?> mainOutputTag = tupleTag(parDo.getMultiOutputInfos().get(0)); |
| ImmutableMap.Builder<TupleTag<?>, Integer> outputTagsToReceiverIndicesBuilder = |
| ImmutableMap.builder(); |
| int successorOffset = 0; |
| for (Node successor : network.successors(node)) { |
| for (Edge edge : network.edgesConnecting(node, successor)) { |
| outputTagsToReceiverIndicesBuilder.put( |
| tupleTag(((MultiOutputInfoEdge) edge).getMultiOutputInfo()), successorOffset); |
| } |
| successorOffset += 1; |
| } |
| ParDoFn fn = |
| parDoFnFactory.create( |
| options, |
| CloudObject.fromSpec(parDo.getUserFn()), |
| parDo.getSideInputs(), |
| mainOutputTag, |
| outputTagsToReceiverIndicesBuilder.build(), |
| executionContext, |
| operationContext); |
| |
| OutputReceiver[] receivers = getOutputReceivers(network, node); |
| return OperationNode.create(new ParDoOperation(fn, receivers, operationContext)); |
| } |
| |
| private static <V> TupleTag<V> tupleTag(MultiOutputInfo multiOutputInfo) { |
| return new TupleTag<>(multiOutputInfo.getTag()); |
| } |
| |
| <K> OperationNode createPartialGroupByKeyOperation( |
| Network<Node, Edge> network, |
| ParallelInstructionNode node, |
| PipelineOptions options, |
| DataflowExecutionContext<?> executionContext, |
| DataflowOperationContext operationContext) |
| throws Exception { |
| |
| ParallelInstruction instruction = node.getParallelInstruction(); |
| PartialGroupByKeyInstruction pgbk = instruction.getPartialGroupByKey(); |
| OutputReceiver[] receivers = getOutputReceivers(network, node); |
| |
| Coder<?> windowedCoder = |
| CloudObjects.coderFromCloudObject(CloudObject.fromSpec(pgbk.getInputElementCodec())); |
| if (!(windowedCoder instanceof WindowedValueCoder)) { |
| throw new IllegalArgumentException( |
| String.format( |
| "unexpected kind of input coder for PartialGroupByKeyOperation: %s", windowedCoder)); |
| } |
| Coder<?> elemCoder = ((WindowedValueCoder<?>) windowedCoder).getValueCoder(); |
| if (!(elemCoder instanceof KvCoder)) { |
| throw new IllegalArgumentException( |
| String.format( |
| "unexpected kind of input element coder for PartialGroupByKeyOperation: %s", |
| elemCoder)); |
| } |
| |
| @SuppressWarnings("unchecked") |
| KvCoder<K, ?> keyedElementCoder = (KvCoder<K, ?>) elemCoder; |
| |
| CloudObject cloudUserFn = |
| pgbk.getValueCombiningFn() != null |
| ? CloudObject.fromSpec(pgbk.getValueCombiningFn()) |
| : null; |
| ParDoFn fn = |
| PartialGroupByKeyParDoFns.create( |
| options, |
| keyedElementCoder, |
| cloudUserFn, |
| pgbk.getSideInputs(), |
| Arrays.<Receiver>asList(receivers), |
| executionContext, |
| operationContext); |
| |
| return OperationNode.create(new ParDoOperation(fn, receivers, operationContext)); |
| } |
| |
| OperationNode createFlattenOperation( |
| Network<Node, Edge> network, ParallelInstructionNode node, OperationContext context) { |
| OutputReceiver[] receivers = getOutputReceivers(network, node); |
| return OperationNode.create(new FlattenOperation(receivers, context)); |
| } |
| |
| /** |
| * 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()); |
| } |
| }; |
| } |
| |
| private static final OutputReceiver[] EMPTY_OUTPUT_RECEIVER_ARRAY = new OutputReceiver[0]; |
| |
| static OutputReceiver[] getOutputReceivers(Network<Node, Edge> network, Node node) { |
| int outDegree = network.outDegree(node); |
| if (outDegree == 0) { |
| return EMPTY_OUTPUT_RECEIVER_ARRAY; |
| } |
| |
| OutputReceiver[] receivers = new OutputReceiver[outDegree]; |
| Iterator<Node> receiverNodes = network.successors(node).iterator(); |
| int i = 0; |
| do { |
| receivers[i] = ((OutputReceiverNode) receiverNodes.next()).getOutputReceiver(); |
| i += 1; |
| } while (receiverNodes.hasNext()); |
| |
| return receivers; |
| } |
| |
| /** Adapts a Coder to the ElementByteSizeObservable interface. */ |
| public static class ElementByteSizeObservableCoder<T> implements ElementByteSizeObservable<T> { |
| final Coder<T> coder; |
| |
| public ElementByteSizeObservableCoder(Coder<T> coder) { |
| this.coder = coder; |
| } |
| |
| @Override |
| public boolean isRegisterByteSizeObserverCheap(T value) { |
| return coder.isRegisterByteSizeObserverCheap(value); |
| } |
| |
| @Override |
| public void registerByteSizeObserver(T value, ElementByteSizeObserver observer) |
| throws Exception { |
| coder.registerByteSizeObserver(value, observer); |
| } |
| } |
| } |
