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> 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 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(); } }
public static Coder<?> fromProto(RunnerApi.Coder protoCoder, RehydratedComponents components) throws IOException { String coderSpecUrn = protoCoder.getSpec().getSpec().getUrn(); if (coderSpecUrn.equals(JAVA_SERIALIZED_CODER_URN)) { return fromCustomCoder(protoCoder); } return fromKnownCoder(protoCoder, components); }
public static WindowedValueCoderComponents getWindowedValueCoderComponents(Coder coder) { checkArgument(WINDOWED_VALUE_CODER_URN.equals(coder.getSpec().getSpec().getUrn())); return new AutoValue_ModelCoders_WindowedValueCoderComponents( coder.getComponentCoderIds(0), coder.getComponentCoderIds(1)); }
/** Returns the URN for the transform if it is known, otherwise {@code null}. */ @Nullable public static String urnForTransformOrNull(RunnerApi.PTransform transform) { return transform.getSpec() == null ? null : transform.getSpec().getUrn(); }
private void urnNotFound( String id, RunnerApi.Pipeline pipeline, FlinkStreamingPortablePipelineTranslator.TranslationContext context) { throw new IllegalArgumentException( String.format( "Unknown type of URN %s for PTransform with id %s.", pipeline.getComponents().getTransformsOrThrow(id).getSpec().getUrn(), id)); }
/** Returns true if the PTransform node for the given input PCollection can be fused across. */ public static boolean canFuse( PTransformNode transformNode, Environment environment, PCollectionNode candidate, Collection<PCollectionNode> stagePCollections, QueryablePipeline pipeline) { return URN_FUSIBILITY_CHECKERS .getOrDefault(transformNode.getTransform().getSpec().getUrn(), DEFAULT_FUSIBILITY_CHECKER) .canFuse(transformNode, environment, candidate, stagePCollections, pipeline); }
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"); }
@Override public FlinkPortablePipelineTranslator.Executor translate( StreamingTranslationContext context, RunnerApi.Pipeline pipeline) { QueryablePipeline p = QueryablePipeline.forTransforms( pipeline.getRootTransformIdsList(), pipeline.getComponents()); for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) { urnToTransformTranslator .getOrDefault(transform.getTransform().getSpec().getUrn(), this::urnNotFound) .translate(transform.getId(), pipeline, context); } return context; }
@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 void urnNotFound( PTransformNode transform, RunnerApi.Pipeline pipeline, BatchTranslationContext context) { throw new IllegalArgumentException( String.format( "Unknown type of URN %s for PTransform with id %s.", transform.getTransform().getSpec().getUrn(), transform.getId())); }
/** 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); }
public static RunnerApi.PCollection getMainInput( RunnerApi.PTransform ptransform, Components components) throws IOException { checkArgument( ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN), "Unexpected payload type %s", ptransform.getSpec().getUrn()); return components.getPcollectionsOrThrow( ptransform.getInputsOrThrow(getMainInputName(ptransform))); }
/** The URN for this transform, if standardized. */ @Nullable public String getUrn() { return getSpec() == null ? null : getSpec().getUrn(); }
private RunnerApi.Pipeline makeKnownUrnsPrimitives( RunnerApi.Pipeline pipeline, Set<String> knownUrns) { RunnerApi.Pipeline.Builder trimmedPipeline = pipeline.toBuilder(); for (String ptransformId : pipeline.getComponents().getTransformsMap().keySet()) { if (knownUrns.contains( pipeline.getComponents().getTransformsOrThrow(ptransformId).getSpec().getUrn())) { LOG.debug("Removing descendants of known PTransform {}" + ptransformId); removeDescendants(trimmedPipeline, ptransformId); } } return trimmedPipeline.build(); }
private void urnNotFound( String id, RunnerApi.Pipeline pipeline, FlinkStreamingPortablePipelineTranslator.TranslationContext context) { throw new IllegalArgumentException( String.format( "Unknown type of URN %s for PTransform with id %s.", pipeline.getComponents().getTransformsOrThrow(id).getSpec().getUrn(), id)); }
@Override public void translate(StreamingTranslationContext context, RunnerApi.Pipeline pipeline) { QueryablePipeline p = QueryablePipeline.forTransforms( pipeline.getRootTransformIdsList(), pipeline.getComponents()); for (PipelineNode.PTransformNode transform : p.getTopologicallyOrderedTransforms()) { urnToTransformTranslator .getOrDefault(transform.getTransform().getSpec().getUrn(), this::urnNotFound) .translate(transform.getId(), pipeline, context); } }
/** 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()); } }
public static KvCoderComponents getKvCoderComponents(Coder coder) { checkArgument(KV_CODER_URN.equals(coder.getSpec().getSpec().getUrn())); return new AutoValue_ModelCoders_KvCoderComponents( coder.getComponentCoderIds(0), coder.getComponentCoderIds(1)); }