org.apache.beam.sdk.coders.CoderRegistry Java Examples

Coder<DestinationT> getDestinationCoderWithDefault(CoderRegistry registry)
    throws CannotProvideCoderException {
  Coder<DestinationT> destinationCoder = getDestinationCoder();
  if (destinationCoder != null) {
    return destinationCoder;
  }
  // If dynamicDestinations doesn't provide a coder, try to find it in the coder registry.
  TypeDescriptor<DestinationT> descriptor =
      extractFromTypeParameters(
          this,
          DynamicDestinations.class,
          new TypeDescriptors.TypeVariableExtractor<
              DynamicDestinations<T, DestinationT>, DestinationT>() {});
  try {
    return registry.getCoder(descriptor);
  } catch (CannotProvideCoderException e) {
    throw new CannotProvideCoderException(
        "Failed to infer coder for DestinationT from type "
            + descriptor
            + ", please provide it explicitly by overriding getDestinationCoder()",
        e);
  }
}
 

@Override
public PCollection<KV<K, V>> expand(PBegin input) {
  validateTransform();
  // Get the key and value coders based on the key and value classes.
  CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
  Coder<K> keyCoder = getDefaultCoder(getKeyTypeDescriptor(), coderRegistry);
  Coder<V> valueCoder = getDefaultCoder(getValueTypeDescriptor(), coderRegistry);
  HadoopInputFormatBoundedSource<K, V> source =
      new HadoopInputFormatBoundedSource<>(
          getConfiguration(),
          keyCoder,
          valueCoder,
          getKeyTranslationFunction(),
          getValueTranslationFunction());
  return input.getPipeline().apply(org.apache.beam.sdk.io.Read.from(source));
}
 

private SerializableCoderReplacement(final ClassVisitor delegate, final String plugin, final Class<?> clazz) {
    super(ASM7, delegate);
    this.plugin = plugin;

    Type accumulatorType = null;
    if (Combine.CombineFn.class.isAssignableFrom(clazz)) { // not the best impl but user code should handle it
        try {
            if (clazz
                    .getMethod("getAccumulatorCoder", CoderRegistry.class, Coder.class)
                    .getDeclaringClass() != clazz) {
                accumulatorType = Type.getType(clazz.getMethod("createAccumulator").getReturnType());
            }
        } catch (final NoSuchMethodException e) {
            // no-op
        }
    }
    this.accumulatorType = accumulatorType;
}
 

/**
 * 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());
}
 

@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));
}
 

@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
  SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry();
  CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
  finishSpecifyingStateSpecs(fn, coderRegistry, schemaRegistry, input.getCoder());
  TupleTag<OutputT> mainOutput = new TupleTag<>(MAIN_OUTPUT_TAG);
  PCollection<OutputT> res =
      input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);

  TypeDescriptor<OutputT> outputTypeDescriptor = getFn().getOutputTypeDescriptor();
  try {
    res.setSchema(
        schemaRegistry.getSchema(outputTypeDescriptor),
        outputTypeDescriptor,
        schemaRegistry.getToRowFunction(outputTypeDescriptor),
        schemaRegistry.getFromRowFunction(outputTypeDescriptor));
  } catch (NoSuchSchemaException e) {
    try {
      res.setCoder(
          coderRegistry.getCoder(
              outputTypeDescriptor,
              getFn().getInputTypeDescriptor(),
              ((PCollection<InputT>) input).getCoder()));
    } catch (CannotProvideCoderException e2) {
      // Ignore and leave coder unset.
    }
  }

  return res;
}
 

@Test
public void testAccumulatorCombiningStateWithUnderlying() throws CannotProvideCoderException {
  CopyOnAccessInMemoryStateInternals<String> underlying =
      CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
  CombineFn<Long, long[], Long> sumLongFn = Sum.ofLongs();

  StateNamespace namespace = new StateNamespaceForTest("foo");
  CoderRegistry reg = pipeline.getCoderRegistry();
  StateTag<CombiningState<Long, long[], Long>> stateTag =
      StateTags.combiningValue(
          "summer", sumLongFn.getAccumulatorCoder(reg, reg.getCoder(Long.class)), sumLongFn);
  GroupingState<Long, Long> underlyingValue = underlying.state(namespace, stateTag);
  assertThat(underlyingValue.read(), equalTo(0L));

  underlyingValue.add(1L);
  assertThat(underlyingValue.read(), equalTo(1L));

  CopyOnAccessInMemoryStateInternals<String> internals =
      CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
  GroupingState<Long, Long> copyOnAccessState = internals.state(namespace, stateTag);
  assertThat(copyOnAccessState.read(), equalTo(1L));

  copyOnAccessState.add(4L);
  assertThat(copyOnAccessState.read(), equalTo(5L));
  assertThat(underlyingValue.read(), equalTo(1L));

  GroupingState<Long, Long> reReadUnderlyingValue = underlying.state(namespace, stateTag);
  assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}
 

@Override
public PCollection<KV<K, V>> expand(PCollection<V> in) {
  PCollection<KV<K, V>> result =
      in.apply(
          "AddKeys",
          MapElements.via(
              new SimpleFunction<V, KV<K, V>>() {
                @Override
                public KV<K, V> apply(V element) {
                  return KV.of(fn.apply(element), element);
                }
              }));

  try {
    Coder<K> keyCoder;
    CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();
    if (keyType == null) {
      keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
    } else {
      keyCoder = coderRegistry.getCoder(keyType);
    }
    // TODO: Remove when we can set the coder inference context.
    result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
  } catch (CannotProvideCoderException exc) {
    // let lazy coder inference have a try
  }

  return result;
}
 

@Test
public void testCombinesElementsInSlidingWindows() throws Exception {
  CombineFn<Long, ?, Long> combineFn = Sum.ofLongs();
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          combineFn,
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));

  GroupAlsoByWindowProperties.combinesElementsInSlidingWindows(
      new CombiningGABWViaOutputBufferDoFnFactory<>(StringUtf8Coder.of(), appliedFn), combineFn);
}
 

@Test
public void testCreateCombiningNonmerging() throws Exception {
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofLongs(),
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));

  assertThat(
      BatchGroupAlsoByWindowsDoFns.create(windowingStrategy, appliedFn),
      instanceOf(BatchGroupAlsoByWindowAndCombineFn.class));
}
 

@VisibleForTesting
Coder<T> inferCoder(CoderRegistry coderRegistry) {
  try {
    return getCoder() != null
        ? getCoder()
        : coderRegistry.getCoder(TypeDescriptors.outputOf(getParseFn()));
  } catch (CannotProvideCoderException e) {
    throw new IllegalArgumentException(
        "Unable to infer coder for output of parseFn ("
            + TypeDescriptors.outputOf(getParseFn())
            + "). Specify it explicitly using withCoder().",
        e);
  }
}
 

@Override
public Coder<Object[]> getAccumulatorCoder(CoderRegistry registry, Coder<DataT> dataCoder)
    throws CannotProvideCoderException {
  List<Coder<Object>> coders = Lists.newArrayList();
  for (int i = 0; i < combineFnCount; ++i) {
    Coder<Object> inputCoder =
        combineInputCoders.get(i).isPresent()
            ? combineInputCoders.get(i).get()
            : registry.getOutputCoder(extractInputFns.get(i), dataCoder);
    coders.add(combineFns.get(i).getAccumulatorCoder(registry, inputCoder));
  }
  return new ComposedAccumulatorCoder(coders);
}
 

@Test
public void testCreateCombiningMerging() throws Exception {
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          Sum.ofLongs(),
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));
  WindowingStrategy<?, IntervalWindow> windowingStrategy =
      WindowingStrategy.of(Sessions.withGapDuration(Duration.millis(10)));

  assertThat(
      BatchGroupAlsoByWindowsDoFns.create(windowingStrategy, appliedFn),
      instanceOf(BatchGroupAlsoByWindowAndCombineFn.class));
}
 

@VisibleForTesting
Coder<T> inferCoder(CoderRegistry coderRegistry) {
  if (getCoder() != null) {
    return getCoder();
  }

  try {
    return coderRegistry.getCoder(TypeDescriptors.outputOf(getParseFn()));
  } catch (CannotProvideCoderException e) {
    throw new IllegalArgumentException(
        "Unable to infer coder for output of parseFn. Specify it explicitly using withCoder().",
        e);
  }
}
 

@Test
public void testCombinesIntoSessionsWithEndOfWindowTimestamp() throws Exception {
  CombineFn<Long, ?, Long> combineFn = Sum.ofLongs();
  AppliedCombineFn<String, Long, ?, Long> appliedFn =
      AppliedCombineFn.withInputCoder(
          combineFn,
          CoderRegistry.createDefault(),
          KvCoder.of(StringUtf8Coder.of(), VarLongCoder.of()));

  GroupAlsoByWindowProperties.combinesElementsPerSessionWithEndOfWindowTimestamp(
      new CombiningGABWViaOutputBufferDoFnFactory<>(StringUtf8Coder.of(), appliedFn), combineFn);
}
 

@VisibleForTesting
public static <K, InputT, AccumT, OutputT>
    AppliedCombineFn<K, InputT, AccumT, OutputT> withInputCoder(
        GlobalCombineFn<? super InputT, AccumT, OutputT> fn,
        CoderRegistry registry,
        KvCoder<K, InputT> kvCoder) {
  return withInputCoder(fn, registry, kvCoder, null, null);
}
 

@Test
public void testCreateWithCombinerAndBatchSideInputs() throws Exception {
  PipelineOptions options = PipelineOptionsFactory.create();

  Coder keyCoder = StringUtf8Coder.of();
  Coder valueCoder = BigEndianIntegerCoder.of();
  KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder);

  TestOutputReceiver receiver =
      new TestOutputReceiver(
          new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)),
          counterSet,
          NameContextsForTests.nameContextForTest());

  StepContext stepContext =
      BatchModeExecutionContext.forTesting(options, "testStage")
          .getStepContext(TestOperationContext.create(counterSet));

  when(mockSideInputReader.isEmpty()).thenReturn(false);

  ParDoFn pgbk =
      PartialGroupByKeyParDoFns.create(
          options,
          kvCoder,
          AppliedCombineFn.withInputCoder(
              Sum.ofIntegers(),
              CoderRegistry.createDefault(),
              kvCoder,
              ImmutableList.<PCollectionView<?>>of(),
              WindowingStrategy.globalDefault()),
          mockSideInputReader,
          receiver,
          stepContext);
  assertTrue(pgbk instanceof BatchSideInputPGBKParDoFn);
}
 

/** Tests that if different headers are found an exception is thrown. */
@Test(expected = RuntimeException.class)
public void testDifferentHeaders() {

  FailsafeElementCoder<String, String> coder = FAILSAFE_ELEMENT_CODER;

  CoderRegistry coderRegistry = pipeline.getCoderRegistry();
  coderRegistry.registerCoderForType(coder.getEncodedTypeDescriptor(), coder);

  PCollection<String> headers =
      pipeline.apply("CreateInput", Create.of(HEADER_STRING, "wrong,header,thing\n"));
  PCollection<String> lines = pipeline.apply("Create lines", Create.of(RECORD_STRING));

  PCollectionTuple readCsvHeadersOut =
      PCollectionTuple.of(CSV_HEADERS, headers).and(CSV_LINES, lines);

  PCollectionTuple test =
      readCsvHeadersOut.apply(
          "TestDifferentHeaders",
          CsvConverters.LineToFailsafeJson.newBuilder()
              .setDelimiter(",")
              .setUdfFileSystemPath(null)
              .setUdfFunctionName(null)
              .setJsonSchemaPath(null)
              .setHeaderTag(CSV_HEADERS)
              .setLineTag(CSV_LINES)
              .setUdfDeadletterTag(PROCESSING_DEADLETTER_OUT)
              .setUdfOutputTag(PROCESSING_OUT)
              .build());

  pipeline.run();
}
 

/** Tests the {@link BigQueryConverters.FailsafeJsonToTableRow} transform with good input. */
@Test
@Category(NeedsRunner.class)
public void testFailsafeJsonToTableRowValidInput() {
  // Test input
  final String payload = "{\"ticker\": \"GOOGL\", \"price\": 1006.94}";
  final Map<String, String> attributes = ImmutableMap.of("id", "0xDb12", "type", "stock");
  final PubsubMessage message = new PubsubMessage(payload.getBytes(), attributes);

  final FailsafeElement<PubsubMessage, String> input = FailsafeElement.of(message, payload);

  // Expected Output
  TableRow expectedRow = new TableRow().set("ticker", "GOOGL").set("price", 1006.94);

  // Register the coder for the pipeline. This prevents having to invoke .setCoder() on
  // many transforms.
  FailsafeElementCoder<PubsubMessage, String> coder =
      FailsafeElementCoder.of(PubsubMessageWithAttributesCoder.of(), StringUtf8Coder.of());

  CoderRegistry coderRegistry = pipeline.getCoderRegistry();
  coderRegistry.registerCoderForType(coder.getEncodedTypeDescriptor(), coder);

  // Build the pipeline
  PCollectionTuple output =
      pipeline
          .apply("CreateInput", Create.of(input).withCoder(coder))
          .apply(
              "JsonToTableRow",
              FailsafeJsonToTableRow.<PubsubMessage>newBuilder()
                  .setSuccessTag(TABLE_ROW_TAG)
                  .setFailureTag(FAILSAFE_ELM_TAG)
                  .build());

  // Assert
  PAssert.that(output.get(TABLE_ROW_TAG)).containsInAnyOrder(expectedRow);
  PAssert.that(output.get(FAILSAFE_ELM_TAG)).empty();

  // Execute the test
  pipeline.run();
}
 

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

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

/**
 * Attempt to infer a {@link Coder} by extracting the type of the deserialized-class from the
 * deserializer argument using the {@link Coder} registry.
 */
@VisibleForTesting
static <T> NullableCoder<T> inferCoder(
    CoderRegistry coderRegistry, Class<? extends Deserializer<T>> deserializer) {
  checkNotNull(deserializer);

  for (Type type : deserializer.getGenericInterfaces()) {
    if (!(type instanceof ParameterizedType)) {
      continue;
    }

    // This does not recurse: we will not infer from a class that extends
    // a class that extends Deserializer<T>.
    ParameterizedType parameterizedType = (ParameterizedType) type;

    if (parameterizedType.getRawType() == Deserializer.class) {
      Type parameter = parameterizedType.getActualTypeArguments()[0];

      @SuppressWarnings("unchecked")
      Class<T> clazz = (Class<T>) parameter;

      try {
        return NullableCoder.of(coderRegistry.getCoder(clazz));
      } catch (CannotProvideCoderException e) {
        throw new RuntimeException(
            String.format(
                "Unable to automatically infer a Coder for "
                    + "the Kafka Deserializer %s: no coder registered for type %s",
                deserializer, clazz));
      }
    }
  }

  throw new RuntimeException(
      String.format("Could not extract the Kafka Deserializer type from %s", deserializer));
}
 

@Override
public Coder<List<T>> getDefaultOutputCoder(CoderRegistry registry, Coder<T> inputCoder) {
  return ListCoder.of(inputCoder);
}
 

@Override
public Coder<Row> getDefaultOutputCoder(CoderRegistry registry, Coder<Row> inputCoder) {
  return SchemaCoder.of(EMPTY_SCHEMA);
}
 

@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)))));
}
 

@Override
public Coder<BoundedHeap<T, ComparatorT>> getAccumulatorCoder(
    CoderRegistry registry, Coder<T> inputCoder) {
  return new BoundedHeapCoder<>(count, compareFn, inputCoder);
}
 

@Override
public Coder<Object> getAccumulatorCoder(CoderRegistry registry, Coder<T> inputCoder)
    throws CannotProvideCoderException {
  return combineFn.getAccumulatorCoder(registry, inputCoder);
}
 

@Override
public Coder<HyperLogLogPlusPlus<HllT>> getAccumulatorCoder(
    CoderRegistry registry, Coder<byte[]> inputCoder) {
  // Use null to represent the "identity element" of the merge operation.
  return NullableCoder.of(HyperLogLogPlusPlusCoder.of());
}
 

@Override
public Coder<CovarianceAccumulator> getAccumulatorCoder(
    CoderRegistry registry, Coder<Row> inputCoder) {
  return SerializableCoder.of(CovarianceAccumulator.class);
}
 

@Override
public Coder<Accumulator> getAccumulatorCoder(
    CoderRegistry registry, Coder<Integer> inputCoder) {
  return Accumulator.getCoder();
}
 

@Override
public Coder<TestCombineFn.Accumulator> getAccumulatorCoder(
    CoderRegistry registry, Coder<Integer> inputCoder) {
  return TestCombineFn.Accumulator.getCoder();
}