Java Code Examples for org.apache.beam.sdk.runners.AppliedPTransform#of()

/**
 * Gets the {@link AppliedPTransform} that has a created {@code PCollection<KV<String, Integer>>}
 * as input.
 */
private AppliedPTransform<?, ?, ?> getAppliedTransform(PTransform pardo) {
  PCollection<KV<String, Integer>> input =
      PCollection.createPrimitiveOutputInternal(
          p,
          WindowingStrategy.globalDefault(),
          IsBounded.BOUNDED,
          KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()));
  input.setName("dummy input");

  PCollection<Integer> output =
      PCollection.createPrimitiveOutputInternal(
          p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of());
  output.setName("dummy output");

  return AppliedPTransform.of("pardo", input.expand(), output.expand(), pardo, p);
}
 

@Test
public void emptyFlattenWithNonEmptyFlatten() {
  AppliedPTransform application =
      AppliedPTransform.of(
          "Flatten",
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Flatten.pCollections(),
          p);

  assertThat(PTransformMatchers.emptyFlatten().matches(application), is(false));
}
 

@Test
public void flattenWithDuplicateInputsNonFlatten() {
  AppliedPTransform application =
      AppliedPTransform
          .<PCollection<Iterable<Integer>>, PCollection<Integer>, Flatten.Iterables<Integer>>of(
              "EmptyFlatten",
              Collections.emptyMap(),
              Collections.singletonMap(
                  new TupleTag<Integer>(),
                  PCollection.createPrimitiveOutputInternal(
                      p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
              /* This isn't actually possible to construct, but for the sake of example */
              Flatten.iterables(),
              p);

  assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false));
}
 

@Test
public void getReplacementTransformGetFn() {
  DoFn<Integer, Long> originalFn = new ToLongFn();
  ParDo.SingleOutput<Integer, Long> originalTransform = ParDo.of(originalFn);
  PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3));
  AppliedPTransform<
          PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>>
      application =
          AppliedPTransform.of(
              "original",
              input.expand(),
              input.apply(originalTransform).expand(),
              originalTransform,
              pipeline);

  PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacementTransform =
      factory.getReplacementTransform(application);
  ParDoSingle<Integer, Long> parDoSingle =
      (ParDoSingle<Integer, Long>) replacementTransform.getTransform();

  assertThat(parDoSingle.getFn(), equalTo(originalTransform.getFn()));
  assertThat(parDoSingle.getFn(), equalTo(originalFn));
}
 

@Test
public void testExtractionDirectFromTransform() throws Exception {
  PCollection<String> input = p.apply(Create.of("hello"));
  WriteFilesResult<Void> output = input.apply(writeFiles);

  AppliedPTransform<PCollection<String>, WriteFilesResult<Void>, WriteFiles<String, Void, String>>
      appliedPTransform =
          AppliedPTransform.of("foo", input.expand(), output.expand(), writeFiles, p);

  assertThat(
      WriteFilesTranslation.isRunnerDeterminedSharding(appliedPTransform),
      equalTo(
          writeFiles.getNumShardsProvider() == null && writeFiles.getComputeNumShards() == null));

  assertThat(
      WriteFilesTranslation.isWindowedWrites(appliedPTransform),
      equalTo(writeFiles.getWindowedWrites()));
  assertThat(
      WriteFilesTranslation.<String, Void, String>getSink(appliedPTransform),
      equalTo(writeFiles.getSink()));
}
 

@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 testRunnerDeterminedSharding() {
  FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);
  options.setRunner(TestFlinkRunner.class);
  options.setFlinkMaster("[auto]");
  options.setParallelism(5);

  TestPipeline p = TestPipeline.fromOptions(options);

  StreamingShardedWriteFactory<Object, Void, Object> factory =
      new StreamingShardedWriteFactory<>(p.getOptions());

  WriteFiles<Object, Void, Object> original = WriteFiles.to(new TestSink(tmpFolder.toString()));
  @SuppressWarnings("unchecked")
  PCollection<Object> objs = (PCollection) p.apply(Create.empty(VoidCoder.of()));
  AppliedPTransform<PCollection<Object>, WriteFilesResult<Void>, WriteFiles<Object, Void, Object>>
      originalApplication =
          AppliedPTransform.of("writefiles", objs.expand(), Collections.emptyMap(), original, p);

  WriteFiles<Object, Void, Object> replacement =
      (WriteFiles<Object, Void, Object>)
          factory.getReplacementTransform(originalApplication).getTransform();

  assertThat(replacement, not(equalTo((Object) original)));
  assertThat(replacement.getNumShardsProvider().get(), is(10));
}
 

@Test
public void emptyFlattenWithNonFlatten() {
  AppliedPTransform application =
      AppliedPTransform
          .<PCollection<Iterable<Integer>>, PCollection<Integer>, Flatten.Iterables<Integer>>of(
              "EmptyFlatten",
              Collections.emptyMap(),
              Collections.singletonMap(
                  new TupleTag<Integer>(),
                  PCollection.createPrimitiveOutputInternal(
                      p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
              /* This isn't actually possible to construct, but for the sake of example */
              Flatten.iterables(),
              p);

  assertThat(PTransformMatchers.emptyFlatten().matches(application), is(false));
}
 

private static AppliedPTransform<?, ?, ?> rawPTransformWithNullSpec(Pipeline pipeline) {
  PTransformTranslation.RawPTransform<PBegin, PDone> rawPTransform =
      new PTransformTranslation.RawPTransform<PBegin, PDone>() {
        @Override
        public String getUrn() {
          return "fake/urn";
        }

        @Nullable
        @Override
        public RunnerApi.FunctionSpec getSpec() {
          return null;
        }
      };
  return AppliedPTransform.<PBegin, PDone, PTransform<PBegin, PDone>>of(
      "RawPTransformWithNoSpec",
      pipeline.begin().expand(),
      PDone.in(pipeline).expand(),
      rawPTransform,
      pipeline);
}
 

private void processElement(final TestUnboundedSource<String> source) throws Exception {
  final EvaluationContext context =
      EvaluationContext.create(
          MockClock.fromInstant(Instant.now()),
          CloningBundleFactory.create(),
          DirectGraph.create(
              emptyMap(), emptyMap(), LinkedListMultimap.create(), emptySet(), emptyMap()),
          emptySet(),
          Executors.newCachedThreadPool());
  final UnboundedReadEvaluatorFactory factory =
      new UnboundedReadEvaluatorFactory(context, options);

  final Read.Unbounded<String> unbounded = Read.from(source);
  final Pipeline pipeline = Pipeline.create(options);
  final PCollection<String> pCollection = pipeline.apply(unbounded);
  final AppliedPTransform<PBegin, PCollection<String>, Read.Unbounded<String>> application =
      AppliedPTransform.of(
          "test",
          new HashMap<>(),
          singletonMap(new TupleTag(), pCollection),
          unbounded,
          pipeline);
  final TransformEvaluator<UnboundedSourceShard<String, TestCheckpointMark>> evaluator =
      factory.forApplication(application, null);
  final UnboundedSource.UnboundedReader<String> reader = source.createReader(options, null);
  final UnboundedSourceShard<String, TestCheckpointMark> shard =
      UnboundedSourceShard.of(source, new NeverDeduplicator(), reader, null);
  final WindowedValue<UnboundedSourceShard<String, TestCheckpointMark>> value =
      WindowedValue.of(
          shard, BoundedWindow.TIMESTAMP_MAX_VALUE, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING);
  TestUnboundedSource.readerClosedCount = 0;
  evaluator.processElement(value);
}
 

@Test
public void nonAdditionalInputsWithMultipleNonAdditionalInputsSucceeds() {
  Map<TupleTag<?>, PValue> allInputs = new HashMap<>();
  PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3));
  allInputs.put(new TupleTag<Integer>() {}, mainInts);
  PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of()));
  allInputs.put(new TupleTag<Void>() {}, voids);
  AppliedPTransform<PInput, POutput, TestTransform> transform =
      AppliedPTransform.of(
          "additional-free", allInputs, Collections.emptyMap(), new TestTransform(), pipeline);

  assertThat(
      TransformInputs.nonAdditionalInputs(transform),
      Matchers.containsInAnyOrder(voids, mainInts));
}
 

private AppliedPTransform<?, ?, ?> appliedPTransform(
    String fullName, PTransform<PInput, POutput> transform, Pipeline p) {
  PInput input = mock(PInput.class);
  when(input.getPipeline()).thenReturn(p);
  return AppliedPTransform.of(
      fullName, Collections.emptyMap(), Collections.emptyMap(), transform, p);
}
 

private PCollection<String> applySplittableParDo(
    String name, PCollection<Integer> input, DoFn<Integer, String> fn) {
  ParDo.MultiOutput<Integer, String> multiOutput =
      ParDo.of(fn).withOutputTags(MAIN_OUTPUT_TAG, TupleTagList.empty());
  PCollectionTuple output = multiOutput.expand(input);
  output.get(MAIN_OUTPUT_TAG).setName("main");
  AppliedPTransform<PCollection<Integer>, PCollectionTuple, ?> transform =
      AppliedPTransform.of("ParDo", input.expand(), output.expand(), multiOutput, pipeline);
  return input.apply(name, SplittableParDo.forAppliedParDo(transform)).get(MAIN_OUTPUT_TAG);
}
 

@Test
public void getReplacementTransformGetSideInputs() {
  PCollectionView<Long> sideLong =
      pipeline
          .apply("LongSideInputVals", Create.of(-1L, -2L, -4L))
          .apply("SideLongView", Sum.longsGlobally().asSingletonView());
  PCollectionView<List<String>> sideStrings =
      pipeline
          .apply("StringSideInputVals", Create.of("foo", "bar", "baz"))
          .apply("SideStringsView", View.asList());
  ParDo.SingleOutput<Integer, Long> originalTransform =
      ParDo.of(new ToLongFn()).withSideInputs(sideLong, sideStrings);

  PCollection<? extends Integer> input = pipeline.apply(Create.of(1, 2, 3));
  AppliedPTransform<
          PCollection<? extends Integer>, PCollection<Long>, ParDo.SingleOutput<Integer, Long>>
      application =
          AppliedPTransform.of(
              "original",
              input.expand(),
              input.apply(originalTransform).expand(),
              originalTransform,
              pipeline);

  PTransformReplacement<PCollection<? extends Integer>, PCollection<Long>> replacementTransform =
      factory.getReplacementTransform(application);
  ParDoSingle<Integer, Long> parDoSingle =
      (ParDoSingle<Integer, Long>) replacementTransform.getTransform();
  assertThat(parDoSingle.getSideInputs().values(), containsInAnyOrder(sideStrings, sideLong));
}
 

@Test
public void registerTransformNullComponents() throws IOException {
  Create.Values<Integer> create = Create.of(1, 2, 3);
  PCollection<Integer> pt = pipeline.apply(create);
  String userName = "my_transform/my_nesting";
  AppliedPTransform<?, ?, ?> transform =
      AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline);
  thrown.expect(NullPointerException.class);
  thrown.expectMessage("child nodes may not be null");
  components.registerPTransform(transform, null);
}
 

@Test
public void registerTransformNoChildren() throws IOException {
  Create.Values<Integer> create = Create.of(1, 2, 3);
  PCollection<Integer> pt = pipeline.apply(create);
  String userName = "my_transform/my_nesting";
  AppliedPTransform<?, ?, ?> transform =
      AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline);
  String componentName = components.registerPTransform(transform, Collections.emptyList());
  assertThat(componentName, equalTo(userName));
  assertThat(components.getExistingPTransformId(transform), equalTo(componentName));
}
 

@Test
public void emptyFlattenWithEmptyFlatten() {
  AppliedPTransform application =
      AppliedPTransform.of(
          "EmptyFlatten",
          Collections.emptyMap(),
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Flatten.pCollections(),
          p);

  assertThat(PTransformMatchers.emptyFlatten().matches(application), is(true));
}
 

@Test
public void getMainInputSingleOutputSideInputs() {
  AppliedPTransform<PCollection<Long>, ?, ?> application =
      AppliedPTransform.of(
          "application",
          ImmutableMap.<TupleTag<?>, PValue>builder()
              .put(new TupleTag<Long>(), mainInput)
              .put(sideInput.getTagInternal(), sideInput.getPCollection())
              .build(),
          Collections.singletonMap(new TupleTag<Long>(), output),
          ParDo.of(new TestDoFn()).withSideInputs(sideInput),
          pipeline);
  PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application);
  assertThat(input, equalTo(mainInput));
}
 

private AppliedPTransform<?, ?, ?> appliedWrite(WriteFiles<Integer, Void, Integer> write) {
  return AppliedPTransform.of(
      "WriteFiles", Collections.emptyMap(), Collections.emptyMap(), write, p);
}
 

private static AppliedPTransform<?, ?, ?> generateSequence(Pipeline pipeline) {
  GenerateSequence sequence = GenerateSequence.from(0);
  PCollection<Long> pcollection = pipeline.apply(sequence);
  return AppliedPTransform.of(
      "Count", pipeline.begin().expand(), pcollection.expand(), sequence, pipeline);
}