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

public static WindowingStrategy getWindowingStrategy(
    String pCollectionId, RunnerApi.Components components) {
  RunnerApi.WindowingStrategy windowingStrategyProto =
      components.getWindowingStrategiesOrThrow(
          components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId());

  try {
    return WindowingStrategyTranslation.fromProto(
        windowingStrategyProto, RehydratedComponents.forComponents(components));
  } catch (InvalidProtocolBufferException e) {
    throw new IllegalStateException(
        String.format(
            "Unable to hydrate windowing strategy %s for %s.",
            windowingStrategyProto, pCollectionId),
        e);
  }
}
 

/**
 * Instantiates a runner-side wire coder for the given PCollection. Any component coders that are
 * unknown by the runner are replaced with length-prefixed byte arrays.
 *
 * @return a full or parameterized windowed value coder containing the PCollection's element coder
 */
public static <T> Coder<WindowedValue<T>> instantiateRunnerWireCoder(
    PCollectionNode pCollectionNode,
    RunnerApi.Components components,
    WireCoderSetting wireCoderSetting)
    throws IOException {
  // NOTE: We discard the new set of components so we don't bother to ensure it's consistent with
  // the caller's view.
  RunnerApi.Components.Builder builder = components.toBuilder();
  String protoCoderId = addRunnerWireCoder(pCollectionNode, builder, wireCoderSetting);
  Coder<?> javaCoder = RehydratedComponents.forComponents(builder.build()).getCoder(protoCoderId);
  checkArgument(
      javaCoder instanceof WindowedValue.FullWindowedValueCoder,
      "Unexpected Deserialized %s type, expected %s, got %s",
      RunnerApi.Coder.class.getSimpleName(),
      FullWindowedValueCoder.class.getSimpleName(),
      javaCoder.getClass());
  return (Coder<WindowedValue<T>>) javaCoder;
}
 

/**
 * Broadcast the side inputs of an executable stage. *This can be expensive.*
 *
 * @return Map from PCollection ID to Spark broadcast variable and coder to decode its contents.
 */
private static <SideInputT>
    ImmutableMap<String, Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>>>
        broadcastSideInputs(
            RunnerApi.ExecutableStagePayload stagePayload, SparkTranslationContext context) {
  Map<String, Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>>>
      broadcastVariables = new HashMap<>();
  for (SideInputId sideInputId : stagePayload.getSideInputsList()) {
    RunnerApi.Components stagePayloadComponents = stagePayload.getComponents();
    String collectionId =
        stagePayloadComponents
            .getTransformsOrThrow(sideInputId.getTransformId())
            .getInputsOrThrow(sideInputId.getLocalName());
    if (broadcastVariables.containsKey(collectionId)) {
      // This PCollection has already been broadcast.
      continue;
    }
    Tuple2<Broadcast<List<byte[]>>, WindowedValueCoder<SideInputT>> tuple2 =
        broadcastSideInput(collectionId, stagePayloadComponents, context);
    broadcastVariables.put(collectionId, tuple2);
  }
  return ImmutableMap.copyOf(broadcastVariables);
}
 

/** Creates a coder for a given PCollection id from the Proto definition. */
public static <T> Coder<WindowedValue<T>> instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException("Could not instantiate Coder", e);
  }
}
 

public WindowedValue.WindowedValueCoder instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return (WindowedValue.WindowedValueCoder)
        WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

private <T> void translateTestStream(
    String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext context) {
  RunnerApi.Components components = pipeline.getComponents();

  SerializableFunction<byte[], TestStream<T>> testStreamDecoder =
      bytes -> {
        try {
          RunnerApi.TestStreamPayload testStreamPayload =
              RunnerApi.TestStreamPayload.parseFrom(bytes);
          @SuppressWarnings("unchecked")
          TestStream<T> testStream =
              (TestStream<T>)
                  TestStreamTranslation.testStreamFromProtoPayload(
                      testStreamPayload, RehydratedComponents.forComponents(components));
          return testStream;
        } catch (Exception e) {
          throw new RuntimeException("Can't decode TestStream payload.", e);
        }
      };

  RunnerApi.PTransform transform = components.getTransformsOrThrow(id);
  String outputPCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values());
  Coder<WindowedValue<T>> coder = instantiateCoder(outputPCollectionId, components);

  DataStream<WindowedValue<T>> source =
      context
          .getExecutionEnvironment()
          .addSource(
              new TestStreamSource<>(
                  testStreamDecoder, transform.getSpec().getPayload().toByteArray()),
              new CoderTypeInformation<>(coder));

  context.addDataStream(outputPCollectionId, source);
}
 

public static WindowedValue.WindowedValueCoder instantiateCoder(
    String collectionId, RunnerApi.Components components) {
  PipelineNode.PCollectionNode collectionNode =
      PipelineNode.pCollection(collectionId, components.getPcollectionsOrThrow(collectionId));
  try {
    return (WindowedValue.WindowedValueCoder)
        WireCoders.instantiateRunnerWireCoder(collectionNode, components);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}
 

/** Create a timer reference from a TimerId proto and components. */
public static TimerReference fromTimerId(
    RunnerApi.ExecutableStagePayload.TimerId timerId, RunnerApi.Components components) {
  String transformId = timerId.getTransformId();
  String localName = timerId.getLocalName();
  RunnerApi.PTransform transform = components.getTransformsOrThrow(transformId);
  return of(PipelineNode.pTransform(transformId, transform), localName);
}
 

/** Create a side input reference from a SideInputId proto and components. */
public static SideInputReference fromSideInputId(
    SideInputId sideInputId, RunnerApi.Components components) {
  String transformId = sideInputId.getTransformId();
  String localName = sideInputId.getLocalName();
  String collectionId = components.getTransformsOrThrow(transformId).getInputsOrThrow(localName);
  PTransform transform = components.getTransformsOrThrow(transformId);
  PCollection collection = components.getPcollectionsOrThrow(collectionId);
  return SideInputReference.of(
      PipelineNode.pTransform(transformId, transform),
      localName,
      PipelineNode.pCollection(collectionId, collection));
}
 

public Set<String> allIds(RunnerApi.Components components) {
  Set<String> all = new HashSet<>();
  all.addAll(components.getTransformsMap().keySet());
  all.addAll(components.getPcollectionsMap().keySet());
  all.addAll(components.getCodersMap().keySet());
  all.addAll(components.getWindowingStrategiesMap().keySet());
  all.addAll(components.getEnvironmentsMap().keySet());
  return all;
}
 

@Test
public void testToProto() throws Exception {
  PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3));
  input.apply(Combine.globally(combineFn));
  final AtomicReference<AppliedPTransform<?, ?, Combine.Globally<?, ?>>> combine =
      new AtomicReference<>();
  pipeline.traverseTopologically(
      new PipelineVisitor.Defaults() {
        @Override
        public void leaveCompositeTransform(Node node) {
          if (node.getTransform() instanceof Combine.Globally) {
            checkState(combine.get() == null);
            combine.set((AppliedPTransform) node.toAppliedPTransform(getPipeline()));
          }
        }
      });
  checkState(combine.get() != null);
  assertEquals(combineFn, combine.get().getTransform().getFn());

  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
  CombinePayload combineProto =
      CombineTranslation.CombineGloballyPayloadTranslator.payloadForCombineGlobally(
          (AppliedPTransform) combine.get(), sdkComponents);
  RunnerApi.Components componentsProto = sdkComponents.toComponents();

  assertEquals(
      combineFn.getAccumulatorCoder(pipeline.getCoderRegistry(), input.getCoder()),
      getAccumulatorCoder(combineProto, RehydratedComponents.forComponents(componentsProto)));
  assertEquals(
      combineFn,
      SerializableUtils.deserializeFromByteArray(
          combineProto.getCombineFn().getPayload().toByteArray(), "CombineFn"));
}
 

/** Produces a {@link RunnerApi.Components} which contains only primitive transforms. */
@VisibleForTesting
static Collection<String> getPrimitiveTransformIds(RunnerApi.Components components) {
  Collection<String> ids = new LinkedHashSet<>();

  for (Map.Entry<String, PTransform> transformEntry : components.getTransformsMap().entrySet()) {
    PTransform transform = transformEntry.getValue();
    boolean isPrimitive = isPrimitiveTransform(transform);
    if (isPrimitive) {
      // Sometimes "primitive" transforms have sub-transforms (and even deeper-nested
      // descendents), due to runners
      // either rewriting them in terms of runner-specific transforms, or SDKs constructing them
      // in terms of other
      // underlying transforms (see https://issues.apache.org/jira/browse/BEAM-5441).
      // We consider any "leaf" descendents of these "primitive" transforms to be the true
      // "primitives" that we
      // preserve here; in the common case, this is just the "primitive" itself, which has no
      // descendents).
      Deque<String> transforms = new ArrayDeque<>();
      transforms.push(transformEntry.getKey());
      while (!transforms.isEmpty()) {
        String id = transforms.pop();
        PTransform next = components.getTransformsMap().get(id);
        List<String> subtransforms = next.getSubtransformsList();
        if (subtransforms.isEmpty()) {
          ids.add(id);
        } else {
          transforms.addAll(subtransforms);
        }
      }
    }
  }
  return ids;
}
 

/** Create a new {@link RehydratedComponents} from a Runner API {@link Components}. */
public static RehydratedComponents forComponents(RunnerApi.Components components) {
  return new RehydratedComponents(components, null);
}
 

@Override
@SuppressWarnings("unchecked")
public ExecutableStage createExecutableStage() throws Exception {
	RunnerApi.Components components =
		RunnerApi.Components.newBuilder()
			.putPcollections(
				INPUT_ID,
				RunnerApi.PCollection.newBuilder()
					.setWindowingStrategyId(WINDOW_STRATEGY)
					.setCoderId(INPUT_CODER_ID)
					.build())
			.putPcollections(
				OUTPUT_ID,
				RunnerApi.PCollection.newBuilder()
					.setWindowingStrategyId(WINDOW_STRATEGY)
					.setCoderId(OUTPUT_CODER_ID)
					.build())
			.putTransforms(
				TRANSFORM_ID,
				RunnerApi.PTransform.newBuilder()
					.setUniqueName(TRANSFORM_ID)
					.setSpec(RunnerApi.FunctionSpec.newBuilder()
						.setUrn(functionUrn)
						.setPayload(
							org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString.copyFrom(
								getUserDefinedFunctionsProto().toByteArray()))
						.build())
					.putInputs(MAIN_INPUT_NAME, INPUT_ID)
					.putOutputs(MAIN_OUTPUT_NAME, OUTPUT_ID)
					.build())
			.putWindowingStrategies(
				WINDOW_STRATEGY,
				RunnerApi.WindowingStrategy.newBuilder()
					.setWindowCoderId(WINDOW_CODER_ID)
					.build())
			.putCoders(
				INPUT_CODER_ID,
				getInputCoderProto())
			.putCoders(
				OUTPUT_CODER_ID,
				getOutputCoderProto())
			.putCoders(
				WINDOW_CODER_ID,
				getWindowCoderProto())
			.build();

	PipelineNode.PCollectionNode input =
		PipelineNode.pCollection(INPUT_ID, components.getPcollectionsOrThrow(INPUT_ID));
	List<SideInputReference> sideInputs = Collections.EMPTY_LIST;
	List<UserStateReference> userStates = Collections.EMPTY_LIST;
	List<TimerReference> timers = Collections.EMPTY_LIST;
	List<PipelineNode.PTransformNode> transforms =
		Collections.singletonList(
			PipelineNode.pTransform(TRANSFORM_ID, components.getTransformsOrThrow(TRANSFORM_ID)));
	List<PipelineNode.PCollectionNode> outputs =
		Collections.singletonList(
			PipelineNode.pCollection(OUTPUT_ID, components.getPcollectionsOrThrow(OUTPUT_ID)));
	return ImmutableExecutableStage.of(
		components, createPythonExecutionEnvironment(), input, sideInputs, userStates, timers, transforms, outputs, createValueOnlyWireCoderSetting());
}
 

@Test
public void replacesOnlyMatching() {
  RunnerApi.Pipeline p =
      Pipeline.newBuilder()
          .addAllRootTransformIds(ImmutableList.of("first", "second"))
          .setComponents(
              Components.newBuilder()
                  .putTransforms(
                      "first",
                      PTransform.newBuilder()
                          .setSpec(FunctionSpec.newBuilder().setUrn("beam:first"))
                          .build())
                  .putTransforms(
                      "second",
                      PTransform.newBuilder()
                          .setSpec(FunctionSpec.newBuilder().setUrn("beam:second"))
                          .build())
                  .putPcollections(
                      "intermediatePc",
                      PCollection.newBuilder().setUniqueName("intermediate").build())
                  .putCoders(
                      "coder",
                      Coder.newBuilder().setSpec(FunctionSpec.getDefaultInstance()).build()))
          .build();

  PTransform secondReplacement =
      PTransform.newBuilder()
          .addSubtransforms("second_sub")
          .setSpec(
              FunctionSpec.newBuilder()
                  .setUrn("beam:second:replacement")
                  .setPayload(
                      ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8))))
          .build();
  WindowingStrategy introducedWS =
      WindowingStrategy.newBuilder().setAccumulationMode(Enum.ACCUMULATING).build();
  RunnerApi.Components extraComponents =
      Components.newBuilder()
          .putPcollections(
              "intermediatePc",
              PCollection.newBuilder().setUniqueName("intermediate_replacement").build())
          .putWindowingStrategies("new_ws", introducedWS)
          .putTransforms("second_sub", PTransform.getDefaultInstance())
          .build();

  Pipeline updated =
      ProtoOverrides.updateTransform(
          "beam:second", p, new TestReplacer(secondReplacement, extraComponents));
  PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second");

  assertThat(updatedSecond, equalTo(secondReplacement));
  assertThat(
      updated.getComponents().getWindowingStrategiesOrThrow("new_ws"), equalTo(introducedWS));
  assertThat(
      updated.getComponents().getTransformsOrThrow("second_sub"),
      equalTo(PTransform.getDefaultInstance()));

  // TODO: This might not be appropriate. Merging in the other direction might force that callers
  // are well behaved.
  assertThat(
      updated.getComponents().getPcollectionsOrThrow("intermediatePc").getUniqueName(),
      equalTo("intermediate_replacement"));

  // Assert that the untouched components are unchanged.
  assertThat(
      updated.getComponents().getTransformsOrThrow("first"),
      equalTo(p.getComponents().getTransformsOrThrow("first")));
  assertThat(
      updated.getComponents().getCodersOrThrow("coder"),
      equalTo(p.getComponents().getCodersOrThrow("coder")));
  assertThat(updated.getRootTransformIdsList(), equalTo(p.getRootTransformIdsList()));
}
 

@Test
public void toTransformProto() throws Exception {
  Map<TupleTag<?>, PValue> inputs = new HashMap<>();
  inputs.put(new TupleTag<KV<Long, String>>("mainInputName") {}, mainInput);
  inputs.putAll(parDo.getAdditionalInputs());
  PCollectionTuple output = mainInput.apply(parDo);

  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));

  // Encode
  RunnerApi.PTransform protoTransform =
      PTransformTranslation.toProto(
          AppliedPTransform.<PCollection<KV<Long, String>>, PCollection<Void>, MultiOutput>of(
              "foo", inputs, output.expand(), parDo, p),
          sdkComponents);
  RunnerApi.Components components = sdkComponents.toComponents();
  RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components);

  // Decode
  ParDoPayload parDoPayload = ParDoPayload.parseFrom(protoTransform.getSpec().getPayload());
  for (PCollectionView<?> view : parDo.getSideInputs().values()) {
    SideInput sideInput = parDoPayload.getSideInputsOrThrow(view.getTagInternal().getId());
    PCollectionView<?> restoredView =
        PCollectionViewTranslation.viewFromProto(
            sideInput,
            view.getTagInternal().getId(),
            view.getPCollection(),
            protoTransform,
            rehydratedComponents);
    assertThat(restoredView.getTagInternal(), equalTo(view.getTagInternal()));
    assertThat(restoredView.getViewFn(), instanceOf(view.getViewFn().getClass()));
    assertThat(
        restoredView.getWindowMappingFn(), instanceOf(view.getWindowMappingFn().getClass()));
    assertThat(
        restoredView.getWindowingStrategyInternal(),
        equalTo(view.getWindowingStrategyInternal().fixDefaults()));
    assertThat(restoredView.getCoderInternal(), equalTo(view.getCoderInternal()));
  }
  String mainInputId = sdkComponents.registerPCollection(mainInput);
  assertThat(
      ParDoTranslation.getMainInput(protoTransform, components),
      equalTo(components.getPcollectionsOrThrow(mainInputId)));
  assertThat(ParDoTranslation.getMainInputName(protoTransform), equalTo("mainInputName"));

  // Ensure the correct timer coder components are used from the main input PCollection's key
  // and window coders.
  for (RunnerApi.TimerFamilySpec timerFamilySpec :
      parDoPayload.getTimerFamilySpecsMap().values()) {
    Coder<?> timerCoder =
        CoderTranslation.fromProto(
            components.getCodersOrThrow(timerFamilySpec.getTimerFamilyCoderId()),
            rehydratedComponents,
            TranslationContext.DEFAULT);
    assertEquals(
        org.apache.beam.runners.core.construction.Timer.Coder.of(
            VarLongCoder.of(), GlobalWindow.Coder.INSTANCE),
        timerCoder);
  }
}
 

@Test
public void testConstructPubsubRead() throws Exception {
  String topic = "projects/project-1234/topics/topic_name";
  String idAttribute = "id_foo";
  Boolean needsAttributes = true;

  ExternalTransforms.ExternalConfigurationPayload payload =
      ExternalTransforms.ExternalConfigurationPayload.newBuilder()
          .putConfiguration(
              "topic",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(topic)))
                  .build())
          .putConfiguration(
              "id_label",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(idAttribute)))
                  .build())
          .putConfiguration(
              "with_attributes",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:bool:v1")
                  .setPayload(ByteString.copyFrom(encodeBoolean(needsAttributes)))
                  .build())
          .build();

  RunnerApi.Components defaultInstance = RunnerApi.Components.getDefaultInstance();
  ExpansionApi.ExpansionRequest request =
      ExpansionApi.ExpansionRequest.newBuilder()
          .setComponents(defaultInstance)
          .setTransform(
              RunnerApi.PTransform.newBuilder()
                  .setUniqueName("test")
                  .setSpec(
                      RunnerApi.FunctionSpec.newBuilder()
                          .setUrn("beam:external:java:pubsub:read:v1")
                          .setPayload(payload.toByteString())))
          .setNamespace("test_namespace")
          .build();

  ExpansionService expansionService = new ExpansionService();
  TestStreamObserver<ExpansionApi.ExpansionResponse> observer = new TestStreamObserver<>();
  expansionService.expand(request, observer);

  ExpansionApi.ExpansionResponse result = observer.result;
  RunnerApi.PTransform transform = result.getTransform();
  assertThat(
      transform.getSubtransformsList(),
      Matchers.contains(
          "test_namespacetest/PubsubUnboundedSource", "test_namespacetest/MapElements"));
  assertThat(transform.getInputsCount(), Matchers.is(0));
  assertThat(transform.getOutputsCount(), Matchers.is(1));

  RunnerApi.PTransform pubsubComposite =
      result.getComponents().getTransformsOrThrow(transform.getSubtransforms(0));
  RunnerApi.PTransform pubsubRead =
      result.getComponents().getTransformsOrThrow(pubsubComposite.getSubtransforms(0));
  RunnerApi.ReadPayload readPayload =
      RunnerApi.ReadPayload.parseFrom(pubsubRead.getSpec().getPayload());
  PubsubUnboundedSource.PubsubSource source =
      (PubsubUnboundedSource.PubsubSource) ReadTranslation.unboundedSourceFromProto(readPayload);
  PubsubUnboundedSource spec = source.outer;

  assertThat(
      spec.getTopicProvider() == null ? null : String.valueOf(spec.getTopicProvider()),
      Matchers.is(topic));
  assertThat(spec.getIdAttribute(), Matchers.is(idAttribute));
  assertThat(spec.getNeedsAttributes(), Matchers.is(true));
}
 

/**
 * Convert this {@link SdkComponents} into a {@link RunnerApi.Components}, including all of the
 * contained {@link Coder coders}, {@link WindowingStrategy windowing strategies}, {@link
 * PCollection PCollections}, and {@link PTransform PTransforms}.
 */
public RunnerApi.Components toComponents() {
  return componentsBuilder.build();
}
 

/**
 * Create new {@link SdkComponents} importing all items from provided {@link Components} object.
 *
 * <p>WARNING: This action might cause some of duplicate items created.
 */
public static SdkComponents create(
    RunnerApi.Components components, Collection<String> requirements) {
  return new SdkComponents(components, requirements, "");
}
 

abstract RunnerApi.Components getDeduplicatedComponents();