Examples with RemoteBundle org.apache.beam.runners.fnexecution.control.RemoteBundle used on opensource projects

Example 11 with RemoteBundle

use of org.apache.beam.runners.fnexecution.control.RemoteBundle in project beam by apache.

the class FlinkExecutableStageFunctionTest method setUpMocks.

@Before
public void setUpMocks() throws Exception {
    MockitoAnnotations.initMocks(this);
    when(runtimeContext.getDistributedCache()).thenReturn(distributedCache);
    when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
    RemoteBundle remoteBundle = Mockito.mock(RemoteBundle.class);
    when(stageBundleFactory.getBundle(any(), any(StateRequestHandler.class), any(BundleProgressHandler.class), any(BundleFinalizationHandler.class), any(BundleCheckpointHandler.class))).thenReturn(remoteBundle);
    when(stageBundleFactory.getBundle(any(), any(TimerReceiverFactory.class), any(StateRequestHandler.class), any(BundleProgressHandler.class))).thenReturn(remoteBundle);
    ImmutableMap input = ImmutableMap.builder().put("input", Mockito.mock(FnDataReceiver.class)).build();
    when(remoteBundle.getInputReceivers()).thenReturn(input);
    when(processBundleDescriptor.getTimerSpecs()).thenReturn(Collections.emptyMap());
}

Example 12 with RemoteBundle

use of org.apache.beam.runners.fnexecution.control.RemoteBundle in project beam by apache.

the class FlinkExecutableStageFunctionTest method outputsAreTaggedCorrectly.

@Test
public void outputsAreTaggedCorrectly() throws Exception {
    WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
    WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4);
    WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);
    Map<String, Integer> outputTagMap = ImmutableMap.of("one", 1, "two", 2, "three", 3);
    // We use a real StageBundleFactory here in order to exercise the output receiver factory.
    StageBundleFactory stageBundleFactory = new StageBundleFactory() {

        private boolean once;

        @Override
        public RemoteBundle getBundle(OutputReceiverFactory receiverFactory, TimerReceiverFactory timerReceiverFactory, StateRequestHandler stateRequestHandler, BundleProgressHandler progressHandler, BundleFinalizationHandler finalizationHandler, BundleCheckpointHandler checkpointHandler) {
            return new RemoteBundle() {

                @Override
                public String getId() {
                    return "bundle-id";
                }

                @Override
                public Map<String, FnDataReceiver> getInputReceivers() {
                    return ImmutableMap.of("input", input -> {
                    /* Ignore input*/
                    });
                }

                @Override
                public Map<KV<String, String>, FnDataReceiver<Timer>> getTimerReceivers() {
                    return Collections.emptyMap();
                }

                @Override
                public void requestProgress() {
                    throw new UnsupportedOperationException();
                }

                @Override
                public void split(double fractionOfRemainder) {
                    throw new UnsupportedOperationException();
                }

                @Override
                public void close() throws Exception {
                    if (once) {
                        return;
                    }
                    // Emit all values to the runner when the bundle is closed.
                    receiverFactory.create("one").accept(three);
                    receiverFactory.create("two").accept(four);
                    receiverFactory.create("three").accept(five);
                    once = true;
                }
            };
        }

        @Override
        public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor getProcessBundleDescriptor() {
            return processBundleDescriptor;
        }

        @Override
        public InstructionRequestHandler getInstructionRequestHandler() {
            return null;
        }

        @Override
        public void close() throws Exception {
        }
    };
    // Wire the stage bundle factory into our context.
    when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
    FlinkExecutableStageFunction<Integer> function = getFunction(outputTagMap);
    function.open(new Configuration());
    if (isStateful) {
        function.reduce(Collections.emptyList(), collector);
    } else {
        function.mapPartition(Collections.emptyList(), collector);
    }
    // Ensure that the tagged values sent to the collector have the correct union tags as specified
    // in the output map.
    verify(collector).collect(new RawUnionValue(1, three));
    verify(collector).collect(new RawUnionValue(2, four));
    verify(collector).collect(new RawUnionValue(3, five));
    verifyNoMoreInteractions(collector);
}

Example 13 with RemoteBundle

use of org.apache.beam.runners.fnexecution.control.RemoteBundle in project beam by apache.

the class ProcessBundleBenchmark method testTinyBundle.

@Benchmark
// Use several threads since we expect contention during bundle processing.
@Threads(16)
public void testTinyBundle(TrivialTransform trivialTransform) throws Exception {
    Map<String, ? super Coder<WindowedValue<?>>> remoteOutputCoders = trivialTransform.descriptor.getRemoteOutputCoders();
    Map<String, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
    AtomicInteger outputValuesCount = new AtomicInteger();
    for (Entry<String, ? super Coder<WindowedValue<?>>> remoteOutputCoder : remoteOutputCoders.entrySet()) {
        outputReceivers.put(remoteOutputCoder.getKey(), RemoteOutputReceiver.of((Coder) remoteOutputCoder.getValue(), (FnDataReceiver<? super WindowedValue<?>>) (WindowedValue<?> value) -> outputValuesCount.incrementAndGet()));
    }
    try (RemoteBundle bundle = trivialTransform.processor.newBundle(outputReceivers, BundleProgressHandler.ignored())) {
        Iterables.getOnlyElement(bundle.getInputReceivers().values()).accept(valueInGlobalWindow(new byte[0]));
    }
    assertEquals(3, outputValuesCount.getAndSet(0));
}

Example 14 with RemoteBundle

use of org.apache.beam.runners.fnexecution.control.RemoteBundle in project beam by apache.

the class ProcessBundleBenchmark method testLargeBundle.

@Benchmark
// Use several threads since we expect contention during bundle processing.
@Threads(16)
public void testLargeBundle(TrivialTransform trivialTransform) throws Exception {
    Map<String, ? super Coder<WindowedValue<?>>> remoteOutputCoders = trivialTransform.descriptor.getRemoteOutputCoders();
    Map<String, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
    AtomicInteger outputValuesCount = new AtomicInteger();
    for (Entry<String, ? super Coder<WindowedValue<?>>> remoteOutputCoder : remoteOutputCoders.entrySet()) {
        outputReceivers.put(remoteOutputCoder.getKey(), RemoteOutputReceiver.of((Coder) remoteOutputCoder.getValue(), (FnDataReceiver<? super WindowedValue<?>>) (WindowedValue<?> value) -> outputValuesCount.incrementAndGet()));
    }
    try (RemoteBundle bundle = trivialTransform.processor.newBundle(outputReceivers, BundleProgressHandler.ignored())) {
        for (int i = 0; i < 1_000; i++) {
            Iterables.getOnlyElement(bundle.getInputReceivers().values()).accept(valueInGlobalWindow(new byte[0]));
        }
    }
    assertEquals(3_000, outputValuesCount.getAndSet(0));
}

Example 15 with RemoteBundle

use of org.apache.beam.runners.fnexecution.control.RemoteBundle in project beam by apache.

the class FlinkExecutableStageFunctionTest method expectedInputsAreSent.

@Test
public void expectedInputsAreSent() throws Exception {
    FlinkExecutableStageFunction<Integer> function = getFunction(Collections.emptyMap());
    function.open(new Configuration());
    @SuppressWarnings("unchecked") RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(stageBundleFactory.getBundle(any(), any(StateRequestHandler.class), any(BundleProgressHandler.class), any(BundleFinalizationHandler.class), any(BundleCheckpointHandler.class))).thenReturn(bundle);
    @SuppressWarnings("unchecked") FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
    when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));
    WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
    WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
    WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
    function.mapPartition(Arrays.asList(one, two, three), collector);
    verify(receiver).accept(one);
    verify(receiver).accept(two);
    verify(receiver).accept(three);
    verifyNoMoreInteractions(receiver);
}