Examples with Node org.apache.beam.runners.dataflow.worker.graph.Nodes.Node used on opensource projects

Example 6 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.

the class IntrinsicMapTaskExecutorFactoryTest method testCreateFlattenOperation.

@Test
public void testCreateFlattenOperation() throws Exception {
    int producerIndex1 = 1;
    int producerOutputNum1 = 2;
    int producerIndex2 = 0;
    int producerOutputNum2 = 1;
    ParallelInstructionNode instructionNode = ParallelInstructionNode.create(createFlattenInstruction(producerIndex1, producerOutputNum1, producerIndex2, producerOutputNum2, "Flatten"), ExecutionLocation.UNKNOWN);
    when(network.successors(instructionNode)).thenReturn(ImmutableSet.<Node>of(IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet).apply(InstructionOutputNode.create(instructionNode.getParallelInstruction().getOutputs().get(0), PCOLLECTION_ID))));
    when(network.outDegree(instructionNode)).thenReturn(1);
    Node operationNode = mapTaskExecutorFactory.createOperationTransformForParallelInstructionNodes(STAGE, network, options, readerRegistry, sinkRegistry, BatchModeExecutionContext.forTesting(options, counterSet, "testStage")).apply(instructionNode);
    assertThat(operationNode, instanceOf(OperationNode.class));
    assertThat(((OperationNode) operationNode).getOperation(), instanceOf(FlattenOperation.class));
    FlattenOperation flattenOperation = (FlattenOperation) ((OperationNode) operationNode).getOperation();
    assertEquals(1, flattenOperation.receivers.length);
    assertEquals(0, flattenOperation.receivers[0].getReceiverCount());
    assertEquals(Operation.InitializationState.UNSTARTED, flattenOperation.initializationState);
}

Example 7 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.

the class IntrinsicMapTaskExecutorFactoryTest method testCreateWriteOperation.

@SuppressWarnings("unchecked")
@Test
public void testCreateWriteOperation() throws Exception {
    int producerIndex = 1;
    int producerOutputNum = 2;
    ParallelInstructionNode instructionNode = ParallelInstructionNode.create(createWriteInstruction(producerIndex, producerOutputNum, "WriteOperation"), ExecutionLocation.UNKNOWN);
    Node operationNode = mapTaskExecutorFactory.createOperationTransformForParallelInstructionNodes(STAGE, network, options, readerRegistry, sinkRegistry, BatchModeExecutionContext.forTesting(options, counterSet, "testStage")).apply(instructionNode);
    assertThat(operationNode, instanceOf(OperationNode.class));
    assertThat(((OperationNode) operationNode).getOperation(), instanceOf(WriteOperation.class));
    WriteOperation writeOperation = (WriteOperation) ((OperationNode) operationNode).getOperation();
    assertEquals(0, writeOperation.receivers.length);
    assertEquals(Operation.InitializationState.UNSTARTED, writeOperation.initializationState);
    assertThat(writeOperation.sink, instanceOf(SizeReportingSinkWrapper.class));
    assertThat(((SizeReportingSinkWrapper<?>) writeOperation.sink).getUnderlyingSink(), instanceOf(TestSink.class));
    counterSet.extractUpdates(false, updateExtractor);
    verify(updateExtractor).longSum(eq(named("WriteOperation-ByteCount")), anyBoolean(), anyLong());
    verifyNoMoreInteractions(updateExtractor);
}

Example 8 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.

the class IntrinsicMapTaskExecutorFactory method create.

/**
 * 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"), "experiment beam_fn_api turned on but non-Fn API MapTaskExecutorFactory invoked");
    // Swap out all the InstructionOutput nodes with OutputReceiver nodes
    Networks.replaceDirectedNetworkNodes(network, createOutputReceiversTransform(stageName, counterSet));
    // 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 IntrinsicMapTaskExecutor.withSharedCounterSet(topoSortedOperations, counterSet, executionContext.getExecutionStateTracker());
}

Example 9 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.

the class CreateRegisterFnOperationFunctionTest method testAllRunnerGraph.

@Test
public void testAllRunnerGraph() {
    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node parDoNode = createParDoNode("ParDo", Nodes.ExecutionLocation.RUNNER_HARNESS);
    Edge parDoNodeEdge = DefaultEdge.create();
    Node parDoNodeOut = createInstructionOutputNode("ParDo.out");
    // Read -out-> ParDo
    MutableNetwork<Node, Edge> expectedNetwork = createEmptyNetwork();
    expectedNetwork.addNode(readNode);
    expectedNetwork.addNode(readNodeOut);
    expectedNetwork.addNode(parDoNode);
    expectedNetwork.addNode(parDoNodeOut);
    expectedNetwork.addEdge(readNode, readNodeOut, readNodeEdge);
    expectedNetwork.addEdge(readNodeOut, parDoNode, readNodeOutEdge);
    expectedNetwork.addEdge(parDoNode, parDoNodeOut, parDoNodeEdge);
    MutableNetwork<Node, Edge> appliedNetwork = createRegisterFnOperation.apply(Graphs.copyOf(expectedNetwork));
    assertNetworkMaintainsBipartiteStructure(appliedNetwork);
    assertEquals(String.format("Expected network %s but got network %s", expectedNetwork, appliedNetwork), expectedNetwork, appliedNetwork);
}

Example 10 with Node

use of org.apache.beam.runners.dataflow.worker.graph.Nodes.Node in project beam by apache.

the class CreateRegisterFnOperationFunctionTest method testAllSdkGraph.

@Test
public void testAllSdkGraph() {
    Node sdkPortionNode = TestNode.create("SdkPortion");
    @SuppressWarnings({ "unchecked", "rawtypes" }) ArgumentCaptor<MutableNetwork<Node, Edge>> networkCapture = ArgumentCaptor.forClass((Class) MutableNetwork.class);
    when(registerFnOperationFunction.apply(networkCapture.capture())).thenReturn(sdkPortionNode);
    // Read -out-> ParDo
    Node readNode = createReadNode("Read", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge readNodeEdge = DefaultEdge.create();
    Node readNodeOut = createInstructionOutputNode("Read.out");
    Edge readNodeOutEdge = DefaultEdge.create();
    Node parDoNode = createParDoNode("ParDo", Nodes.ExecutionLocation.SDK_HARNESS);
    Edge parDoNodeEdge = DefaultEdge.create();
    Node parDoNodeOut = createInstructionOutputNode("ParDo.out");
    MutableNetwork<Node, Edge> network = createEmptyNetwork();
    network.addNode(readNode);
    network.addNode(readNodeOut);
    network.addNode(parDoNode);
    network.addNode(parDoNodeOut);
    network.addEdge(readNode, readNodeOut, readNodeEdge);
    network.addEdge(readNodeOut, parDoNode, readNodeOutEdge);
    network.addEdge(parDoNode, parDoNodeOut, parDoNodeEdge);
    MutableNetwork<Node, Edge> expectedNetwork = createEmptyNetwork();
    expectedNetwork.addNode(sdkPortionNode);
    MutableNetwork<Node, Edge> appliedNetwork = createRegisterFnOperation.apply(Graphs.copyOf(network));
    assertNetworkMaintainsBipartiteStructure(appliedNetwork);
    assertNetworkMaintainsBipartiteStructure(networkCapture.getValue());
    assertEquals(String.format("Expected network %s but got network %s", expectedNetwork, appliedNetwork), expectedNetwork, appliedNetwork);
    assertEquals(String.format("Expected network %s but got network %s", network, networkCapture.getValue()), network, networkCapture.getValue());
}