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

@Test
public void testFetchReadySideInputsAndFilterBlockedStreamingSideInputsNode() {
  WindowingStrategy windowingStrategy = WindowingStrategy.globalDefault();
  Map<PCollectionView<?>, RunnerApi.FunctionSpec> pcollectionViewsToWindowMappingFns =
      ImmutableMap.of(
          mock(PCollectionView.class),
          FunctionSpec.newBuilder().setUrn("beam:test:urn:1.0").build());
  NameContext nameContext = NameContextsForTests.nameContextForTest();
  assertSame(
      FetchAndFilterStreamingSideInputsNode.create(
              windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext)
          .getWindowingStrategy(),
      windowingStrategy);
  assertSame(
      FetchAndFilterStreamingSideInputsNode.create(
              windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext)
          .getPCollectionViewsToWindowMappingFns(),
      pcollectionViewsToWindowMappingFns);
  assertSame(
      FetchAndFilterStreamingSideInputsNode.create(
              windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext)
          .getNameContext(),
      nameContext);
}
 

@Test
public void testRegistrarEncodedProto() throws Exception {
  PCollection<String> output = p.apply(testStream);

  AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream =
      AppliedPTransform.of("fakeName", PBegin.in(p).expand(), output.expand(), testStream, p);

  SdkComponents components = SdkComponents.create();
  components.registerEnvironment(Environments.createDockerEnvironment("java"));
  RunnerApi.FunctionSpec spec =
      PTransformTranslation.toProto(appliedTestStream, components).getSpec();

  assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN));

  RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload());

  verifyTestStreamEncoding(
      testStream, payload, RehydratedComponents.forComponents(components.toComponents()));
}
 

@Test
public void testWindowMapping() throws Exception {
  String pTransformId = "pTransformId";

  SdkComponents components = SdkComponents.create();
  components.registerEnvironment(Environments.createDockerEnvironment("java"));
  RunnerApi.FunctionSpec functionSpec =
      RunnerApi.FunctionSpec.newBuilder()
          .setUrn(WindowMappingFnRunner.URN)
          .setPayload(
              ParDoTranslation.translateWindowMappingFn(
                      new GlobalWindows().getDefaultWindowMappingFn(), components)
                  .toByteString())
          .build();
  RunnerApi.PTransform pTransform =
      RunnerApi.PTransform.newBuilder().setSpec(functionSpec).build();

  ThrowingFunction<KV<Object, BoundedWindow>, KV<Object, BoundedWindow>> mapFunction =
      WindowMappingFnRunner.createMapFunctionForPTransform(pTransformId, pTransform);

  KV<Object, BoundedWindow> input =
      KV.of("abc", new IntervalWindow(Instant.now(), Duration.standardMinutes(1)));

  assertEquals(KV.of(input.getKey(), GlobalWindow.INSTANCE), mapFunction.apply(input));
}
 

public FetchAndFilterStreamingSideInputsOperation(
    OutputReceiver[] receivers,
    DataflowOperationContext context,
    InstructionRequestHandler instructionRequestHandler,
    FnDataService beamFnDataService,
    ApiServiceDescriptor dataServiceApiServiceDescriptor,
    IdGenerator idGenerator,
    Coder<WindowedValue<T>> inputCoder,
    WindowingStrategy<?, W> windowingStrategy,
    DataflowExecutionContext.DataflowStepContext stepContext,
    Map<PCollectionView<?>, RunnerApi.FunctionSpec> pCollectionViewToWindowMappingFns) {
  super(receivers, context);

  this.sideInputFetcher =
      new StreamingSideInputFetcher<>(
          buildPCollectionViewsWithSdkSupportedWindowMappingFn(
              idGenerator,
              instructionRequestHandler,
              beamFnDataService,
              dataServiceApiServiceDescriptor,
              ((FullWindowedValueCoder) inputCoder).getWindowCoder(),
              pCollectionViewToWindowMappingFns),
          ((WindowedValueCoder) inputCoder).getValueCoder(),
          windowingStrategy,
          (StreamingModeExecutionContext.StreamingModeStepContext)
              stepContext.namespacedToUser());
  this.elementsToProcess = new LinkedBlockingQueue<>();
  this.singleThreadExecutor = Executors.newSingleThreadExecutor();
}
 

private static <W extends BoundedWindow> RunnerApi.PTransform createMergeTransformForWindowFn(
    WindowFn<?, W> windowFn) throws Exception {
  SdkComponents components = SdkComponents.create();
  components.registerEnvironment(Environments.createDockerEnvironment("test"));
  RunnerApi.FunctionSpec functionSpec =
      RunnerApi.FunctionSpec.newBuilder()
          .setUrn(WindowMergingFnRunner.URN)
          .setPayload(WindowingStrategyTranslation.toProto(windowFn, components).toByteString())
          .build();
  return RunnerApi.PTransform.newBuilder().setSpec(functionSpec).build();
}
 

private <T> RunnerApi.FunctionSpec translateTyped(
    final TestStream<T> testStream, SdkComponents components) throws IOException {
  return RunnerApi.FunctionSpec.newBuilder()
      .setUrn(TEST_STREAM_TRANSFORM_URN)
      .setPayload(payloadForTestStream(testStream, components).toByteString())
      .build();
}
 

public static <T, W extends BoundedWindow> SideInputSpec of(
    String transformId,
    String sideInputId,
    RunnerApi.FunctionSpec accessPattern,
    Coder<T> elementCoder,
    Coder<W> windowCoder) {
  return new AutoValue_ProcessBundleDescriptors_SideInputSpec(
      transformId, sideInputId, accessPattern, elementCoder, windowCoder);
}
 

/**
 * Converts a {@link org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec} into a {@link
 * ViewFn} using the URN.
 */
public static ViewFn<?, ?> viewFnFromProto(RunnerApi.FunctionSpec viewFn)
    throws InvalidProtocolBufferException {
  RunnerApi.FunctionSpec spec = viewFn;
  checkArgument(
      spec.getUrn().equals(ParDoTranslation.CUSTOM_JAVA_VIEW_FN_URN),
      "Can't deserialize unknown %s type %s",
      ViewFn.class.getSimpleName(),
      spec.getUrn());
  return (ViewFn<?, ?>)
      SerializableUtils.deserializeFromByteArray(
          spec.getPayload().toByteArray(), "Custom ViewFn");
}
 

public Twister2BoundedSource(
    BoundedSource<T> boundedSource, Twister2TranslationContext context, PipelineOptions options) {
  source = boundedSource;
  this.options = options;
  this.serializedOptions = new SerializablePipelineOptions(options).toString();
  SdkComponents components = SdkComponents.create();
  components.registerEnvironment(
      Environments.createOrGetDefaultEnvironment(options.as(PortablePipelineOptions.class)));
  RunnerApi.FunctionSpec sourceProto = ReadTranslation.toProto(source, components);
  sourceBytes = sourceProto.getPayload().toByteArray();
}
 

@Override
public RunnerApi.FunctionSpec translate(
    AppliedPTransform<?, ?, ParDoSingle<?, ?>> transform, SdkComponents components)
    throws IOException {
  RunnerApi.ParDoPayload payload = payloadForParDoSingle(transform, components);

  return RunnerApi.FunctionSpec.newBuilder()
      .setUrn(PAR_DO_TRANSFORM_URN)
      .setPayload(payload.toByteString())
      .build();
}
 

private static RunnerApi.FunctionSpec getAccessPattern(SideInputReference sideInputReference) {
  try {
    return RunnerApi.ParDoPayload.parseFrom(
            sideInputReference.transform().getTransform().getSpec().getPayload())
        .getSideInputsMap()
        .get(sideInputReference.localName())
        .getAccessPattern();
  } catch (InvalidProtocolBufferException e) {
    throw new RuntimeException(e);
  }
}
 

@Override
public PTransform<PCollectionList<Long>, PCollection<Long>> getTransform(
    RunnerApi.FunctionSpec spec) {
  return Flatten.pCollections();
}
 

@Nullable
@Override
public RunnerApi.FunctionSpec getSpec() {
  return null;
}
 

@Override
public PTransform<KeyedPCollectionTuple<Long>, PCollection<KV<Long, Iterable<String>>>>
    getTransform(RunnerApi.FunctionSpec spec) {
  return new TestCoGroupByKeyTransform();
}
 

/**
 * Create a Flatten that consumes data from the same PCollection duplicated through two outputs
 * and validates that inputs are flattened together and directed to the output.
 */
@Test
public void testFlattenWithDuplicateInputCollectionProducesMultipleOutputs() throws Exception {
  String pTransformId = "pTransformId";
  String mainOutputId = "101";

  RunnerApi.FunctionSpec functionSpec =
      RunnerApi.FunctionSpec.newBuilder()
          .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)
          .build();
  RunnerApi.PTransform pTransform =
      RunnerApi.PTransform.newBuilder()
          .setSpec(functionSpec)
          .putInputs("inputA", "inputATarget")
          .putInputs("inputAAgain", "inputATarget")
          .putOutputs(mainOutputId, "mainOutputTarget")
          .build();

  List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
  MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
  PCollectionConsumerRegistry consumers =
      new PCollectionConsumerRegistry(
          metricsContainerRegistry, mock(ExecutionStateTracker.class));
  consumers.register(
      "mainOutputTarget",
      pTransformId,
      (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>) mainOutputValues::add);

  new FlattenRunner.Factory<>()
      .createRunnerForPTransform(
          PipelineOptionsFactory.create(),
          null /* beamFnDataClient */,
          null /* beamFnStateClient */,
          null /* beamFnTimerClient */,
          pTransformId,
          pTransform,
          Suppliers.ofInstance("57L")::get,
          Collections.emptyMap(),
          Collections.emptyMap(),
          Collections.emptyMap(),
          consumers,
          null /* startFunctionRegistry */,
          null, /* finishFunctionRegistry */
          null, /* tearDownRegistry */
          null /* addProgressRequestCallback */,
          null /* splitListener */,
          null /* bundleFinalizer */);

  mainOutputValues.clear();
  assertThat(consumers.keySet(), containsInAnyOrder("inputATarget", "mainOutputTarget"));

  assertThat(consumers.getUnderlyingConsumers("inputATarget"), hasSize(2));

  FnDataReceiver<WindowedValue<?>> input = consumers.getMultiplexingConsumer("inputATarget");

  input.accept(WindowedValue.valueInGlobalWindow("A1"));
  input.accept(WindowedValue.valueInGlobalWindow("A2"));

  assertThat(
      mainOutputValues,
      containsInAnyOrder(
          valueInGlobalWindow("A1"),
          valueInGlobalWindow("A1"),
          valueInGlobalWindow("A2"),
          valueInGlobalWindow("A2")));
}
 

/**
 * Create a Flatten that has 4 inputs (inputATarget1, inputATarget2, inputBTarget, inputCTarget)
 * and one output (mainOutput). Validate that inputs are flattened together and directed to the
 * output.
 */
@Test
public void testCreatingAndProcessingDoFlatten() throws Exception {
  String pTransformId = "pTransformId";
  String mainOutputId = "101";

  RunnerApi.FunctionSpec functionSpec =
      RunnerApi.FunctionSpec.newBuilder()
          .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)
          .build();
  RunnerApi.PTransform pTransform =
      RunnerApi.PTransform.newBuilder()
          .setSpec(functionSpec)
          .putInputs("inputA", "inputATarget")
          .putInputs("inputB", "inputBTarget")
          .putInputs("inputC", "inputCTarget")
          .putOutputs(mainOutputId, "mainOutputTarget")
          .build();

  List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
  MetricsContainerStepMap metricsContainerRegistry = new MetricsContainerStepMap();
  PCollectionConsumerRegistry consumers =
      new PCollectionConsumerRegistry(
          metricsContainerRegistry, mock(ExecutionStateTracker.class));
  consumers.register(
      "mainOutputTarget",
      pTransformId,
      (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>) mainOutputValues::add);

  new FlattenRunner.Factory<>()
      .createRunnerForPTransform(
          PipelineOptionsFactory.create(),
          null /* beamFnDataClient */,
          null /* beamFnStateClient */,
          null /* beamFnTimerClient */,
          pTransformId,
          pTransform,
          Suppliers.ofInstance("57L")::get,
          Collections.emptyMap(),
          Collections.emptyMap(),
          Collections.emptyMap(),
          consumers,
          null /* startFunctionRegistry */,
          null, /* finishFunctionRegistry */
          null, /* tearDownRegistry */
          null /* addProgressRequestCallback */,
          null /* splitListener */,
          null /* bundleFinalizer */);

  mainOutputValues.clear();
  assertThat(
      consumers.keySet(),
      containsInAnyOrder("inputATarget", "inputBTarget", "inputCTarget", "mainOutputTarget"));

  consumers.getMultiplexingConsumer("inputATarget").accept(valueInGlobalWindow("A1"));
  consumers.getMultiplexingConsumer("inputATarget").accept(valueInGlobalWindow("A2"));
  consumers.getMultiplexingConsumer("inputBTarget").accept(valueInGlobalWindow("B"));
  consumers.getMultiplexingConsumer("inputCTarget").accept(valueInGlobalWindow("C"));
  assertThat(
      mainOutputValues,
      contains(
          valueInGlobalWindow("A1"),
          valueInGlobalWindow("A2"),
          valueInGlobalWindow("B"),
          valueInGlobalWindow("C")));

  mainOutputValues.clear();
}
 

public abstract Map<PCollectionView<?>, RunnerApi.FunctionSpec>
getPCollectionViewsToWindowMappingFns();
 

default Map<String, PCollection<?>> apply(
    Pipeline p, String name, RunnerApi.FunctionSpec spec, Map<String, PCollection<?>> inputs) {
  return extractOutputs(
      Pipeline.applyTransform(name, createInput(p, inputs), getTransform(spec)));
}
 

PTransform<InputT, OutputT> getTransform(RunnerApi.FunctionSpec spec); 

public abstract RunnerApi.FunctionSpec accessPattern();