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))); }
/** 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); }
/** Create a user state reference from a UserStateId proto and components. */ public static UserStateReference fromUserStateId( UserStateId userStateId, RunnerApi.Components components) { PTransform transform = components.getTransformsOrThrow(userStateId.getTransformId()); String mainInputCollectionId; try { mainInputCollectionId = transform.getInputsOrThrow(ParDoTranslation.getMainInputName(transform)); } catch (IOException e) { throw new RuntimeException(e); } return UserStateReference.of( PipelineNode.pTransform(userStateId.getTransformId(), transform), userStateId.getLocalName(), PipelineNode.pCollection( mainInputCollectionId, components.getPcollectionsOrThrow(mainInputCollectionId))); }
ParDoTranslation.getMainInput(protoTransform, components), equalTo(components.getPcollectionsOrThrow(mainInputId))); assertThat(ParDoTranslation.getMainInputName(protoTransform), equalTo("mainInputName"));