Examples with ImmutableMap org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap used on opensource projects

Example 6 with ImmutableMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap in project beam by apache.

the class IsmSideInputReaderTest method testListInWindow.

@Test
public void testListInWindow() throws Exception {
    Coder<WindowedValue<Long>> valueCoder = WindowedValue.getFullCoder(VarLongCoder.of(), INTERVAL_WINDOW_CODER);
    IsmRecordCoder<WindowedValue<Long>> ismCoder = IsmRecordCoder.of(1, 0, ImmutableList.of(INTERVAL_WINDOW_CODER, BigEndianLongCoder.of()), valueCoder);
    final List<KV<Long, WindowedValue<Long>>> firstElements = Arrays.asList(KV.of(0L, valueInIntervalWindow(12, 10)), KV.of(1L, valueInIntervalWindow(22, 10)), KV.of(2L, valueInIntervalWindow(32, 10)));
    final List<KV<Long, WindowedValue<Long>>> secondElements = Arrays.asList(KV.of(0L, valueInIntervalWindow(42, 20)), KV.of(1L, valueInIntervalWindow(52, 20)), KV.of(2L, valueInIntervalWindow(62, 20)));
    final List<KV<Long, WindowedValue<Long>>> thirdElements = Arrays.asList(KV.of(0L, valueInIntervalWindow(42L, 30)), KV.of(1L, valueInIntervalWindow(52L, 30)), KV.of(2L, valueInIntervalWindow(62L, 30)));
    final PCollectionView<List<Long>> view = Pipeline.create().apply(Create.empty(VarLongCoder.of())).apply(Window.into(FixedWindows.of(Duration.millis(10)))).apply(View.asList());
    Source sourceA = initInputFile(fromKvsForList(concat(firstElements, secondElements)), ismCoder);
    Source sourceB = initInputFile(fromKvsForList(thirdElements), ismCoder);
    final IsmSideInputReader reader = sideInputReader(view.getTagInternal().getId(), sourceA, sourceB);
    List<Callable<Map<BoundedWindow, List<Long>>>> tasks = new ArrayList<>();
    for (int i = 0; i < NUM_THREADS; ++i) {
        tasks.add(() -> {
            // Store a strong reference to the returned value so that the logical reference
            // cache is not cleared for this test.
            List<Long> firstValues = reader.get(view, intervalWindow(10));
            List<Long> secondValues = reader.get(view, intervalWindow(20));
            List<Long> thirdValues = reader.get(view, intervalWindow(30));
            verifyList(toValueList(firstElements), firstValues);
            verifyList(toValueList(secondElements), secondValues);
            verifyList(toValueList(thirdElements), thirdValues);
            // Assert that the same value reference was returned showing that it was cached.
            assertSame(firstValues, reader.get(view, intervalWindow(10)));
            assertSame(secondValues, reader.get(view, intervalWindow(20)));
            assertSame(thirdValues, reader.get(view, intervalWindow(30)));
            // Also verify when requesting a window that is not part of the side input
            assertEquals(Collections.EMPTY_LIST, reader.get(view, intervalWindow(40)));
            return ImmutableMap.<BoundedWindow, List<Long>>of(intervalWindow(10), firstValues, intervalWindow(20), secondValues, intervalWindow(30), thirdValues);
        });
    }
    List<Future<Map<BoundedWindow, List<Long>>>> results = pipelineOptions.getExecutorService().invokeAll(tasks);
    Map<BoundedWindow, List<Long>> value = results.get(0).get();
    // Assert that all threads got back the same reference
    for (Future<Map<BoundedWindow, List<Long>>> result : results) {
        assertEquals(value, result.get());
        for (Map.Entry<BoundedWindow, List<Long>> entry : result.get().entrySet()) {
            assertSame(value.get(entry.getKey()), entry.getValue());
        }
    }
}

Example 7 with ImmutableMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap in project beam by apache.

the class IsmSideInputReaderTest method testMapInWindow.

@Test
public void testMapInWindow() throws Exception {
    // Note that we purposely use byte[]s as keys to force structural equality testing
    // versus using java equality testing.
    Coder<WindowedValue<Long>> valueCoder = WindowedValue.getFullCoder(VarLongCoder.of(), INTERVAL_WINDOW_CODER);
    final ListMultimap<byte[], WindowedValue<Long>> firstWindow = ImmutableListMultimap.<byte[], WindowedValue<Long>>builder().put(new byte[] { 0x00 }, valueInIntervalWindow(12L, 10)).put(new byte[] { 0x01 }, valueInIntervalWindow(22L, 10)).put(new byte[] { 0x02 }, valueInIntervalWindow(32L, 10)).build();
    final ListMultimap<byte[], WindowedValue<Long>> secondWindow = ImmutableListMultimap.<byte[], WindowedValue<Long>>builder().put(new byte[] { 0x00 }, valueInIntervalWindow(42L, 20)).put(new byte[] { 0x03 }, valueInIntervalWindow(52L, 20)).put(new byte[] { 0x02 }, valueInIntervalWindow(62L, 20)).build();
    final ListMultimap<byte[], WindowedValue<Long>> thirdWindow = ImmutableListMultimap.<byte[], WindowedValue<Long>>builder().put(new byte[] { 0x02 }, valueInIntervalWindow(72L, 30)).put(new byte[] { 0x04 }, valueInIntervalWindow(82L, 30)).put(new byte[] { 0x05 }, valueInIntervalWindow(92L, 30)).build();
    final PCollectionView<Map<byte[], Long>> view = Pipeline.create().apply(Create.empty(KvCoder.of(ByteArrayCoder.of(), VarLongCoder.of()))).apply(Window.into(FixedWindows.of(Duration.millis(10)))).apply(View.asMap());
    IsmRecordCoder<WindowedValue<Long>> ismCoder = IsmRecordCoder.of(1, 2, ImmutableList.of(MetadataKeyCoder.of(ByteArrayCoder.of()), INTERVAL_WINDOW_CODER, BigEndianLongCoder.of()), valueCoder);
    Multimap<Integer, IsmRecord<WindowedValue<Long>>> elementsPerShard = forMap(ismCoder, firstWindow);
    elementsPerShard.putAll(forMap(ismCoder, secondWindow));
    elementsPerShard.putAll(forMap(ismCoder, thirdWindow));
    List<IsmRecord<WindowedValue<Long>>> firstElements = new ArrayList<>();
    List<IsmRecord<WindowedValue<Long>>> secondElements = new ArrayList<>();
    for (Map.Entry<Integer, Collection<IsmRecord<WindowedValue<Long>>>> entry : elementsPerShard.asMap().entrySet()) {
        if (entry.getKey() % 2 == 0) {
            firstElements.addAll(entry.getValue());
        } else {
            secondElements.addAll(entry.getValue());
        }
    }
    // Ensure that each file will have some records.
    checkState(!firstElements.isEmpty());
    checkState(!secondElements.isEmpty());
    Source sourceA = initInputFile(firstElements, ismCoder);
    Source sourceB = initInputFile(secondElements, ismCoder);
    List<IsmRecord<WindowedValue<Long>>> firstWindowMapMetadata = forMapMetadata(ByteArrayCoder.of(), firstWindow.keySet(), intervalWindow(10));
    List<IsmRecord<WindowedValue<Long>>> secondWindowMapMetadata = forMapMetadata(ByteArrayCoder.of(), secondWindow.keySet(), intervalWindow(20));
    List<IsmRecord<WindowedValue<Long>>> thirdWindowMapMetadata = forMapMetadata(ByteArrayCoder.of(), thirdWindow.keySet(), intervalWindow(30));
    Source sourceMetaA = initInputFile(firstWindowMapMetadata, ismCoder);
    Source sourceMetaB = initInputFile(concat(secondWindowMapMetadata, thirdWindowMapMetadata), ismCoder);
    final IsmSideInputReader reader = sideInputReader(view.getTagInternal().getId(), sourceA, sourceB, sourceMetaA, sourceMetaB);
    List<Callable<Map<BoundedWindow, Map<byte[], Long>>>> tasks = new ArrayList<>();
    for (int i = 0; i < NUM_THREADS; ++i) {
        tasks.add(() -> {
            // Store a strong reference to the returned value so that the logical reference
            // cache is not cleared for this test.
            Map<byte[], Long> firstValues = reader.get(view, intervalWindow(10));
            Map<byte[], Long> secondValues = reader.get(view, intervalWindow(20));
            Map<byte[], Long> thirdValues = reader.get(view, intervalWindow(30));
            verifyMap(Maps.transformValues(firstWindow.asMap(), new TransformForMap<Long>()), firstValues, new ComparatorForMap<Long>());
            verifyMap(Maps.transformValues(secondWindow.asMap(), new TransformForMap<Long>()), secondValues, new ComparatorForMap<Long>());
            verifyMap(Maps.transformValues(thirdWindow.asMap(), new TransformForMap<Long>()), thirdValues, new ComparatorForMap<Long>());
            // Assert that the same value reference was returned showing that it was cached.
            assertSame(firstValues, reader.get(view, intervalWindow(10)));
            assertSame(secondValues, reader.get(view, intervalWindow(20)));
            assertSame(thirdValues, reader.get(view, intervalWindow(30)));
            // Also verify when requesting a window that is not part of the side input
            assertEquals(Collections.EMPTY_MAP, reader.get(view, intervalWindow(40)));
            return ImmutableMap.<BoundedWindow, Map<byte[], Long>>of(intervalWindow(10), firstValues, intervalWindow(20), secondValues, intervalWindow(30), thirdValues);
        });
    }
    List<Future<Map<BoundedWindow, Map<byte[], Long>>>> results = pipelineOptions.getExecutorService().invokeAll(tasks);
    Map<BoundedWindow, Map<byte[], Long>> value = results.get(0).get();
    // Assert that all threads got back the same reference
    for (Future<Map<BoundedWindow, Map<byte[], Long>>> result : results) {
        assertEquals(value, result.get());
        for (Map.Entry<BoundedWindow, Map<byte[], Long>> entry : result.get().entrySet()) {
            assertSame(value.get(entry.getKey()), entry.getValue());
        }
    }
}

Example 8 with ImmutableMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap in project beam by apache.

the class DoFnOperatorTest method keyedParDoSideInputCheckpointing.

@Test
public void keyedParDoSideInputCheckpointing() throws Exception {
    sideInputCheckpointing(() -> {
        StringUtf8Coder keyCoder = StringUtf8Coder.of();
        Coder<WindowedValue<String>> coder = WindowedValue.getFullCoder(keyCoder, IntervalWindow.getCoder());
        TupleTag<String> outputTag = new TupleTag<>("main-output");
        KeySelector<WindowedValue<String>, ByteBuffer> keySelector = e -> FlinkKeyUtils.encodeKey(e.getValue(), keyCoder);
        ImmutableMap<Integer, PCollectionView<?>> sideInputMapping = ImmutableMap.<Integer, PCollectionView<?>>builder().put(1, view1).put(2, view2).build();
        DoFnOperator<String, String> doFnOperator = new DoFnOperator<>(new IdentityDoFn<>(), "stepName", coder, Collections.emptyMap(), outputTag, Collections.emptyList(), new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())), WindowingStrategy.of(FixedWindows.of(Duration.millis(100))), sideInputMapping, /* side-input mapping */
        ImmutableList.of(view1, view2), /* side inputs */
        FlinkPipelineOptions.defaults(), keyCoder, keySelector, DoFnSchemaInformation.create(), Collections.emptyMap());
        return new KeyedTwoInputStreamOperatorTestHarness<>(doFnOperator, keySelector, // we use a dummy key for the second input since it is considered to be broadcast
        null, new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of(), FlinkPipelineOptions.defaults()));
    });
}

Example 9 with ImmutableMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap 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 10 with ImmutableMap

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap in project beam by apache.

the class ProcessBundleHandler method monitoringData.

private ImmutableMap<String, ByteString> monitoringData(BundleProcessor bundleProcessor) throws Exception {
    ImmutableMap.Builder<String, ByteString> result = ImmutableMap.builder();
    // Get start bundle Execution Time Metrics.
    result.putAll(bundleProcessor.getStartFunctionRegistry().getExecutionTimeMonitoringData(shortIds));
    // Get process bundle Execution Time Metrics.
    result.putAll(bundleProcessor.getpCollectionConsumerRegistry().getExecutionTimeMonitoringData(shortIds));
    // Get finish bundle Execution Time Metrics.
    result.putAll(bundleProcessor.getFinishFunctionRegistry().getExecutionTimeMonitoringData(shortIds));
    // Extract MonitoringInfos that come from the metrics container registry.
    result.putAll(bundleProcessor.getMetricsContainerRegistry().getMonitoringData(shortIds));
    // Add any additional monitoring infos that the "runners" report explicitly.
    for (ProgressRequestCallback progressRequestCallback : bundleProcessor.getProgressRequestCallbacks()) {
        // TODO(BEAM-6597): Plumb reporting monitoring infos using the short id system upstream.
        for (MetricsApi.MonitoringInfo monitoringInfo : progressRequestCallback.getMonitoringInfos()) {
            ByteString payload = monitoringInfo.getPayload();
            String shortId = shortIds.getOrCreateShortId(monitoringInfo.toBuilder().clearPayload().build());
            result.put(shortId, payload);
        }
    }
    return result.build();
}