@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) { try { return PTransformReplacement.of( transform.getPipeline().begin(), new DirectTestStream<>(runner, TestStreamTranslation.getTestStream(transform))); } catch (IOException exc) { throw new RuntimeException( String.format( "Transform could not be converted to %s", TestStream.class.getSimpleName()), exc); } }
@Override public PTransformReplacement<PCollectionList<T>, PCollection<T>> getReplacementTransform( AppliedPTransform<PCollectionList<T>, PCollection<T>, PCollections<T>> transform) { checkArgument( transform.getInputs().isEmpty(), "Unexpected nonempty input %s for %s", transform.getInputs(), getClass().getSimpleName()); return PTransformReplacement.of( PCollectionList.empty(transform.getPipeline()), new CreateEmptyFromList<T>()); }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) { PBegin input = PBegin.in(transform.getPipeline()); BoundedSource<T> source; try { source = ReadTranslation.boundedSourceFromTransform(transform); } catch (IOException e) { throw new RuntimeException(e); } return PTransformReplacement.of(input, bounded(source)); }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Read.Bounded<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new StreamingBoundedRead<>(transform.getTransform())); }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Read.Unbounded<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new StreamingUnboundedRead<>(transform.getTransform())); }
@Override public PTransformReplacement<PBegin, PCollection<PubsubMessage>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<PubsubMessage>, PubsubUnboundedSource> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new StreamingPubsubIORead(transform.getTransform())); }
@Override public PTransformReplacement<PCollectionList<T>, PCollection<T>> getReplacementTransform( AppliedPTransform<PCollectionList<T>, PCollection<T>, PCollections<T>> transform) { return PTransformReplacement.of( getInput(transform.getInputs(), transform.getPipeline()), new FlattenWithoutDuplicateInputs<>()); }
private static <T> ReadPayload getReadPayload( AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return ReadPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()); }
private static <T> ParDoPayload getParDoPayload(AppliedPTransform<?, ?, ?> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); RunnerApi.PTransform parDoPTransform = PTransformTranslation.toProto(transform, Collections.emptyList(), components); return ParDoPayload.parseFrom(parDoPTransform.getSpec().getPayload()); }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Create.Values<T>> transform) { Create.Values<T> original = transform.getTransform(); PCollection<T> output = (PCollection) Iterables.getOnlyElement(transform.getOutputs().values()); return PTransformReplacement.of( transform.getPipeline().begin(), new StreamingFnApiCreate<>(original, output)); }
private static <T, DestinationT> WriteFilesPayload getWriteFilesPayload( AppliedPTransform< PCollection<T>, WriteFilesResult<DestinationT>, ? extends PTransform<PCollection<T>, WriteFilesResult<DestinationT>>> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return WriteFilesPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()); }
public static PCollection.IsBounded sourceIsBounded(AppliedPTransform<?, ?, ?> transform) { try { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return PCollectionTranslation.fromProto( ReadPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()) .getIsBounded()); } catch (IOException e) { throw new RuntimeException("Internal error determining boundedness of Read", e); } }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Values<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), new PrimitiveCreate<>( transform.getTransform(), ((PCollection<T>) Iterables.getOnlyElement(transform.getOutputs().values())) .getCoder())); }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Read.Bounded<T>> transform) { return PTransformReplacement.of( transform.getPipeline().begin(), JavaReadViaImpulse.bounded(transform.getTransform().getSource())); }
/** Extracts an {@link ExecutableStagePayload} from the given transform. */ public static ExecutableStagePayload getExecutableStagePayload( AppliedPTransform<?, ?, ?> appliedTransform) throws IOException { RunnerApi.PTransform transform = PTransformTranslation.toProto( appliedTransform, SdkComponents.create(appliedTransform.getPipeline().getOptions())); checkArgument(ExecutableStage.URN.equals(transform.getSpec().getUrn())); return ExecutableStagePayload.parseFrom(transform.getSpec().getPayload()); } }
@Override public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<T>, Create.Values<T>> transform) { return PTransformReplacement.of(transform.getPipeline().begin(), new EmptyFlatten<T>()); }
@Override public PTransformReplacement<PBegin, PCollection<Long>> getReplacementTransform( AppliedPTransform<PBegin, PCollection<Long>, GenerateSequence> transform) { return PTransformReplacement.of(transform.getPipeline().begin(), Create.of(0L)); }
private static Coder<?> getAccumulatorCoder(AppliedPTransform<?, ?, ?> transform) throws IOException { SdkComponents sdkComponents = SdkComponents.create(transform.getPipeline().getOptions()); String id = getCombinePayload(transform, sdkComponents) .map(CombinePayload::getAccumulatorCoderId) .orElseThrow(() -> new IOException("Transform does not contain an AccumulatorCoder")); Components components = sdkComponents.toComponents(); return CoderTranslation.fromProto( components.getCodersOrThrow(id), RehydratedComponents.forComponents(components)); }
private static <K, InputT, AccumT> Coder<AccumT> extractAccumulatorCoder( GlobalCombineFn<InputT, AccumT, ?> combineFn, AppliedPTransform<PCollection<KV<K, InputT>>, ?, Combine.PerKey<K, InputT, ?>> transform) throws CannotProvideCoderException { @SuppressWarnings("unchecked") PCollection<KV<K, InputT>> mainInput = (PCollection<KV<K, InputT>>) Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(transform)); KvCoder<K, InputT> inputCoder = (KvCoder<K, InputT>) mainInput.getCoder(); return AppliedCombineFn.withInputCoder( combineFn, transform.getPipeline().getCoderRegistry(), inputCoder, transform.getTransform().getSideInputs(), ((PCollection<?>) Iterables.getOnlyElement(transform.getOutputs().values())) .getWindowingStrategy()) .getAccumulatorCoder(); }
@Test public void toAndFromProto() throws IOException { SdkComponents components = SdkComponents.create(spec.getTransform().getPipeline().getOptions()); RunnerApi.PTransform converted = convert(spec, components); Components protoComponents = components.toComponents(); // Sanity checks assertThat(converted.getInputsCount(), equalTo(spec.getTransform().getInputs().size())); assertThat(converted.getOutputsCount(), equalTo(spec.getTransform().getOutputs().size())); assertThat(converted.getSubtransformsCount(), equalTo(spec.getChildren().size())); assertThat(converted.getUniqueName(), equalTo(spec.getTransform().getFullName())); for (PValue inputValue : spec.getTransform().getInputs().values()) { PCollection<?> inputPc = (PCollection<?>) inputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(inputPc)); } for (PValue outputValue : spec.getTransform().getOutputs().values()) { PCollection<?> outputPc = (PCollection<?>) outputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(outputPc)); } }