Search in sources :

Example 1 with TransformedMap

use of org.apache.beam.runners.dataflow.BatchViewOverrides.TransformedMap in project beam by apache.

the class BatchViewOverridesTest method testToMultimapDoFn.

@Test
public void testToMultimapDoFn() throws Exception {
    Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();
    DoFnTester<KV<Integer, Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>>, IsmRecord<WindowedValue<TransformedMap<Long, Iterable<WindowedValue<Long>>, Iterable<Long>>>>> doFnTester = DoFnTester.of(new BatchViewOverrides.BatchViewAsMultimap.ToMultimapDoFn<Long, Long, IntervalWindow>(windowCoder));
    IntervalWindow windowA = new IntervalWindow(new Instant(0), new Instant(10));
    IntervalWindow windowB = new IntervalWindow(new Instant(10), new Instant(20));
    IntervalWindow windowC = new IntervalWindow(new Instant(20), new Instant(30));
    Iterable<KV<Integer, Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>>> inputElements = ImmutableList.of(KV.of(1, (Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>) ImmutableList.of(KV.of(windowA, WindowedValue.of(KV.of(1L, 11L), new Instant(3), windowA, PaneInfo.NO_FIRING)), KV.of(windowA, WindowedValue.of(KV.of(1L, 12L), new Instant(5), windowA, PaneInfo.NO_FIRING)), KV.of(windowA, WindowedValue.of(KV.of(2L, 21L), new Instant(7), windowA, PaneInfo.NO_FIRING)), KV.of(windowB, WindowedValue.of(KV.of(2L, 21L), new Instant(13), windowB, PaneInfo.NO_FIRING)), KV.of(windowB, WindowedValue.of(KV.of(3L, 31L), new Instant(15), windowB, PaneInfo.NO_FIRING)))), KV.of(2, (Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>) ImmutableList.of(KV.of(windowC, WindowedValue.of(KV.of(4L, 41L), new Instant(25), windowC, PaneInfo.NO_FIRING)))));
    // The order of the output elements is important relative to processing order
    List<IsmRecord<WindowedValue<TransformedMap<Long, Iterable<WindowedValue<Long>>, Iterable<Long>>>>> output = doFnTester.processBundle(inputElements);
    assertEquals(3, output.size());
    Map<Long, Iterable<Long>> outputMap;
    outputMap = output.get(0).getValue().getValue();
    assertEquals(2, outputMap.size());
    assertThat(outputMap.get(1L), containsInAnyOrder(11L, 12L));
    assertThat(outputMap.get(2L), containsInAnyOrder(21L));
    outputMap = output.get(1).getValue().getValue();
    assertEquals(2, outputMap.size());
    assertThat(outputMap.get(2L), containsInAnyOrder(21L));
    assertThat(outputMap.get(3L), containsInAnyOrder(31L));
    outputMap = output.get(2).getValue().getValue();
    assertEquals(1, outputMap.size());
    assertThat(outputMap.get(4L), containsInAnyOrder(41L));
}
Also used : Instant(org.joda.time.Instant) IsmRecord(org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord) KV(org.apache.beam.sdk.values.KV) TransformedMap(org.apache.beam.runners.dataflow.BatchViewOverrides.TransformedMap) WindowedValue(org.apache.beam.sdk.util.WindowedValue) IntervalWindow(org.apache.beam.sdk.transforms.windowing.IntervalWindow) Test(org.junit.Test)

Example 2 with TransformedMap

use of org.apache.beam.runners.dataflow.BatchViewOverrides.TransformedMap in project beam by apache.

the class BatchViewOverridesTest method testToMapDoFn.

@Test
public void testToMapDoFn() throws Exception {
    Coder<IntervalWindow> windowCoder = IntervalWindow.getCoder();
    DoFnTester<KV<Integer, Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>>, IsmRecord<WindowedValue<TransformedMap<Long, WindowedValue<Long>, Long>>>> doFnTester = DoFnTester.of(new BatchViewOverrides.BatchViewAsMap.ToMapDoFn<Long, Long, IntervalWindow>(windowCoder));
    IntervalWindow windowA = new IntervalWindow(new Instant(0), new Instant(10));
    IntervalWindow windowB = new IntervalWindow(new Instant(10), new Instant(20));
    IntervalWindow windowC = new IntervalWindow(new Instant(20), new Instant(30));
    Iterable<KV<Integer, Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>>> inputElements = ImmutableList.of(KV.of(1, (Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>) ImmutableList.of(KV.of(windowA, WindowedValue.of(KV.of(1L, 11L), new Instant(3), windowA, PaneInfo.NO_FIRING)), KV.of(windowA, WindowedValue.of(KV.of(2L, 21L), new Instant(7), windowA, PaneInfo.NO_FIRING)), KV.of(windowB, WindowedValue.of(KV.of(2L, 21L), new Instant(13), windowB, PaneInfo.NO_FIRING)), KV.of(windowB, WindowedValue.of(KV.of(3L, 31L), new Instant(15), windowB, PaneInfo.NO_FIRING)))), KV.of(2, (Iterable<KV<IntervalWindow, WindowedValue<KV<Long, Long>>>>) ImmutableList.of(KV.of(windowC, WindowedValue.of(KV.of(4L, 41L), new Instant(25), windowC, PaneInfo.NO_FIRING)))));
    // The order of the output elements is important relative to processing order
    List<IsmRecord<WindowedValue<TransformedMap<Long, WindowedValue<Long>, Long>>>> output = doFnTester.processBundle(inputElements);
    assertEquals(3, output.size());
    Map<Long, Long> outputMap;
    outputMap = output.get(0).getValue().getValue();
    assertEquals(2, outputMap.size());
    assertEquals(ImmutableMap.of(1L, 11L, 2L, 21L), outputMap);
    outputMap = output.get(1).getValue().getValue();
    assertEquals(2, outputMap.size());
    assertEquals(ImmutableMap.of(2L, 21L, 3L, 31L), outputMap);
    outputMap = output.get(2).getValue().getValue();
    assertEquals(1, outputMap.size());
    assertEquals(ImmutableMap.of(4L, 41L), outputMap);
}
Also used : Instant(org.joda.time.Instant) IsmRecord(org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord) KV(org.apache.beam.sdk.values.KV) TransformedMap(org.apache.beam.runners.dataflow.BatchViewOverrides.TransformedMap) WindowedValue(org.apache.beam.sdk.util.WindowedValue) IntervalWindow(org.apache.beam.sdk.transforms.windowing.IntervalWindow) Test(org.junit.Test)

Aggregations

TransformedMap (org.apache.beam.runners.dataflow.BatchViewOverrides.TransformedMap)2 IsmRecord (org.apache.beam.runners.dataflow.internal.IsmFormat.IsmRecord)2 IntervalWindow (org.apache.beam.sdk.transforms.windowing.IntervalWindow)2 WindowedValue (org.apache.beam.sdk.util.WindowedValue)2 KV (org.apache.beam.sdk.values.KV)2 Instant (org.joda.time.Instant)2 Test (org.junit.Test)2