Java Code Examples for org.apache.beam.sdk.coders.VarIntCoder#of()

@Test
public void merge() {
  double eps = 0.01;
  double conf = 0.8;
  long nOccurrences = 2L;
  int size = 3;

  List<Sketch<Integer>> sketches = new ArrayList<>();
  Coder<Integer> coder = VarIntCoder.of();

  // n sketches each containing [0, 1, 2]
  for (int i = 0; i < nOccurrences; i++) {
    Sketch<Integer> sketch = Sketch.create(eps, conf);
    for (int j = 0; j < size; j++) {
      sketch.add(j, coder);
    }
    sketches.add(sketch);
  }

  CountMinSketchFn<Integer> fn = CountMinSketchFn.create(coder).withAccuracy(eps, conf);
  Sketch<Integer> merged = fn.mergeAccumulators(sketches);
  for (int i = 0; i < size; i++) {
    assertEquals(nOccurrences, merged.estimateCount(i, coder));
  }
}
 

private static Coder resolveCoder(Class deserializer) {
  for (Method method : deserializer.getDeclaredMethods()) {
    if (method.getName().equals("deserialize")) {
      Class<?> returnType = method.getReturnType();
      if (returnType.equals(Object.class)) {
        continue;
      }
      if (returnType.equals(byte[].class)) {
        return ByteArrayCoder.of();
      } else if (returnType.equals(Integer.class)) {
        return VarIntCoder.of();
      } else if (returnType.equals(Long.class)) {
        return VarLongCoder.of();
      } else {
        throw new RuntimeException("Couldn't infer Coder from " + deserializer);
      }
    }
  }
  throw new RuntimeException("Couldn't resolve coder for Deserializer: " + deserializer);
}
 

@Test
public void testCompoundCoder() throws Exception {
  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
  Coder<?> coder = VarIntCoder.of();
  Coder<?> compoundCoder = NullableCoder.of(coder);
  String compoundCoderId = sdkComponents.registerCoder(compoundCoder);
  String coderId = sdkComponents.registerCoder(coder);

  RehydratedComponents rehydratedComponents =
      RehydratedComponents.forComponents(sdkComponents.toComponents());

  Coder<?> rehydratedCoder = rehydratedComponents.getCoder(coderId);
  Coder<?> rehydratedCompoundCoder = rehydratedComponents.getCoder(compoundCoderId);

  assertThat(rehydratedCoder, equalTo((Coder) coder));
  assertThat(rehydratedCompoundCoder, equalTo((Coder) compoundCoder));

  assertThat(rehydratedComponents.getCoder(coderId), theInstance((Coder) rehydratedCoder));
  assertThat(
      rehydratedComponents.getCoder(compoundCoderId),
      theInstance((Coder) rehydratedCompoundCoder));
}
 

@Test
public void testCoder() throws Exception {
  Sketch<Integer> cMSketch = Sketch.create(0.01, 0.8);
  Coder<Integer> coder = VarIntCoder.of();
  for (int i = 0; i < 3; i++) {
    cMSketch.add(i, coder);
  }

  CoderProperties.coderDecodeEncodeEqual(new SketchFrequencies.CountMinSketchCoder<>(), cMSketch);
}
 

@Test
public void testSourceGetOutputCoderReturnsConstructorCoder() throws Exception {
  Coder<Integer> coder = VarIntCoder.of();
  CreateSource<Integer> source =
      CreateSource.fromIterable(ImmutableList.of(1, 2, 3, 4, 5, 6, 7, 8), coder);

  Coder<Integer> defaultCoder = source.getOutputCoder();
  assertThat(defaultCoder, equalTo(coder));
}
 

@Test
public void testSimpleCoder() throws Exception {
  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
  Coder<?> coder = VarIntCoder.of();
  String id = sdkComponents.registerCoder(coder);
  RehydratedComponents rehydratedComponents =
      RehydratedComponents.forComponents(sdkComponents.toComponents());

  Coder<?> rehydratedCoder = rehydratedComponents.getCoder(id);
  assertThat(rehydratedCoder, equalTo((Coder) coder));
  assertThat(rehydratedComponents.getCoder(id), theInstance((Coder) rehydratedCoder));
}
 

@Before
public void setup() {
  MockitoAnnotations.initMocks(this);
  PCollection<Integer> created = p.apply(Create.of(1, 2, 3));
  singletonView =
      created
          .apply(Window.into(new IdentitySideInputWindowFn()))
          .apply(Sum.integersGlobally().asSingletonView());

  underlying = new TestDoFnRunner<>(VarIntCoder.of());

  DoFn<KV<String, Integer>, Integer> fn = new MyDoFn();

  MockitoAnnotations.initMocks(this);
  when(mockStepContext.timerInternals()).thenReturn(timerInternals);

  stateInternals = new InMemoryStateInternals<>("hello");
  timerInternals = new InMemoryTimerInternals();

  when(mockStepContext.stateInternals()).thenReturn((StateInternals) stateInternals);
  when(mockStepContext.timerInternals()).thenReturn(timerInternals);

  statefulRunner =
      DoFnRunners.defaultStatefulDoFnRunner(
          fn,
          KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()),
          getDoFnRunner(fn),
          asStepContext(stateInternals, timerInternals),
          WINDOWING_STRATEGY,
          new StatefulDoFnRunner.TimeInternalsCleanupTimer(timerInternals, WINDOWING_STRATEGY),
          new StatefulDoFnRunner.StateInternalsStateCleaner<>(
              fn, stateInternals, (Coder) WINDOWING_STRATEGY.getWindowFn().windowCoder()));
}
 

@SuppressWarnings({"unchecked", "rawtypes"})
@Test
public void testCombiningValueEquality() {
  Combine.BinaryCombineIntegerFn maxFn = Max.ofIntegers();
  Coder<Integer> input1 = VarIntCoder.of();
  Coder<Integer> input2 = BigEndianIntegerCoder.of();
  Combine.BinaryCombineIntegerFn minFn = Min.ofIntegers();

  StateTag<?> fooCoder1Max1 = StateTags.combiningValueFromInputInternal("foo", input1, maxFn);
  StateTag<?> fooCoder1Max2 = StateTags.combiningValueFromInputInternal("foo", input1, maxFn);
  StateTag<?> fooCoder1Min = StateTags.combiningValueFromInputInternal("foo", input1, minFn);

  StateTag<?> fooCoder2Max = StateTags.combiningValueFromInputInternal("foo", input2, maxFn);
  StateTag<?> barCoder1Max = StateTags.combiningValueFromInputInternal("bar", input1, maxFn);

  // Same name, coder and combineFn
  assertEquals(fooCoder1Max1, fooCoder1Max2);
  assertEquals(
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Max1),
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Max2));

  // Different combineFn, but we treat them as equal since we only serialize the bits.
  assertEquals(fooCoder1Max1, fooCoder1Min);
  assertEquals(
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Max1),
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Min));

  // Different input coder coder.
  assertEquals(fooCoder1Max1, fooCoder2Max);
  assertEquals(
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Max1),
      StateTags.convertToBagTagInternal((StateTag) fooCoder2Max));

  // These StateTags have different IDs.
  assertNotEquals(fooCoder1Max1, barCoder1Max);
  assertNotEquals(
      StateTags.convertToBagTagInternal((StateTag) fooCoder1Max1),
      StateTags.convertToBagTagInternal((StateTag) barCoder1Max));
}
 

@Test
@SuppressWarnings("unchecked")
public void testEnsureStateCleanupWithKeyedInput() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VarIntCoder.of());
  VarIntCoder keyCoder = VarIntCoder.of();
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(
          mainOutput,
          Collections.emptyList(),
          outputManagerFactory,
          WindowingStrategy.globalDefault(),
          keyCoder,
          WindowedValue.getFullCoder(keyCoder, GlobalWindow.Coder.INSTANCE));

  KeyedOneInputStreamOperatorTestHarness<Integer, WindowedValue<Integer>, WindowedValue<Integer>>
      testHarness =
          new KeyedOneInputStreamOperatorTestHarness(
              operator, val -> val, new CoderTypeInformation<>(keyCoder));

  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(bundle.getInputReceivers())
      .thenReturn(
          ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder()
              .put("input", Mockito.mock(FnDataReceiver.class))
              .build());
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  testHarness.open();

  Object doFnRunner = Whitebox.getInternalState(operator, "doFnRunner");
  assertThat(doFnRunner, instanceOf(DoFnRunnerWithMetricsUpdate.class));

  // There should be a StatefulDoFnRunner installed which takes care of clearing state
  Object statefulDoFnRunner = Whitebox.getInternalState(doFnRunner, "delegate");
  assertThat(statefulDoFnRunner, instanceOf(StatefulDoFnRunner.class));
}
 

@Override
public Coder<Integer> getStateCoder() {
  return VarIntCoder.of();
}
 

@Test
public void testSplitBySameKey() {
  VarIntCoder coder = VarIntCoder.of();
  WindowedValue.WindowedValueCoder<Integer> wvCoder =
      WindowedValue.FullWindowedValueCoder.of(coder, GlobalWindow.Coder.INSTANCE);
  Instant now = Instant.now();
  List<GlobalWindow> window = Arrays.asList(GlobalWindow.INSTANCE);
  PaneInfo paneInfo = PaneInfo.NO_FIRING;
  List<Tuple2<ByteArray, byte[]>> firstKey =
      Arrays.asList(
          new Tuple2(
              new ByteArray(CoderHelpers.toByteArrayWithTs(1, coder, now)),
              CoderHelpers.toByteArray(WindowedValue.of(1, now, window, paneInfo), wvCoder)),
          new Tuple2(
              new ByteArray(CoderHelpers.toByteArrayWithTs(1, coder, now.plus(1))),
              CoderHelpers.toByteArray(
                  WindowedValue.of(2, now.plus(1), window, paneInfo), wvCoder)));

  List<Tuple2<ByteArray, byte[]>> secondKey =
      Arrays.asList(
          new Tuple2(
              new ByteArray(CoderHelpers.toByteArrayWithTs(2, coder, now)),
              CoderHelpers.toByteArray(WindowedValue.of(3, now, window, paneInfo), wvCoder)),
          new Tuple2(
              new ByteArray(CoderHelpers.toByteArrayWithTs(2, coder, now.plus(2))),
              CoderHelpers.toByteArray(
                  WindowedValue.of(4, now.plus(2), window, paneInfo), wvCoder)));

  Iterable<Tuple2<ByteArray, byte[]>> concat = Iterables.concat(firstKey, secondKey);
  Iterator<Iterator<WindowedValue<KV<Integer, Integer>>>> keySplit;
  keySplit = TransformTranslator.splitBySameKey(concat.iterator(), coder, wvCoder);

  for (int i = 0; i < 2; i++) {
    Iterator<WindowedValue<KV<Integer, Integer>>> iter = keySplit.next();
    List<WindowedValue<KV<Integer, Integer>>> list = new ArrayList<>();
    Iterators.addAll(list, iter);
    if (i == 0) {
      // first key
      assertEquals(
          Arrays.asList(
              WindowedValue.of(KV.of(1, 1), now, window, paneInfo),
              WindowedValue.of(KV.of(1, 2), now.plus(1), window, paneInfo)),
          list);
    } else {
      // second key
      assertEquals(
          Arrays.asList(
              WindowedValue.of(KV.of(2, 3), now, window, paneInfo),
              WindowedValue.of(KV.of(2, 4), now.plus(2), window, paneInfo)),
          list);
    }
  }
}
 

@Test
public void testLateDroppingForStatefulFn() throws Exception {

  WindowingStrategy<Object, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(new Duration(10)));

  DoFn<Integer, String> fn =
      new DoFn<Integer, String>() {

        @StateId("state")
        private final StateSpec<ValueState<String>> stateSpec =
            StateSpecs.value(StringUtf8Coder.of());

        @ProcessElement
        public void processElement(ProcessContext context) {
          context.output(context.element().toString());
        }
      };

  VarIntCoder keyCoder = VarIntCoder.of();
  Coder<WindowedValue<Integer>> inputCoder =
      WindowedValue.getFullCoder(keyCoder, windowingStrategy.getWindowFn().windowCoder());
  Coder<WindowedValue<String>> outputCoder =
      WindowedValue.getFullCoder(
          StringUtf8Coder.of(), windowingStrategy.getWindowFn().windowCoder());

  KeySelector<WindowedValue<Integer>, ByteBuffer> keySelector =
      e -> FlinkKeyUtils.encodeKey(e.getValue(), keyCoder);

  TupleTag<String> outputTag = new TupleTag<>("main-output");

  DoFnOperator<Integer, String> doFnOperator =
      new DoFnOperator<>(
          fn,
          "stepName",
          inputCoder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, outputCoder),
          windowingStrategy,
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          PipelineOptionsFactory.as(FlinkPipelineOptions.class),
          keyCoder, /* key coder */
          keySelector,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
      new KeyedOneInputStreamOperatorTestHarness<>(
          doFnOperator,
          keySelector,
          new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of()));

  testHarness.open();

  testHarness.processWatermark(0);

  IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10));

  // this should not be late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(WindowedValue.of("13", new Instant(0), window1, PaneInfo.NO_FIRING)));

  testHarness.getOutput().clear();

  testHarness.processWatermark(9);

  // this should still not be considered late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(
      stripStreamRecordFromWindowedValue(testHarness.getOutput()),
      contains(WindowedValue.of("17", new Instant(0), window1, PaneInfo.NO_FIRING)));

  testHarness.getOutput().clear();

  testHarness.processWatermark(10);

  // this should now be considered late
  testHarness.processElement(
      new StreamRecord<>(WindowedValue.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));

  assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());

  testHarness.close();
}