Java Code Examples for org.apache.beam.sdk.coders.CoderRegistry#createDefault()

@Test
public void testDefaultWatermarkEstimatorStateAndCoder() throws Exception {
  class MockFn extends DoFn<String, String> {
    @ProcessElement
    public void processElement(
        ProcessContext c, RestrictionTracker<RestrictionWithDefaultTracker, Void> tracker) {}

    @GetInitialRestriction
    public RestrictionWithDefaultTracker getInitialRestriction(@Element String element) {
      return null;
    }
  }

  MockFn fn = mock(MockFn.class);
  DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);

  CoderRegistry coderRegistry = CoderRegistry.createDefault();
  coderRegistry.registerCoderProvider(
      CoderProviders.fromStaticMethods(
          RestrictionWithDefaultTracker.class, CoderForDefaultTracker.class));
  assertEquals(VoidCoder.of(), invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry));
  assertNull(invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>()));
}
 

@Test
public void testDefaultCoderHandlesNull() throws CannotProvideCoderException {
  Latest.LatestFn<Long> fn = new Latest.LatestFn<>();

  CoderRegistry registry = CoderRegistry.createDefault();
  TimestampedValue.TimestampedValueCoder<Long> inputCoder =
      TimestampedValue.TimestampedValueCoder.of(VarLongCoder.of());

  assertThat(
      "Default output coder should handle null values",
      fn.getDefaultOutputCoder(registry, inputCoder),
      instanceOf(NullableCoder.class));
  assertThat(
      "Default accumulator coder should handle null values",
      fn.getAccumulatorCoder(registry, inputCoder),
      instanceOf(NullableCoder.class));
}
 

@Test
public void testInferKeyCoder() {
  CoderRegistry registry = CoderRegistry.createDefault();
  assertTrue(
      LocalDeserializerProvider.of(LongDeserializer.class).getCoder(registry).getValueCoder()
          instanceof VarLongCoder);
  assertTrue(
      LocalDeserializerProvider.of(StringDeserializer.class).getCoder(registry).getValueCoder()
          instanceof StringUtf8Coder);
  assertTrue(
      LocalDeserializerProvider.of(InstantDeserializer.class).getCoder(registry).getValueCoder()
          instanceof InstantCoder);
  assertTrue(
      LocalDeserializerProvider.of(DeserializerWithInterfaces.class)
              .getCoder(registry)
              .getValueCoder()
          instanceof VarLongCoder);
}
 

@Test
public void testGetCoder() {
  String schemaRegistryUrl = "mock://my-scope-name";
  String subject = "mytopic";
  SchemaRegistryClient mockRegistryClient = mockSchemaRegistryClient(schemaRegistryUrl, subject);
  CoderRegistry coderRegistry = CoderRegistry.createDefault();

  AvroCoder coderV0 =
      (AvroCoder)
          mockDeserializerProvider(schemaRegistryUrl, subject, null).getCoder(coderRegistry);
  assertEquals(AVRO_SCHEMA, coderV0.getSchema());

  try {
    Integer version = mockRegistryClient.register(subject, AVRO_SCHEMA_V1);
    AvroCoder coderV1 =
        (AvroCoder)
            mockDeserializerProvider(schemaRegistryUrl, subject, version).getCoder(coderRegistry);
    assertEquals(AVRO_SCHEMA_V1, coderV1.getSchema());
  } catch (IOException | RestClientException e) {
    throw new RuntimeException("Unable to register schema for subject: " + subject, e);
  }
}
 

/**
 * Creates a {@link ReduceFnTester} for the given {@link WindowingStrategy} and {@link CombineFn},
 * creating a {@link TriggerStateMachine} from the {@link Trigger} in the {@link
 * WindowingStrategy}.
 */
public static <W extends BoundedWindow, AccumT, OutputT>
    ReduceFnTester<Integer, OutputT, W> combining(
        WindowingStrategy<?, W> strategy,
        CombineFn<Integer, AccumT, OutputT> combineFn,
        Coder<OutputT> outputCoder)
        throws Exception {

  CoderRegistry registry = CoderRegistry.createDefault();
  // Ensure that the CombineFn can be converted into an AppliedCombineFn
  AppliedCombineFn.withInputCoder(
      combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));

  return combining(
      strategy,
      TriggerStateMachines.stateMachineForTrigger(
          TriggerTranslation.toProto(strategy.getTrigger())),
      combineFn,
      outputCoder);
}
 

/**
 * Creates a {@link ReduceFnTester} for the given {@link WindowingStrategy}, {@link CombineFn},
 * and {@link TriggerStateMachine}, for mocking the interaction between {@link ReduceFnRunner} and
 * the {@link TriggerStateMachine}. Ignores the {@link Trigger} in the {@link WindowingStrategy}.
 */
public static <W extends BoundedWindow, AccumT, OutputT>
    ReduceFnTester<Integer, OutputT, W> combining(
        WindowingStrategy<?, W> strategy,
        TriggerStateMachine triggerStateMachine,
        CombineFn<Integer, AccumT, OutputT> combineFn,
        Coder<OutputT> outputCoder)
        throws Exception {

  CoderRegistry registry = CoderRegistry.createDefault();
  AppliedCombineFn<String, Integer, AccumT, OutputT> fn =
      AppliedCombineFn.withInputCoder(
          combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));

  return new ReduceFnTester<>(
      strategy,
      triggerStateMachine,
      SystemReduceFn.combining(StringUtf8Coder.of(), fn),
      outputCoder,
      PipelineOptionsFactory.create(),
      NullSideInputReader.empty());
}
 

public static <W extends BoundedWindow, AccumT, OutputT>
    ReduceFnTester<Integer, OutputT, W> combining(
        WindowingStrategy<?, W> strategy,
        CombineFnWithContext<Integer, AccumT, OutputT> combineFn,
        Coder<OutputT> outputCoder,
        PipelineOptions options,
        SideInputReader sideInputReader)
        throws Exception {
  CoderRegistry registry = CoderRegistry.createDefault();
  // Ensure that the CombineFn can be converted into an AppliedCombineFn
  AppliedCombineFn.withInputCoder(
      combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));

  return combining(
      strategy,
      TriggerStateMachines.stateMachineForTrigger(
          TriggerTranslation.toProto(strategy.getTrigger())),
      combineFn,
      outputCoder,
      options,
      sideInputReader);
}
 

public static <W extends BoundedWindow, AccumT, OutputT>
    ReduceFnTester<Integer, OutputT, W> combining(
        WindowingStrategy<?, W> strategy,
        TriggerStateMachine triggerStateMachine,
        CombineFnWithContext<Integer, AccumT, OutputT> combineFn,
        Coder<OutputT> outputCoder,
        PipelineOptions options,
        SideInputReader sideInputReader)
        throws Exception {
  CoderRegistry registry = CoderRegistry.createDefault();
  AppliedCombineFn<String, Integer, AccumT, OutputT> fn =
      AppliedCombineFn.withInputCoder(
          combineFn, registry, KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));

  return new ReduceFnTester<>(
      strategy,
      triggerStateMachine,
      SystemReduceFn.combining(StringUtf8Coder.of(), fn),
      outputCoder,
      options,
      sideInputReader);
}
 

/** Returns the {@link CoderRegistry} that this {@link Pipeline} uses. */
public CoderRegistry getCoderRegistry() {
  if (coderRegistry == null) {
    coderRegistry = CoderRegistry.createDefault();
  }
  return coderRegistry;
}
 

@Test
public void testInferKeyCoderFailure() throws Exception {
  cannotInferException.expect(RuntimeException.class);

  CoderRegistry registry = CoderRegistry.createDefault();
  LocalDeserializerProvider.of(NonInferableObjectDeserializer.class).getCoder(registry);
}
 

@Test
public void testCombineValuesFnCoders() throws Exception {
  CoderRegistry registry = CoderRegistry.createDefault();

  MeanInts meanInts = new MeanInts();
  CountSum countSum = new CountSum(6, 27);

  Coder<CountSum> coder =
      meanInts.getAccumulatorCoder(registry, registry.getCoder(TypeDescriptor.of(Integer.class)));

  assertEquals(
      countSum,
      CoderUtils.decodeFromByteArray(coder, CoderUtils.encodeToByteArray(coder, countSum)));
}
 

@SuppressWarnings("unchecked")
private StreamingKeyedWorkItemSideInputDoFnRunner<
        String, Integer, KV<String, Integer>, IntervalWindow>
    createRunner(DoFnRunners.OutputManager outputManager) throws Exception {
  CoderRegistry registry = CoderRegistry.createDefault();
  Coder<String> keyCoder = StringUtf8Coder.of();
  Coder<Integer> inputCoder = BigEndianIntegerCoder.of();

  AppliedCombineFn<String, Integer, ?, Integer> combineFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofIntegers(), registry, KvCoder.of(keyCoder, inputCoder));

  WindowingStrategy<Object, IntervalWindow> windowingStrategy = WindowingStrategy.of(WINDOW_FN);
  @SuppressWarnings("rawtypes")
  StreamingGroupAlsoByWindowViaWindowSetFn doFn =
      (StreamingGroupAlsoByWindowViaWindowSetFn)
          StreamingGroupAlsoByWindowsDoFns.create(
              windowingStrategy, key -> state, combineFn, keyCoder);

  DoFnRunner<KeyedWorkItem<String, Integer>, KV<String, Integer>> simpleDoFnRunner =
      new GroupAlsoByWindowFnRunner<>(
          PipelineOptionsFactory.create(),
          doFn.asDoFn(),
          mockSideInputReader,
          outputManager,
          mainOutputTag,
          stepContext);
  return new StreamingKeyedWorkItemSideInputDoFnRunner<
      String, Integer, KV<String, Integer>, IntervalWindow>(
      simpleDoFnRunner, keyCoder, sideInputFetcher, stepContext);
}
 

@SuppressWarnings({"unchecked", "rawtypes"})
@Test
public void testCombiningValueWithContextEquality() {
  CoderRegistry registry = CoderRegistry.createDefault();

  Combine.BinaryCombineIntegerFn maxFn = Max.ofIntegers();
  Combine.BinaryCombineIntegerFn minFn = Min.ofIntegers();

  Coder<int[]> accum1 = maxFn.getAccumulatorCoder(registry, VarIntCoder.of());
  Coder<int[]> accum2 = minFn.getAccumulatorCoder(registry, BigEndianIntegerCoder.of());

  StateTag<?> fooCoder1Max1 =
      StateTags.combiningValueWithContext("foo", accum1, CombineFnUtil.toFnWithContext(maxFn));
  StateTag<?> fooCoder1Max2 =
      StateTags.combiningValueWithContext("foo", accum1, CombineFnUtil.toFnWithContext(maxFn));
  StateTag<?> fooCoder1Min =
      StateTags.combiningValueWithContext("foo", accum1, CombineFnUtil.toFnWithContext(minFn));

  StateTag<?> fooCoder2Max =
      StateTags.combiningValueWithContext("foo", accum2, CombineFnUtil.toFnWithContext(maxFn));
  StateTag<?> barCoder1Max =
      StateTags.combiningValueWithContext("bar", accum1, CombineFnUtil.toFnWithContext(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
public void testSessionsCombine() throws Exception {
  TupleTag<KV<String, Long>> outputTag = new TupleTag<>();
  CombineFn<Long, ?, Long> combineFn = new SumLongs();
  CoderRegistry registry = CoderRegistry.createDefault();

  AppliedCombineFn<String, Long, ?, Long> appliedCombineFn =
      AppliedCombineFn.withInputCoder(
          combineFn, registry, KvCoder.of(StringUtf8Coder.of(), BigEndianLongCoder.of()));

  ListOutputManager outputManager = new ListOutputManager();
  DoFnRunner<KeyedWorkItem<String, Long>, KV<String, Long>> runner =
      makeRunner(
          outputTag,
          outputManager,
          WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10))),
          appliedCombineFn);

  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(0));

  runner.startBundle();

  WorkItem.Builder workItem1 = WorkItem.newBuilder();
  workItem1.setKey(ByteString.copyFromUtf8(KEY));
  workItem1.setWorkToken(WORK_TOKEN);
  InputMessageBundle.Builder messageBundle = workItem1.addMessageBundlesBuilder();
  messageBundle.setSourceComputationId(SOURCE_COMPUTATION_ID);

  Coder<Long> valueCoder = BigEndianLongCoder.of();
  addElement(messageBundle, Arrays.asList(window(0, 10)), new Instant(0), valueCoder, 1L);
  addElement(messageBundle, Arrays.asList(window(5, 15)), new Instant(5), valueCoder, 2L);
  addElement(messageBundle, Arrays.asList(window(15, 25)), new Instant(15), valueCoder, 3L);
  addElement(messageBundle, Arrays.asList(window(3, 13)), new Instant(3), valueCoder, 4L);

  runner.processElement(createValue(workItem1, valueCoder));

  runner.finishBundle();
  runner.startBundle();

  WorkItem.Builder workItem2 = WorkItem.newBuilder();
  workItem2.setKey(ByteString.copyFromUtf8(KEY));
  workItem2.setWorkToken(WORK_TOKEN);
  // Note that the WATERMARK timer for Instant(9) will have been deleted by
  // ReduceFnRunner when window(0, 10) was merged away.
  addTimer(workItem2, window(0, 15), new Instant(14), Timer.Type.WATERMARK);
  addTimer(workItem2, window(15, 25), new Instant(24), Timer.Type.WATERMARK);
  when(mockTimerInternals.currentInputWatermarkTime()).thenReturn(new Instant(25));

  runner.processElement(createValue(workItem2, valueCoder));

  runner.finishBundle();

  List<WindowedValue<KV<String, Long>>> result = outputManager.getOutput(outputTag);

  assertThat(result.size(), equalTo(2));

  assertThat(
      result,
      containsInAnyOrder(
          WindowMatchers.isSingleWindowedValue(
              isKv(equalTo(KEY), equalTo(7L)),
              equalTo(window(0, 15).maxTimestamp()),
              equalTo(window(0, 15))),
          WindowMatchers.isSingleWindowedValue(
              isKv(equalTo(KEY), equalTo(3L)),
              equalTo(window(15, 25).maxTimestamp()),
              equalTo(window(15, 25)))));
}