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(); } }
ParDoPayload parDoPayload = ParDoPayload.parseFrom(protoTransform.getSpec().getPayload()); for (PCollectionView<?> view : parDo.getSideInputs()) { SideInput sideInput = parDoPayload.getSideInputsOrThrow(view.getTagInternal().getId());
assertEquals( PTransformTranslation.SPLITTABLE_PROCESS_KEYED_URN, spkTransform.getSpec().getUrn()); ParDoPayload payload = ParDoPayload.parseFrom(spkTransform.getSpec().getPayload()); assertThat( ParDoTranslation.doFnAndMainOutputTagFromProto(payload.getDoFn()).getDoFn(),
ParDoPayload payload = ParDoPayload.parseFrom(parDo.getTransform().getSpec().getPayload()); if (Maps.filterKeys( parDo.getTransform().getInputsMap(), s -> payload.getTimerSpecsMap().containsKey(s))
public static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> application) throws IOException { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.MultiOutput) { return ((ParDo.MultiOutput<?, ?>) transform).getSideInputs(); } SdkComponents sdkComponents = SdkComponents.create(application.getPipeline().getOptions()); RunnerApi.PTransform parDoProto = PTransformTranslation.toProto(application, sdkComponents); ParDoPayload payload = ParDoPayload.parseFrom(parDoProto.getSpec().getPayload()); List<PCollectionView<?>> views = new ArrayList<>(); RehydratedComponents components = RehydratedComponents.forComponents(sdkComponents.toComponents()); for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) { String sideInputTag = sideInputEntry.getKey(); RunnerApi.SideInput sideInput = sideInputEntry.getValue(); PCollection<?> originalPCollection = checkNotNull( (PCollection<?>) application.getInputs().get(new TupleTag<>(sideInputTag)), "no input with tag %s", sideInputTag); views.add( PCollectionViewTranslation.viewFromProto( sideInput, sideInputTag, originalPCollection, parDoProto, components)); } return views; }
public static TupleTagList getAdditionalOutputTags(AppliedPTransform<?, ?, ?> application) throws IOException { PTransform<?, ?> transform = application.getTransform(); if (transform instanceof ParDo.MultiOutput) { return ((ParDo.MultiOutput<?, ?>) transform).getAdditionalOutputTags(); } RunnerApi.PTransform protoTransform = PTransformTranslation.toProto( application, SdkComponents.create(application.getPipeline().getOptions())); ParDoPayload payload = ParDoPayload.parseFrom(protoTransform.getSpec().getPayload()); TupleTag<?> mainOutputTag = getMainOutputTag(payload); Set<String> outputTags = Sets.difference( protoTransform.getOutputsMap().keySet(), Collections.singleton(mainOutputTag.getId())); ArrayList<TupleTag<?>> additionalOutputTags = new ArrayList<>(); for (String outputTag : outputTags) { additionalOutputTags.add(new TupleTag<>(outputTag)); } return TupleTagList.of(additionalOutputTags); }
private static void validateParDo(String id, PTransform transform, Components components) throws Exception { ParDoPayload payload = ParDoPayload.parseFrom(transform.getSpec().getPayload()); // side_inputs for (String sideInputId : payload.getSideInputsMap().keySet()) { checkArgument( transform.containsInputs(sideInputId), "Transform %s side input %s is not listed in the transform's inputs", id, sideInputId); } // TODO: Validate state_specs and timer_specs if (!payload.getRestrictionCoderId().isEmpty()) { checkArgument(components.containsCoders(payload.getRestrictionCoderId())); } }
private static String parDoExtractor(PTransform pTransform) throws InvalidProtocolBufferException { return ParDoPayload.parseFrom(pTransform.getSpec().getPayload()).getDoFn().getEnvironmentId(); }
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()); }
/** 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); }