Examples with ParDoFn org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn used on opensource projects

Example 1 with ParDoFn

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn in project beam by apache.

the class CombineValuesFnFactoryTest method testCombineValuesFnAll.

@Test
public void testCombineValuesFnAll() throws Exception {
    TestReceiver receiver = new TestReceiver();
    Combine.CombineFn<Integer, CountSum, String> combiner = (new MeanInts());
    ParDoFn combineParDoFn = createCombineValuesFn(CombinePhase.ALL, combiner, StringUtf8Coder.of(), BigEndianIntegerCoder.of(), new CountSumCoder(), WindowingStrategy.globalDefault());
    combineParDoFn.startBundle(receiver);
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("a", Arrays.asList(5, 6, 7))));
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("b", Arrays.asList(1, 3, 7))));
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("c", Arrays.asList(3, 6, 8, 9))));
    combineParDoFn.finishBundle();
    Object[] expectedReceivedElems = { WindowedValue.valueInGlobalWindow(KV.of("a", String.format("%.1f", 6.0))), WindowedValue.valueInGlobalWindow(KV.of("b", String.format("%.1f", 3.7))), WindowedValue.valueInGlobalWindow(KV.of("c", String.format("%.1f", 6.5))) };
    assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}

Example 2 with ParDoFn

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn in project beam by apache.

the class CombineValuesFnFactoryTest method testCombineValuesFnAdd.

@Test
public void testCombineValuesFnAdd() throws Exception {
    TestReceiver receiver = new TestReceiver();
    MeanInts mean = new MeanInts();
    Combine.CombineFn<Integer, CountSum, String> combiner = mean;
    ParDoFn combineParDoFn = createCombineValuesFn(CombinePhase.ADD, combiner, StringUtf8Coder.of(), BigEndianIntegerCoder.of(), new CountSumCoder(), WindowingStrategy.globalDefault());
    combineParDoFn.startBundle(receiver);
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("a", Arrays.asList(5, 6, 7))));
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("b", Arrays.asList(1, 3, 7))));
    combineParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("c", Arrays.asList(3, 6, 8, 9))));
    combineParDoFn.finishBundle();
    Object[] expectedReceivedElems = { WindowedValue.valueInGlobalWindow(KV.of("a", new CountSum(3, 18))), WindowedValue.valueInGlobalWindow(KV.of("b", new CountSum(3, 11))), WindowedValue.valueInGlobalWindow(KV.of("c", new CountSum(4, 26))) };
    assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}

Example 3 with ParDoFn

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createParDoOperation.

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));
}

Example 4 with ParDoFn

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn in project beam by apache.

the class ToIsmRecordForMultimapDoFnFactoryTest method testConversionOfRecord.

@Test
public void testConversionOfRecord() throws Exception {
    ParDoFn parDoFn = new ToIsmRecordForMultimapDoFnFactory().create(null, /* pipeline options */
    CloudObject.fromSpec(ImmutableMap.of(PropertyNames.OBJECT_TYPE_NAME, "ToIsmRecordForMultimapDoFn", PropertyNames.ENCODING, createIsmRecordEncoding())), null, /* side input infos */
    null, /* main output tag */
    null, /* output tag to receiver index */
    null, /* exection context */
    null);
    List<Object> outputReceiver = new ArrayList<>();
    parDoFn.startBundle(outputReceiver::add);
    parDoFn.processElement(valueInGlobalWindow(KV.of(12, /* shard key */
    ImmutableList.of(KV.of(KV.of(42, /* user key */
    GlobalWindow.INSTANCE), /* sort key */
    4), KV.of(KV.of(42, /* user key */
    GlobalWindow.INSTANCE), /* sort key */
    5), KV.of(KV.of(43, /* user key */
    GlobalWindow.INSTANCE), /* sort key */
    6), KV.of(KV.of(44, /* user key */
    GlobalWindow.INSTANCE), /* sort key */
    7), KV.of(KV.of(44, /* user key */
    GlobalWindow.INSTANCE), /* sort key */
    8)))));
    assertThat(outputReceiver, contains(valueInGlobalWindow(IsmRecord.of(ImmutableList.of(42, GlobalWindow.INSTANCE, 0L), 4)), /* same structural value  as above */
    valueInGlobalWindow(IsmRecord.of(ImmutableList.of(42, GlobalWindow.INSTANCE, 1L), 5)), valueInGlobalWindow(IsmRecord.of(ImmutableList.of(43, GlobalWindow.INSTANCE, 0L), 6)), valueInGlobalWindow(IsmRecord.of(ImmutableList.of(44, GlobalWindow.INSTANCE, 0L), 7)), /* same structural value as above and final value */
    valueInGlobalWindow(IsmRecord.of(ImmutableList.of(44, GlobalWindow.INSTANCE, 1L), 8))));
}

Example 5 with ParDoFn

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn in project beam by apache.

the class UserParDoFnFactoryTest method testFactoryDoesNotReuseAfterAborted.

@Test
public void testFactoryDoesNotReuseAfterAborted() throws Exception {
    PipelineOptions options = PipelineOptionsFactory.create();
    CounterSet counters = new CounterSet();
    TestDoFn initialFn = new TestDoFn(Collections.<TupleTag<String>>emptyList());
    CloudObject cloudObject = getCloudObject(initialFn);
    ParDoFn parDoFn = factory.create(options, cloudObject, null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters));
    Receiver rcvr = new OutputReceiver();
    parDoFn.startBundle(rcvr);
    parDoFn.processElement(WindowedValue.valueInGlobalWindow("foo"));
    TestDoFn fn = (TestDoFn) ((SimpleParDoFn) parDoFn).getDoFnInfo().getDoFn();
    parDoFn.abort();
    assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN));
    // The fn should not be torn down here
    ParDoFn secondParDoFn = factory.create(options, cloudObject.clone(), null, MAIN_OUTPUT, ImmutableMap.<TupleTag<?>, Integer>of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage"), TestOperationContext.create(counters));
    secondParDoFn.startBundle(rcvr);
    secondParDoFn.processElement(WindowedValue.valueInGlobalWindow("foo"));
    TestDoFn secondFn = (TestDoFn) ((SimpleParDoFn) secondParDoFn).getDoFnInfo().getDoFn();
    assertThat(secondFn, not(theInstance(fn)));
    assertThat(fn.state, equalTo(TestDoFn.State.TORN_DOWN));
    assertThat(secondFn.state, equalTo(TestDoFn.State.PROCESSING));
}