private Set<String> getLocalSideInputNames(PTransform transform) { if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { try { return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getSideInputsMap().keySet(); } catch (InvalidProtocolBufferException e) { throw new RuntimeException(e); } } else { return Collections.emptySet(); } }
private Set<String> getLocalTimerNames(PTransform transform) { if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { try { return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getTimerSpecsMap().keySet(); } catch (InvalidProtocolBufferException e) { throw new RuntimeException(e); } } else { return Collections.emptySet(); } }
private Set<String> getLocalUserStateNames(PTransform transform) { if (PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) { try { return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getStateSpecsMap().keySet(); } catch (InvalidProtocolBufferException e) { throw new RuntimeException(e); } } else { return Collections.emptySet(); } }
private static String parDoExtractor(PTransform pTransform) throws InvalidProtocolBufferException { return ParDoPayload.parseFrom(pTransform.getSpec().getPayload()).getDoFn().getEnvironmentId(); }
private static void validateAssignWindows(String id, PTransform transform, Components components) throws Exception { WindowIntoPayload.parseFrom(transform.getSpec().getPayload()); }
public static DoFnAndMainOutput doFnAndMainOutputTagFromProto(SdkFunctionSpec fnSpec) { checkArgument( fnSpec.getSpec().getUrn().equals(CUSTOM_JAVA_DO_FN_URN), "Expected %s to be %s with URN %s, but URN was %s", DoFn.class.getSimpleName(), FunctionSpec.class.getSimpleName(), CUSTOM_JAVA_DO_FN_URN, fnSpec.getSpec().getUrn()); byte[] serializedFn = fnSpec.getSpec().getPayload().toByteArray(); return (DoFnAndMainOutput) SerializableUtils.deserializeFromByteArray(serializedFn, "Custom DoFn And Main Output tag"); }
public static BoundedSource<?> boundedSourceFromProto(ReadPayload payload) throws InvalidProtocolBufferException { checkArgument(payload.getIsBounded().equals(IsBounded.Enum.BOUNDED)); return (BoundedSource<?>) SerializableUtils.deserializeFromByteArray( payload.getSource().getSpec().getPayload().toByteArray(), "BoundedSource"); }
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()); }
public static UnboundedSource<?, ?> unboundedSourceFromProto(ReadPayload payload) throws InvalidProtocolBufferException { checkArgument(payload.getIsBounded().equals(IsBounded.Enum.UNBOUNDED)); return (UnboundedSource<?, ?>) SerializableUtils.deserializeFromByteArray( payload.getSource().getSpec().getPayload().toByteArray(), "UnboundedSource"); }
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()); }
public RawWriteFiles( RunnerApi.PTransform protoTransform, RehydratedComponents rehydratedComponents) throws IOException { this.rehydratedComponents = rehydratedComponents; this.protoTransform = protoTransform; this.spec = protoTransform.getSpec(); this.payload = RunnerApi.WriteFilesPayload.parseFrom(spec.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); } }
@VisibleForTesting static FileBasedSink<?, ?, ?> sinkFromProto(SdkFunctionSpec sinkProto) throws IOException { checkArgument( sinkProto.getSpec().getUrn().equals(CUSTOM_JAVA_FILE_BASED_SINK_URN), "Cannot extract %s instance from %s with URN %s", FileBasedSink.class.getSimpleName(), FunctionSpec.class.getSimpleName(), sinkProto.getSpec().getUrn()); byte[] serializedSink = sinkProto.getSpec().getPayload().toByteArray(); return (FileBasedSink<?, ?, ?>) SerializableUtils.deserializeFromByteArray( serializedSink, FileBasedSink.class.getSimpleName()); }
private static String windowExtractor(PTransform transform) throws InvalidProtocolBufferException { return WindowIntoPayload.parseFrom(transform.getSpec().getPayload()) .getWindowFn() .getEnvironmentId(); }
private static String combineExtractor(PTransform pTransform) throws InvalidProtocolBufferException { return CombinePayload.parseFrom(pTransform.getSpec().getPayload()) .getCombineFn() .getEnvironmentId(); }
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()); }
private static Coder<?> fromCustomCoder(RunnerApi.Coder protoCoder) throws IOException { return (Coder<?>) SerializableUtils.deserializeFromByteArray( protoCoder.getSpec().getSpec().getPayload().toByteArray(), "Custom Coder Bytes"); } }
/** 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()); } }
/** Returns the name of the main input of the ptransform. */ public static String getMainInputName(RunnerApi.PTransformOrBuilder ptransform) throws IOException { checkArgument( ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN), "Unexpected payload type %s", ptransform.getSpec().getUrn()); ParDoPayload payload = ParDoPayload.parseFrom(ptransform.getSpec().getPayload()); return getMainInputName(ptransform, payload); }
private static void validateCombine(String id, PTransform transform, Components components) throws Exception { CombinePayload payload = CombinePayload.parseFrom(transform.getSpec().getPayload()); checkArgument( components.containsCoders(payload.getAccumulatorCoderId()), "Transform %s uses unknown accumulator coder id %s", payload.getAccumulatorCoderId()); }