Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#Coder

/**
 * Patches the input coder of a stateful transform to ensure that the byte representation of a key
 * used to partition the input element at the Runner, matches the key byte representation received
 * for state requests and timers from the SDK Harness. Stateful transforms always have a KvCoder
 * as input.
 */
private static void lengthPrefixKeyCoder(
    String inputColId, Components.Builder componentsBuilder) {
  RunnerApi.PCollection pcollection = componentsBuilder.getPcollectionsOrThrow(inputColId);
  RunnerApi.Coder kvCoder = componentsBuilder.getCodersOrThrow(pcollection.getCoderId());
  Preconditions.checkState(
      ModelCoders.KV_CODER_URN.equals(kvCoder.getSpec().getUrn()),
      "Stateful executable stages must use a KV coder, but is: %s",
      kvCoder.getSpec().getUrn());
  String keyCoderId = ModelCoders.getKvCoderComponents(kvCoder).keyCoderId();
  // Retain the original coder, but wrap in LengthPrefixCoder
  String newKeyCoderId =
      LengthPrefixUnknownCoders.addLengthPrefixedCoder(keyCoderId, componentsBuilder, false);
  // Replace old key coder with LengthPrefixCoder<old_key_coder>
  kvCoder = kvCoder.toBuilder().setComponentCoderIds(0, newKeyCoderId).build();
  componentsBuilder.putCoders(pcollection.getCoderId(), kvCoder);
}
 

@Test
public void toAndFromProto() throws Exception {
  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
  RunnerApi.Coder coderProto = CoderTranslation.toProto(coder, sdkComponents);

  Components encodedComponents = sdkComponents.toComponents();
  Coder<?> decodedCoder =
      CoderTranslation.fromProto(
          coderProto,
          RehydratedComponents.forComponents(encodedComponents),
          TranslationContext.DEFAULT);
  assertThat(decodedCoder, equalTo(coder));

  if (KNOWN_CODERS.contains(coder)) {
    for (RunnerApi.Coder encodedCoder : encodedComponents.getCodersMap().values()) {
      assertThat(
          encodedCoder.getSpec().getUrn(),
          not(equalTo(CoderTranslation.JAVA_SERIALIZED_CODER_URN)));
    }
  }
}
 

private static Coder<?> fromKnownCoder(
    RunnerApi.Coder coder, RehydratedComponents components, TranslationContext context)
    throws IOException {
  String coderUrn = coder.getSpec().getUrn();
  List<Coder<?>> coderComponents = new ArrayList<>();
  for (String componentId : coder.getComponentCoderIdsList()) {
    // Only store coders in RehydratedComponents as long as we are not using a custom
    // translation context.
    Coder<?> innerCoder =
        context == TranslationContext.DEFAULT
            ? components.getCoder(componentId)
            : fromProto(
                components.getComponents().getCodersOrThrow(componentId), components, context);
    coderComponents.add(innerCoder);
  }
  Class<? extends Coder> coderType = KNOWN_CODER_URNS.inverse().get(coderUrn);
  CoderTranslator<?> translator = KNOWN_TRANSLATORS.get(coderType);
  checkArgument(
      translator != null,
      "Unknown Coder URN %s. Known URNs: %s",
      coderUrn,
      KNOWN_CODER_URNS.values());
  return translator.fromComponents(
      coderComponents, coder.getSpec().getPayload().toByteArray(), context);
}
 

/**
 * Registers the provided {@link Coder} into this {@link SdkComponents}, returning a unique ID for
 * the {@link Coder}. Multiple registrations of the same {@link Coder} will return the same unique
 * ID.
 *
 * <p>Coders are stored by identity to ensure that coders with implementations of {@link
 * #equals(Object)} and {@link #hashCode()} but incompatible binary formats are not considered the
 * same coder.
 */
public String registerCoder(Coder<?> coder) throws IOException {
  String existing = coderIds.get(coder);
  if (existing != null) {
    return existing;
  }
  String baseName = NameUtils.approximateSimpleName(coder);
  String name = uniqify(baseName, coderIds.values());
  coderIds.put(coder, name);
  RunnerApi.Coder coderProto = CoderTranslation.toProto(coder, this);
  componentsBuilder.putCoders(name, coderProto);
  return name;
}
 

@Override
public Coder<?> load(String id) throws Exception {
  @Nullable RunnerApi.Coder coder = components.getCodersOrDefault(id, null);
  checkState(coder != null, "No coder with id '%s' in serialized components", id);
  return CoderTranslation.fromProto(
      coder, RehydratedComponents.this, TranslationContext.DEFAULT);
}
 

boolean isJavaSDKCompatible(RunnerApi.Components components, String coderId) {
  RunnerApi.Coder coder = components.getCodersOrThrow(coderId);
  if (!CoderTranslation.JAVA_SERIALIZED_CODER_URN.equals(coder.getSpec().getUrn())
      && !CoderTranslation.KNOWN_CODER_URNS.containsValue(coder.getSpec().getUrn())) {
    return false;
  }
  for (String componentId : coder.getComponentCoderIdsList()) {
    if (!isJavaSDKCompatible(components, componentId)) {
      return false;
    }
  }
  return true;
}
 

public static Coder<?> fromProto(
    RunnerApi.Coder protoCoder, RehydratedComponents components, TranslationContext context)
    throws IOException {
  String coderSpecUrn = protoCoder.getSpec().getUrn();
  if (coderSpecUrn.equals(JAVA_SERIALIZED_CODER_URN)) {
    return fromCustomCoder(protoCoder);
  }
  return fromKnownCoder(protoCoder, components, context);
}
 

private static RunnerApi.Coder toCustomCoder(Coder<?> coder) throws IOException {
  RunnerApi.Coder.Builder coderBuilder = RunnerApi.Coder.newBuilder();
  return coderBuilder
      .setSpec(
          FunctionSpec.newBuilder()
              .setUrn(JAVA_SERIALIZED_CODER_URN)
              .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(coder)))
              .build())
      .build();
}
 

private static RunnerApi.Coder toKnownCoder(Coder<?> coder, SdkComponents components)
    throws IOException {
  CoderTranslator translator = KNOWN_TRANSLATORS.get(coder.getClass());
  List<String> componentIds = registerComponents(coder, translator, components);
  return RunnerApi.Coder.newBuilder()
      .addAllComponentCoderIds(componentIds)
      .setSpec(
          FunctionSpec.newBuilder()
              .setUrn(KNOWN_CODER_URNS.get(coder.getClass()))
              .setPayload(ByteString.copyFrom(translator.getPayload(coder))))
      .build();
}
 

public static RunnerApi.Coder toProto(Coder<?> coder, SdkComponents components)
    throws IOException {
  if (KNOWN_CODER_URNS.containsKey(coder.getClass())) {
    return toKnownCoder(coder, components);
  }
  return toCustomCoder(coder);
}
 

public static RunnerApi.MessageWithComponents toProto(Coder<?> coder) throws IOException {
  SdkComponents components = SdkComponents.create();
  RunnerApi.Coder coderProto = toProto(coder, components);
  return RunnerApi.MessageWithComponents.newBuilder()
      .setCoder(coderProto)
      .setComponents(components.toComponents())
      .build();
}
 

private RunnerApi.Coder getRowCoderProto(RowType rowType) {
	return RunnerApi.Coder.newBuilder()
		.setSpec(
			RunnerApi.FunctionSpec.newBuilder()
				.setUrn(getInputOutputCoderUrn())
				.setPayload(org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString.copyFrom(
					toProtoType(rowType).getRowSchema().toByteArray()))
				.build())
		.build();
}
 

private static void ensureLengthPrefixed(
    RunnerApi.Coder coder,
    RunnerApi.Coder originalCoder,
    Map<String, RunnerApi.Coder> pbsCoderMap) {
  assertThat(coder.getSpec().getUrn(), is(ModelCoders.LENGTH_PREFIX_CODER_URN));
  // Check that the wrapped coder is unchanged
  String lengthPrefixedWrappedCoderId = coder.getComponentCoderIds(0);
  assertThat(pbsCoderMap.get(lengthPrefixedWrappedCoderId), is(originalCoder));
}
 

BeamFnDataWriteRunner(
    String pTransformId,
    RunnerApi.PTransform remoteWriteNode,
    Supplier<String> processBundleInstructionIdSupplier,
    Map<String, RunnerApi.Coder> coders,
    BeamFnDataClient beamFnDataClientFactory,
    BeamFnStateClient beamFnStateClient)
    throws IOException {
  this.pTransformId = pTransformId;
  RemoteGrpcPort port = RemoteGrpcPortWrite.fromPTransform(remoteWriteNode).getPort();
  this.apiServiceDescriptor = port.getApiServiceDescriptor();
  this.beamFnDataClientFactory = beamFnDataClientFactory;
  this.processBundleInstructionIdSupplier = processBundleInstructionIdSupplier;

  RehydratedComponents components =
      RehydratedComponents.forComponents(Components.newBuilder().putAllCoders(coders).build());
  this.coder =
      (Coder<WindowedValue<InputT>>)
          CoderTranslation.fromProto(
              coders.get(port.getCoderId()),
              components,
              new StateBackedIterableTranslationContext() {
                @Override
                public BeamFnStateClient getStateClient() {
                  return beamFnStateClient;
                }

                @Override
                public Supplier<String> getCurrentInstructionId() {
                  return processBundleInstructionIdSupplier;
                }
              });
}
 

@Override
public BeamFnDataWriteRunner<InputT> createRunnerForPTransform(
    PipelineOptions pipelineOptions,
    BeamFnDataClient beamFnDataClient,
    BeamFnStateClient beamFnStateClient,
    BeamFnTimerClient beamFnTimerClient,
    String pTransformId,
    PTransform pTransform,
    Supplier<String> processBundleInstructionId,
    Map<String, PCollection> pCollections,
    Map<String, RunnerApi.Coder> coders,
    Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
    PCollectionConsumerRegistry pCollectionConsumerRegistry,
    PTransformFunctionRegistry startFunctionRegistry,
    PTransformFunctionRegistry finishFunctionRegistry,
    Consumer<ThrowingRunnable> tearDownFunctions,
    Consumer<ProgressRequestCallback> addProgressRequestCallback,
    BundleSplitListener splitListener,
    BundleFinalizer bundleFinalizer)
    throws IOException {

  BeamFnDataWriteRunner<InputT> runner =
      new BeamFnDataWriteRunner<>(
          pTransformId,
          pTransform,
          processBundleInstructionId,
          coders,
          beamFnDataClient,
          beamFnStateClient);
  startFunctionRegistry.register(pTransformId, runner::registerForOutput);
  pCollectionConsumerRegistry.register(
      getOnlyElement(pTransform.getInputsMap().values()),
      pTransformId,
      (FnDataReceiver) (FnDataReceiver<WindowedValue<InputT>>) runner::consume);

  finishFunctionRegistry.register(pTransformId, runner::close);
  return runner;
}
 

private static String addWireCoder(
    PCollectionNode pCollectionNode,
    RunnerApi.Components.Builder components,
    boolean useByteArrayCoder,
    WireCoderSetting wireCoderSetting) {
  String elementCoderId = pCollectionNode.getPCollection().getCoderId();
  String windowingStrategyId = pCollectionNode.getPCollection().getWindowingStrategyId();
  String windowCoderId =
      components.getWindowingStrategiesOrThrow(windowingStrategyId).getWindowCoderId();

  // decide type of windowedValueCoder according to the wire coder setting.
  RunnerApi.Coder windowedValueCoder;
  String wireCoderUrn = wireCoderSetting.getUrn();
  if (wireCoderUrn.equals(getUrn(RunnerApi.StandardCoders.Enum.WINDOWED_VALUE))
      || wireCoderUrn.isEmpty()) {
    windowedValueCoder = ModelCoders.windowedValueCoder(elementCoderId, windowCoderId);
  } else {
    checkArgument(
        wireCoderUrn.equals(getUrn(RunnerApi.StandardCoders.Enum.PARAM_WINDOWED_VALUE)),
        "Unexpected wire coder urn %s, currently only %s or %s are supported!",
        wireCoderUrn,
        getUrn(RunnerApi.StandardCoders.Enum.WINDOWED_VALUE),
        getUrn(RunnerApi.StandardCoders.Enum.PARAM_WINDOWED_VALUE));
    windowedValueCoder =
        ModelCoders.paramWindowedValueCoder(
            elementCoderId, windowCoderId, wireCoderSetting.getPayload().toByteArray());
  }

  // Add the original WindowedValue<T, W> coder to the components;
  String windowedValueId =
      SyntheticComponents.uniqueId(
          String.format("fn/wire/%s", pCollectionNode.getId()), components::containsCoders);
  components.putCoders(windowedValueId, windowedValueCoder);
  return LengthPrefixUnknownCoders.addLengthPrefixedCoder(
      windowedValueId, components, useByteArrayCoder);
}
 

private static String addCoder(
    RunnerApi.Coder coder, RunnerApi.Components.Builder components, String uniqueIdPrefix) {
  for (Entry<String, Coder> entry : components.getCodersMap().entrySet()) {
    if (entry.getValue().equals(coder)) {
      return entry.getKey();
    }
  }
  String id = generateUniqueId(uniqueIdPrefix, components::containsCoders);
  components.putCoders(id, coder);
  return id;
}
 

private static RunnerApi.Coder buildProto(
    Deque<String> coderUrns, RunnerApi.Components.Builder componentsBuilder) {
  Preconditions.checkArgument(coderUrns.size() > 0, "No URNs left to construct coder from");

  final String coderUrn = coderUrns.pop();
  RunnerApi.Coder.Builder coderBuilder =
      RunnerApi.Coder.newBuilder()
          .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn(coderUrn).build());

  if (coderUrn.equals(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.ITERABLE))) {
    RunnerApi.Coder elementCoder = buildProto(coderUrns, componentsBuilder);
    String coderId = UUID.randomUUID().toString();
    componentsBuilder.putCoders(coderId, elementCoder);
    coderBuilder.addComponentCoderIds(coderId);
  } else if (coderUrn.equals(BeamUrns.getUrn(RunnerApi.StandardCoders.Enum.KV))) {
    RunnerApi.Coder element1Coder = buildProto(coderUrns, componentsBuilder);
    RunnerApi.Coder element2Coder = buildProto(coderUrns, componentsBuilder);
    String coderId1 = UUID.randomUUID().toString();
    String coderId2 = UUID.randomUUID().toString();
    componentsBuilder.putCoders(coderId1, element1Coder);
    componentsBuilder.putCoders(coderId2, element2Coder);
    coderBuilder.addComponentCoderIds(coderId1);
    coderBuilder.addComponentCoderIds(coderId2);
  }

  return coderBuilder.build();
}
 

@Override
public final FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT, OutputT>
    createRunnerForPTransform(
        PipelineOptions pipelineOptions,
        BeamFnDataClient beamFnDataClient,
        BeamFnStateClient beamFnStateClient,
        BeamFnTimerClient beamFnTimerClient,
        String pTransformId,
        PTransform pTransform,
        Supplier<String> processBundleInstructionId,
        Map<String, PCollection> pCollections,
        Map<String, RunnerApi.Coder> coders,
        Map<String, RunnerApi.WindowingStrategy> windowingStrategies,
        PCollectionConsumerRegistry pCollectionConsumerRegistry,
        PTransformFunctionRegistry startFunctionRegistry,
        PTransformFunctionRegistry finishFunctionRegistry,
        Consumer<ThrowingRunnable> tearDownFunctions,
        Consumer<ProgressRequestCallback> addProgressRequestCallback,
        BundleSplitListener splitListener,
        BundleFinalizer bundleFinalizer) {

  FnApiDoFnRunner<InputT, RestrictionT, PositionT, WatermarkEstimatorStateT, OutputT> runner =
      new FnApiDoFnRunner<>(
          pipelineOptions,
          beamFnStateClient,
          beamFnTimerClient,
          pTransformId,
          pTransform,
          processBundleInstructionId,
          pCollections,
          coders,
          windowingStrategies,
          pCollectionConsumerRegistry,
          startFunctionRegistry,
          finishFunctionRegistry,
          tearDownFunctions,
          addProgressRequestCallback,
          splitListener,
          bundleFinalizer);

  return runner;
}
 

private static Coder<?> fromCustomCoder(RunnerApi.Coder protoCoder) throws IOException {
  return (Coder<?>)
      SerializableUtils.deserializeFromByteArray(
          protoCoder.getSpec().getPayload().toByteArray(), "Custom Coder Bytes");
}