SplittableParDo extends PTransform<PCollection<InputT>, PCollectionTuple> { public static void convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(Pipeline pipeline) { if (ExperimentalOptions.hasExperiment(pipeline.getOptions(), "beam_fn_api_use_deprecated_read") || ExperimentalOptions.hasExperiment(pipeline.getOptions(), "use_deprecated_read")) { convertReadBasedSplittableDoFnsToPrimitiveReads(pipeline); } } private  SplittableParDo(
      DoFn<InputT, OutputT> doFn,
      List<PCollectionView<?>> sideInputs,
      TupleTag<OutputT> mainOutputTag,
      TupleTagList additionalOutputTags,
      Map<TupleTag<?>, Coder<?>> outputTagsToCoders); @SuppressWarnings({"unchecked", "rawtypes"}) static SplittableParDo<InputT, OutputT, ?, ?> forAppliedParDo(
      AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> parDo); @Override PCollectionTuple expand(PCollection<InputT> input); @Override Map<TupleTag<?>, PValue> getAdditionalInputs(); static void convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(Pipeline pipeline); static void convertReadBasedSplittableDoFnsToPrimitiveReads(Pipeline pipeline); static final String SPLITTABLE_PROCESS_URN; static final String SPLITTABLE_GBKIKWI_URN; static final PTransformOverride PRIMITIVE_BOUNDED_READ_OVERRIDE; static final PTransformOverride PRIMITIVE_UNBOUNDED_READ_OVERRIDE; }
@Test public void testConvertIsSkippedWhenUsingDeprecatedRead() { Pipeline sdfRead = Pipeline.create(); sdfRead.apply(Read.from(new FakeBoundedSource())); sdfRead.apply(Read.from(new BoundedToUnboundedSourceAdapter<>(new FakeBoundedSource()))); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(sdfRead); pipeline.traverseTopologically( new Defaults() { @Override public void visitPrimitiveTransform(Node node) { assertThat( node.getTransform(), not(instanceOf(SplittableParDo.PrimitiveBoundedRead.class))); assertThat( node.getTransform(), not(instanceOf(SplittableParDo.PrimitiveUnboundedRead.class))); } }); }
@Test public void testConvertToPrimitiveReadsHappen() { PipelineOptions deprecatedReadOptions = PipelineOptionsFactory.create(); deprecatedReadOptions.setRunner(CrashingRunner.class); ExperimentalOptions.addExperiment( deprecatedReadOptions.as(ExperimentalOptions.class), "use_deprecated_read"); Pipeline pipeline = Pipeline.create(deprecatedReadOptions); pipeline.apply(Read.from(new FakeBoundedSource())); pipeline.apply(Read.from(new BoundedToUnboundedSourceAdapter<>(new FakeBoundedSource()))); SplittableParDo.convertReadBasedSplittableDoFnsToPrimitiveReadsIfNecessary(pipeline); AtomicBoolean sawPrimitiveBoundedRead = new AtomicBoolean(); AtomicBoolean sawPrimitiveUnboundedRead = new AtomicBoolean(); pipeline.traverseTopologically( new Defaults() { @Override public CompositeBehavior enterCompositeTransform(Node node) { assertThat(node.getTransform(), not(instanceOf(Read.Bounded.class))); assertThat(node.getTransform(), not(instanceOf(Read.Unbounded.class))); return super.enterCompositeTransform(node); } @Override public void visitPrimitiveTransform(Node node) { if (node.getTransform() instanceof SplittableParDo.PrimitiveBoundedRead) { sawPrimitiveBoundedRead.set(true); } else if (node.getTransform() instanceof SplittableParDo.PrimitiveUnboundedRead) { sawPrimitiveUnboundedRead.set(true); } } }); assertTrue(sawPrimitiveBoundedRead.get()); assertTrue(sawPrimitiveUnboundedRead.get()); }