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(); } }
/** Returns the name of the main input of the ptransform. */ private static String getMainInputName( RunnerApi.PTransformOrBuilder ptransform, RunnerApi.ParDoPayload payload) { return Iterables.getOnlyElement( Sets.difference( ptransform.getInputsMap().keySet(), Sets.union(payload.getSideInputsMap().keySet(), payload.getTimerSpecsMap().keySet()))); }
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; }
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())); } }