| /* |
| * 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.runners.dataflow.util.Structs.addString; |
| import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getElementCounterName; |
| import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getMeanByteCounterName; |
| import static org.apache.beam.runners.dataflow.worker.DataflowOutputCounter.getObjectCounterName; |
| import static org.apache.beam.runners.dataflow.worker.counters.CounterName.named; |
| import static org.apache.beam.sdk.util.SerializableUtils.serializeToByteArray; |
| import static org.apache.beam.sdk.util.StringUtils.byteArrayToJsonString; |
| import static org.hamcrest.Matchers.hasItems; |
| import static org.hamcrest.Matchers.instanceOf; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertSame; |
| import static org.junit.Assert.assertThat; |
| import static org.mockito.Matchers.anyBoolean; |
| import static org.mockito.Matchers.anyLong; |
| import static org.mockito.Matchers.eq; |
| import static org.mockito.Mockito.verify; |
| import static org.mockito.Mockito.verifyNoMoreInteractions; |
| import static org.mockito.Mockito.when; |
| |
| import com.google.api.services.dataflow.model.FlattenInstruction; |
| import com.google.api.services.dataflow.model.InstructionInput; |
| 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.io.IOException; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.List; |
| import java.util.function.Function; |
| import javax.annotation.Nullable; |
| import org.apache.beam.runners.dataflow.util.CloudObject; |
| import org.apache.beam.runners.dataflow.util.CloudObjects; |
| import org.apache.beam.runners.dataflow.util.PropertyNames; |
| import org.apache.beam.runners.dataflow.worker.apiary.FixMultiOutputInfosOnParDoInstructions; |
| import org.apache.beam.runners.dataflow.worker.counters.Counter; |
| import org.apache.beam.runners.dataflow.worker.counters.Counter.CounterUpdateExtractor; |
| import org.apache.beam.runners.dataflow.worker.counters.CounterFactory.CounterMean; |
| import org.apache.beam.runners.dataflow.worker.counters.CounterSet; |
| 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.MapTaskToNetworkFunction; |
| import org.apache.beam.runners.dataflow.worker.graph.Nodes.ExecutionLocation; |
| 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.ParallelInstructionNode; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.FlattenOperation; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.Operation; |
| 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.Sink; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.WriteOperation; |
| import org.apache.beam.sdk.coders.BigEndianIntegerCoder; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.CoderRegistry; |
| import org.apache.beam.sdk.coders.IterableCoder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.extensions.gcp.util.Transport; |
| import org.apache.beam.sdk.fn.IdGenerator; |
| import org.apache.beam.sdk.fn.IdGenerators; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.DoFnSchemaInformation; |
| import org.apache.beam.sdk.transforms.Sum; |
| import org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder; |
| import org.apache.beam.sdk.util.AppliedCombineFn; |
| import org.apache.beam.sdk.util.DoFnInfo; |
| import org.apache.beam.sdk.util.SerializableUtils; |
| import org.apache.beam.sdk.util.StringUtils; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.MutableNetwork; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.graph.Network; |
| import org.junit.Before; |
| import org.junit.Test; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| import org.mockito.Mock; |
| import org.mockito.Mockito; |
| import org.mockito.MockitoAnnotations; |
| |
| /** Tests for {@link IntrinsicMapTaskExecutorFactory}. */ |
| @RunWith(JUnit4.class) |
| public class IntrinsicMapTaskExecutorFactoryTest { |
| private static final String STAGE = "test"; |
| |
| private static final IdGenerator idGenerator = IdGenerators.decrementingLongs(); |
| |
| private static final Function<MapTask, MutableNetwork<Node, Edge>> mapTaskToNetwork = |
| new FixMultiOutputInfosOnParDoInstructions(idGenerator) |
| .andThen(new MapTaskToNetworkFunction(idGenerator)); |
| |
| private static final CloudObject windowedStringCoder = |
| CloudObjects.asCloudObject( |
| WindowedValue.getValueOnlyCoder(StringUtf8Coder.of()), /*sdkComponents=*/ null); |
| |
| private IntrinsicMapTaskExecutorFactory mapTaskExecutorFactory; |
| private PipelineOptions options; |
| private ReaderRegistry readerRegistry; |
| private SinkRegistry sinkRegistry; |
| private static final String PCOLLECTION_ID = "fakeId"; |
| |
| @Mock private Network<Node, Edge> network; |
| @Mock private CounterUpdateExtractor<?> updateExtractor; |
| |
| private final CounterSet counterSet = new CounterSet(); |
| |
| @Before |
| public void setUp() { |
| MockitoAnnotations.initMocks(this); |
| mapTaskExecutorFactory = IntrinsicMapTaskExecutorFactory.defaultFactory(); |
| options = PipelineOptionsFactory.create(); |
| readerRegistry = |
| ReaderRegistry.defaultRegistry() |
| .register( |
| ReaderFactoryTest.TestReaderFactory.class.getName(), |
| new ReaderFactoryTest.TestReaderFactory()) |
| .register( |
| ReaderFactoryTest.SingletonTestReaderFactory.class.getName(), |
| new ReaderFactoryTest.SingletonTestReaderFactory()); |
| |
| sinkRegistry = |
| SinkRegistry.defaultRegistry() |
| .register(TestSinkFactory.class.getName(), new TestSinkFactory()); |
| } |
| |
| @Test |
| public void testCreateMapTaskExecutor() throws Exception { |
| List<ParallelInstruction> instructions = |
| Arrays.asList( |
| createReadInstruction("Read"), |
| createParDoInstruction(0, 0, "DoFn1"), |
| createParDoInstruction(0, 0, "DoFnWithContext"), |
| createFlattenInstruction(1, 0, 2, 0, "Flatten"), |
| createWriteInstruction(3, 0, "Write")); |
| |
| MapTask mapTask = new MapTask(); |
| mapTask.setStageName(STAGE); |
| mapTask.setSystemName("systemName"); |
| mapTask.setInstructions(instructions); |
| mapTask.setFactory(Transport.getJsonFactory()); |
| |
| try (DataflowMapTaskExecutor executor = |
| mapTaskExecutorFactory.create( |
| null /* beamFnControlClientHandler */, |
| null /* GrpcFnServer<GrpcDataService> */, |
| null /* ApiServiceDescriptor */, |
| null, /* GrpcFnServer<GrpcStateService> */ |
| mapTaskToNetwork.apply(mapTask), |
| options, |
| STAGE, |
| readerRegistry, |
| sinkRegistry, |
| BatchModeExecutionContext.forTesting(options, counterSet, "testStage"), |
| counterSet, |
| idGenerator)) { |
| // Safe covariant cast not expressible without rawtypes. |
| @SuppressWarnings({"rawtypes", "unchecked"}) |
| List<Object> operations = (List) executor.operations; |
| assertThat( |
| operations, |
| hasItems( |
| instanceOf(ReadOperation.class), |
| instanceOf(ParDoOperation.class), |
| instanceOf(ParDoOperation.class), |
| instanceOf(FlattenOperation.class), |
| instanceOf(WriteOperation.class))); |
| |
| // Verify that the inputs are attached. |
| ReadOperation readOperation = |
| Iterables.getOnlyElement(Iterables.filter(operations, ReadOperation.class)); |
| assertEquals(2, readOperation.receivers[0].getReceiverCount()); |
| |
| FlattenOperation flattenOperation = |
| Iterables.getOnlyElement(Iterables.filter(operations, FlattenOperation.class)); |
| for (ParDoOperation operation : Iterables.filter(operations, ParDoOperation.class)) { |
| assertSame(flattenOperation, operation.receivers[0].getOnlyReceiver()); |
| } |
| WriteOperation writeOperation = |
| Iterables.getOnlyElement(Iterables.filter(operations, WriteOperation.class)); |
| assertSame(writeOperation, flattenOperation.receivers[0].getOnlyReceiver()); |
| } |
| |
| @SuppressWarnings("unchecked") |
| Counter<Long, ?> otherMsecCounter = |
| (Counter<Long, ?>) counterSet.getExistingCounter("test-other-msecs"); |
| |
| // "other" state only got created upon MapTaskExecutor.execute(). |
| assertNull(otherMsecCounter); |
| |
| counterSet.extractUpdates(false, updateExtractor); |
| verifyOutputCounters( |
| updateExtractor, |
| "read_output_name", |
| "DoFn1_output", |
| "DoFnWithContext_output", |
| "flatten_output_name"); |
| verify(updateExtractor).longSum(eq(named("Read-ByteCount")), anyBoolean(), anyLong()); |
| verify(updateExtractor).longSum(eq(named("Write-ByteCount")), anyBoolean(), anyLong()); |
| verifyNoMoreInteractions(updateExtractor); |
| } |
| |
| private static void verifyOutputCounters( |
| CounterUpdateExtractor<?> updateExtractor, String... outputNames) { |
| for (String outputName : outputNames) { |
| verify(updateExtractor) |
| .longSum(eq(named(getElementCounterName(outputName))), anyBoolean(), anyLong()); |
| verify(updateExtractor) |
| .longSum(eq(named(getObjectCounterName(outputName))), anyBoolean(), anyLong()); |
| verify(updateExtractor) |
| .longMean( |
| eq(named(getMeanByteCounterName(outputName))), |
| anyBoolean(), |
| Mockito.<CounterMean<Long>>any()); |
| } |
| } |
| |
| @Test |
| public void testExecutionContextPlumbing() throws Exception { |
| List<ParallelInstruction> instructions = |
| Arrays.asList( |
| createReadInstruction("Read", ReaderFactoryTest.SingletonTestReaderFactory.class), |
| createParDoInstruction(0, 0, "DoFn1", "DoFnUserName"), |
| createParDoInstruction(1, 0, "DoFnWithContext", "DoFnWithContextUserName")); |
| |
| MapTask mapTask = new MapTask(); |
| mapTask.setStageName(STAGE); |
| mapTask.setInstructions(instructions); |
| mapTask.setFactory(Transport.getJsonFactory()); |
| |
| BatchModeExecutionContext context = |
| BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); |
| |
| try (DataflowMapTaskExecutor executor = |
| mapTaskExecutorFactory.create( |
| null /* beamFnControlClientHandler */, |
| null /* beamFnDataService */, |
| null /* beamFnStateService */, |
| null, |
| mapTaskToNetwork.apply(mapTask), |
| options, |
| STAGE, |
| readerRegistry, |
| sinkRegistry, |
| context, |
| counterSet, |
| idGenerator)) { |
| executor.execute(); |
| } |
| |
| List<String> stepNames = new ArrayList<>(); |
| for (BatchModeExecutionContext.StepContext stepContext : context.getAllStepContexts()) { |
| stepNames.add(stepContext.getNameContext().systemName()); |
| } |
| assertThat(stepNames, hasItems("DoFn1", "DoFnWithContext")); |
| } |
| |
| static ParallelInstruction createReadInstruction(String name) { |
| return createReadInstruction(name, ReaderFactoryTest.TestReaderFactory.class); |
| } |
| |
| static ParallelInstruction createReadInstruction( |
| String name, Class<? extends ReaderFactory> readerFactoryClass) { |
| CloudObject spec = CloudObject.forClass(readerFactoryClass); |
| |
| Source cloudSource = new Source(); |
| cloudSource.setSpec(spec); |
| cloudSource.setCodec(windowedStringCoder); |
| |
| ReadInstruction readInstruction = new ReadInstruction(); |
| readInstruction.setSource(cloudSource); |
| |
| InstructionOutput output = new InstructionOutput(); |
| output.setName("read_output_name"); |
| output.setCodec(windowedStringCoder); |
| output.setOriginalName("originalName"); |
| output.setSystemName("systemName"); |
| |
| ParallelInstruction instruction = new ParallelInstruction(); |
| instruction.setSystemName(name); |
| instruction.setOriginalName(name + "OriginalName"); |
| instruction.setRead(readInstruction); |
| instruction.setOutputs(Arrays.asList(output)); |
| |
| return instruction; |
| } |
| |
| @Test |
| public void testCreateReadOperation() throws Exception { |
| ParallelInstructionNode instructionNode = |
| ParallelInstructionNode.create(createReadInstruction("Read"), 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, |
| PipelineOptionsFactory.create(), |
| readerRegistry, |
| sinkRegistry, |
| BatchModeExecutionContext.forTesting(options, counterSet, "testStage")) |
| .apply(instructionNode); |
| assertThat(operationNode, instanceOf(OperationNode.class)); |
| assertThat(((OperationNode) operationNode).getOperation(), instanceOf(ReadOperation.class)); |
| ReadOperation readOperation = (ReadOperation) ((OperationNode) operationNode).getOperation(); |
| |
| assertEquals(1, readOperation.receivers.length); |
| assertEquals(0, readOperation.receivers[0].getReceiverCount()); |
| assertEquals(Operation.InitializationState.UNSTARTED, readOperation.initializationState); |
| assertThat(readOperation.reader, instanceOf(ReaderFactoryTest.TestReader.class)); |
| |
| counterSet.extractUpdates(false, updateExtractor); |
| verifyOutputCounters(updateExtractor, "read_output_name"); |
| verify(updateExtractor).longSum(eq(named("Read-ByteCount")), anyBoolean(), anyLong()); |
| verifyNoMoreInteractions(updateExtractor); |
| } |
| |
| static ParallelInstruction createWriteInstruction( |
| int producerIndex, int producerOutputNum, String systemName) { |
| InstructionInput cloudInput = new InstructionInput(); |
| cloudInput.setProducerInstructionIndex(producerIndex); |
| cloudInput.setOutputNum(producerOutputNum); |
| |
| CloudObject spec = |
| CloudObject.forClass(IntrinsicMapTaskExecutorFactoryTest.TestSinkFactory.class); |
| |
| com.google.api.services.dataflow.model.Sink cloudSink = |
| new com.google.api.services.dataflow.model.Sink(); |
| cloudSink.setSpec(spec); |
| cloudSink.setCodec(windowedStringCoder); |
| |
| WriteInstruction writeInstruction = new WriteInstruction(); |
| writeInstruction.setInput(cloudInput); |
| writeInstruction.setSink(cloudSink); |
| |
| ParallelInstruction instruction = new ParallelInstruction(); |
| instruction.setWrite(writeInstruction); |
| instruction.setSystemName(systemName); |
| instruction.setOriginalName(systemName + "OriginalName"); |
| |
| return instruction; |
| } |
| |
| @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); |
| } |
| |
| static class TestDoFn extends DoFn<String, String> { |
| @ProcessElement |
| public void processElement(ProcessContext c) { |
| c.output(c.element()); |
| } |
| } |
| |
| static class TestSink extends Sink<Integer> { |
| @Override |
| public SinkWriter<Integer> writer() { |
| return new TestSinkWriter(); |
| } |
| |
| /** A sink writer that drops its input values, for testing. */ |
| static class TestSinkWriter implements SinkWriter<Integer> { |
| @Override |
| public long add(Integer outputElem) { |
| return 4; |
| } |
| |
| @Override |
| public void close() {} |
| |
| @Override |
| public void abort() throws IOException { |
| close(); |
| } |
| } |
| } |
| |
| static class TestSinkFactory implements SinkFactory { |
| @Override |
| public TestSink create( |
| CloudObject o, |
| @Nullable Coder<?> coder, |
| @Nullable PipelineOptions options, |
| @Nullable DataflowExecutionContext executionContext, |
| DataflowOperationContext operationContext) { |
| return new TestSink(); |
| } |
| } |
| |
| static ParallelInstruction createParDoInstruction( |
| int producerIndex, int producerOutputNum, String systemName) { |
| return createParDoInstruction(producerIndex, producerOutputNum, systemName, ""); |
| } |
| |
| static ParallelInstruction createParDoInstruction( |
| int producerIndex, int producerOutputNum, String systemName, String userName) { |
| InstructionInput cloudInput = new InstructionInput(); |
| cloudInput.setProducerInstructionIndex(producerIndex); |
| cloudInput.setOutputNum(producerOutputNum); |
| |
| TestDoFn fn = new TestDoFn(); |
| |
| String serializedFn = |
| StringUtils.byteArrayToJsonString( |
| SerializableUtils.serializeToByteArray( |
| DoFnInfo.forFn( |
| fn, |
| WindowingStrategy.globalDefault(), |
| null /* side input views */, |
| null /* input coder */, |
| new TupleTag<>(PropertyNames.OUTPUT) /* main output id */, |
| DoFnSchemaInformation.create()))); |
| |
| CloudObject cloudUserFn = CloudObject.forClassName("DoFn"); |
| addString(cloudUserFn, PropertyNames.SERIALIZED_FN, serializedFn); |
| |
| MultiOutputInfo mainOutputTag = new MultiOutputInfo(); |
| mainOutputTag.setTag("1"); |
| |
| ParDoInstruction parDoInstruction = new ParDoInstruction(); |
| parDoInstruction.setInput(cloudInput); |
| parDoInstruction.setNumOutputs(1); |
| parDoInstruction.setMultiOutputInfos(ImmutableList.of(mainOutputTag)); |
| parDoInstruction.setUserFn(cloudUserFn); |
| |
| InstructionOutput output = new InstructionOutput(); |
| output.setName(systemName + "_output"); |
| output.setCodec(windowedStringCoder); |
| output.setOriginalName("originalName"); |
| output.setSystemName("systemName"); |
| |
| ParallelInstruction instruction = new ParallelInstruction(); |
| instruction.setParDo(parDoInstruction); |
| instruction.setOutputs(Arrays.asList(output)); |
| instruction.setSystemName(systemName); |
| instruction.setOriginalName(systemName + "OriginalName"); |
| instruction.setName(userName); |
| return instruction; |
| } |
| |
| @Test |
| public void testCreateParDoOperation() throws Exception { |
| int producerIndex = 1; |
| int producerOutputNum = 2; |
| |
| BatchModeExecutionContext context = |
| BatchModeExecutionContext.forTesting(options, counterSet, "testStage"); |
| ParallelInstructionNode instructionNode = |
| ParallelInstructionNode.create( |
| createParDoInstruction(producerIndex, producerOutputNum, "DoFn"), |
| ExecutionLocation.UNKNOWN); |
| |
| Node outputReceiverNode = |
| IntrinsicMapTaskExecutorFactory.createOutputReceiversTransform(STAGE, counterSet) |
| .apply( |
| InstructionOutputNode.create( |
| instructionNode.getParallelInstruction().getOutputs().get(0), PCOLLECTION_ID)); |
| |
| when(network.successors(instructionNode)).thenReturn(ImmutableSet.of(outputReceiverNode)); |
| when(network.outDegree(instructionNode)).thenReturn(1); |
| when(network.edgesConnecting(instructionNode, outputReceiverNode)) |
| .thenReturn( |
| ImmutableSet.<Edge>of( |
| MultiOutputInfoEdge.create( |
| instructionNode |
| .getParallelInstruction() |
| .getParDo() |
| .getMultiOutputInfos() |
| .get(0)))); |
| |
| Node operationNode = |
| mapTaskExecutorFactory |
| .createOperationTransformForParallelInstructionNodes( |
| STAGE, network, options, readerRegistry, sinkRegistry, context) |
| .apply(instructionNode); |
| assertThat(operationNode, instanceOf(OperationNode.class)); |
| assertThat(((OperationNode) operationNode).getOperation(), instanceOf(ParDoOperation.class)); |
| ParDoOperation parDoOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); |
| |
| assertEquals(1, parDoOperation.receivers.length); |
| assertEquals(0, parDoOperation.receivers[0].getReceiverCount()); |
| assertEquals(Operation.InitializationState.UNSTARTED, parDoOperation.initializationState); |
| } |
| |
| static ParallelInstruction createPartialGroupByKeyInstruction( |
| int producerIndex, int producerOutputNum) { |
| InstructionInput cloudInput = new InstructionInput(); |
| cloudInput.setProducerInstructionIndex(producerIndex); |
| cloudInput.setOutputNum(producerOutputNum); |
| |
| PartialGroupByKeyInstruction pgbkInstruction = new PartialGroupByKeyInstruction(); |
| pgbkInstruction.setInput(cloudInput); |
| pgbkInstruction.setInputElementCodec( |
| CloudObjects.asCloudObject( |
| FullWindowedValueCoder.of( |
| KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of()), |
| IntervalWindowCoder.of()), |
| /*sdkComponents=*/ null)); |
| |
| InstructionOutput output = new InstructionOutput(); |
| output.setName("pgbk_output_name"); |
| output.setCodec( |
| CloudObjects.asCloudObject( |
| KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(BigEndianIntegerCoder.of())), |
| /*sdkComponents=*/ null)); |
| output.setOriginalName("originalName"); |
| output.setSystemName("systemName"); |
| |
| ParallelInstruction instruction = new ParallelInstruction(); |
| instruction.setOriginalName("pgbk_original_name"); |
| instruction.setSystemName("pgbk_system_name"); |
| instruction.setPartialGroupByKey(pgbkInstruction); |
| instruction.setOutputs(Arrays.asList(output)); |
| |
| return instruction; |
| } |
| |
| @Test |
| public void testCreatePartialGroupByKeyOperation() throws Exception { |
| int producerIndex = 1; |
| int producerOutputNum = 2; |
| |
| ParallelInstructionNode instructionNode = |
| ParallelInstructionNode.create( |
| createPartialGroupByKeyInstruction(producerIndex, producerOutputNum), |
| 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, |
| PipelineOptionsFactory.create(), |
| readerRegistry, |
| sinkRegistry, |
| BatchModeExecutionContext.forTesting(options, counterSet, "testStage")) |
| .apply(instructionNode); |
| assertThat(operationNode, instanceOf(OperationNode.class)); |
| assertThat(((OperationNode) operationNode).getOperation(), instanceOf(ParDoOperation.class)); |
| ParDoOperation pgbkOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); |
| |
| assertEquals(1, pgbkOperation.receivers.length); |
| assertEquals(0, pgbkOperation.receivers[0].getReceiverCount()); |
| assertEquals(Operation.InitializationState.UNSTARTED, pgbkOperation.initializationState); |
| } |
| |
| @Test |
| public void testCreatePartialGroupByKeyOperationWithCombine() throws Exception { |
| int producerIndex = 1; |
| int producerOutputNum = 2; |
| |
| ParallelInstruction instruction = |
| createPartialGroupByKeyInstruction(producerIndex, producerOutputNum); |
| |
| AppliedCombineFn<?, ?, ?, ?> combineFn = |
| AppliedCombineFn.withInputCoder( |
| Sum.ofIntegers(), |
| CoderRegistry.createDefault(), |
| KvCoder.of(StringUtf8Coder.of(), BigEndianIntegerCoder.of())); |
| CloudObject cloudCombineFn = CloudObject.forClassName("CombineFn"); |
| addString( |
| cloudCombineFn, |
| PropertyNames.SERIALIZED_FN, |
| byteArrayToJsonString(serializeToByteArray(combineFn))); |
| instruction.getPartialGroupByKey().setValueCombiningFn(cloudCombineFn); |
| |
| ParallelInstructionNode instructionNode = |
| ParallelInstructionNode.create(instruction, 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(ParDoOperation.class)); |
| ParDoOperation pgbkOperation = (ParDoOperation) ((OperationNode) operationNode).getOperation(); |
| |
| assertEquals(1, pgbkOperation.receivers.length); |
| assertEquals(0, pgbkOperation.receivers[0].getReceiverCount()); |
| assertEquals(Operation.InitializationState.UNSTARTED, pgbkOperation.initializationState); |
| } |
| |
| static ParallelInstruction createFlattenInstruction( |
| int producerIndex1, |
| int producerOutputNum1, |
| int producerIndex2, |
| int producerOutputNum2, |
| String systemName) { |
| List<InstructionInput> cloudInputs = new ArrayList<>(); |
| |
| InstructionInput cloudInput1 = new InstructionInput(); |
| cloudInput1.setProducerInstructionIndex(producerIndex1); |
| cloudInput1.setOutputNum(producerOutputNum1); |
| cloudInputs.add(cloudInput1); |
| |
| InstructionInput cloudInput2 = new InstructionInput(); |
| cloudInput2.setProducerInstructionIndex(producerIndex2); |
| cloudInput2.setOutputNum(producerOutputNum2); |
| cloudInputs.add(cloudInput2); |
| |
| FlattenInstruction flattenInstruction = new FlattenInstruction(); |
| flattenInstruction.setInputs(cloudInputs); |
| |
| InstructionOutput output = new InstructionOutput(); |
| output.setName("flatten_output_name"); |
| output.setCodec(CloudObjects.asCloudObject(StringUtf8Coder.of(), /*sdkComponents=*/ null)); |
| output.setOriginalName("originalName"); |
| output.setSystemName("systemName"); |
| |
| ParallelInstruction instruction = new ParallelInstruction(); |
| instruction.setFlatten(flattenInstruction); |
| instruction.setOutputs(Arrays.asList(output)); |
| instruction.setSystemName(systemName); |
| instruction.setOriginalName(systemName + "OriginalName"); |
| |
| return instruction; |
| } |
| |
| @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); |
| } |
| } |
