@Override public Map<TupleTag<?>, PValue> getAdditionalInputs() { Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>(); for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) { try { additionalInputs.put( new TupleTag<>(sideInputEntry.getKey()), rehydratedComponents.getPCollection( protoTransform.getInputsOrThrow(sideInputEntry.getKey()))); } catch (IOException exc) { throw new IllegalStateException( String.format( "Could not find input with name %s for %s transform", sideInputEntry.getKey(), WriteFiles.class.getSimpleName())); } } return additionalInputs; }
/** * Returns the {@link SideInputReference SideInputReferences} that the provided transform consumes * as side inputs. */ public Collection<SideInputReference> getSideInputs(PTransformNode transform) { return getLocalSideInputNames(transform.getTransform()) .stream() .map( localName -> { String transformId = transform.getId(); PTransform transformProto = components.getTransformsOrThrow(transformId); String collectionId = transform.getTransform().getInputsOrThrow(localName); PCollection collection = components.getPcollectionsOrThrow(collectionId); return SideInputReference.of( PipelineNode.pTransform(transformId, transformProto), localName, PipelineNode.pCollection(collectionId, collection)); }) .collect(Collectors.toSet()); }
components .getTransformsOrThrow(sideInput.getKey().getTransformId()) .getInputsOrThrow(sideInput.getKey().getLocalName()); DataStream<Object> sideInputStream = context.getDataStreamOrThrow(collectionId); TypeInformation<Object> tpe = sideInputStream.getType(); components .getTransformsOrThrow(sideInput.getKey().getTransformId()) .getInputsOrThrow(sideInput.getKey().getLocalName()); DataStream<WindowedValue<?>> sideInputStream = context.getDataStreamOrThrow(collectionId);
components .getTransformsOrThrow(sideInputId.getTransformId()) .getInputsOrThrow(sideInputId.getLocalName());
components .getTransformsOrThrow(sideInputId.getTransformId()) .getInputsOrThrow(sideInputId.getLocalName());
components .getTransformsOrThrow(sideInput.getKey().getTransformId()) .getInputsOrThrow(sideInput.getKey().getLocalName()); DataStream<Object> sideInputStream = context.getDataStreamOrThrow(collectionId); TypeInformation<Object> tpe = sideInputStream.getType(); RunnerApi.PTransform pTransform = components.getTransformsOrThrow(sideInput.getKey().getTransformId()); String collectionId = pTransform.getInputsOrThrow(sideInput.getKey().getLocalName()); DataStream<WindowedValue<?>> sideInputStream = context.getDataStreamOrThrow(collectionId);
.map(Map.Entry::getKey) .collect(Collectors.toSet())); String sideInputCollectionId = parDoTransform.getInputsOrThrow(sideInputLocalName); PCollectionNode sideInput = PipelineNode.pCollection(
components .getTransformsOrThrow(sideInputId.getTransformId()) .getInputsOrThrow(sideInputId.getLocalName()); RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow(
components .getTransformsOrThrow(sideInputId.getTransformId()) .getInputsOrThrow(sideInputId.getLocalName()); RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow(
String pcollectionId = transform.getInputsOrThrow(inputId); checkArgument( components.containsPcollections(pcollectionId),
public Collection<UserStateReference> getUserStates(PTransformNode transform) { return getLocalUserStateNames(transform.getTransform()) .stream() .map( localName -> { String transformId = transform.getId(); PTransform transformProto = components.getTransformsOrThrow(transformId); // Get the main input PCollection id. String collectionId = transform .getTransform() .getInputsOrThrow( Iterables.getOnlyElement( Sets.difference( transform.getTransform().getInputsMap().keySet(), ImmutableSet.builder() .addAll(getLocalSideInputNames(transformProto)) .addAll(getLocalTimerNames(transformProto)) .build()))); PCollection collection = components.getPcollectionsOrThrow(collectionId); return UserStateReference.of( PipelineNode.pTransform(transformId, transformProto), localName, PipelineNode.pCollection(collectionId, collection)); }) .collect(Collectors.toSet()); }
/** 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))); }
/** Create a side input reference from a SideInputId proto and components. */ public static SideInputReference fromSideInputId( SideInputId sideInputId, RunnerApi.Components components) { String transformId = sideInputId.getTransformId(); String localName = sideInputId.getLocalName(); String collectionId = components.getTransformsOrThrow(transformId).getInputsOrThrow(localName); PTransform transform = components.getTransformsOrThrow(transformId); PCollection collection = components.getPcollectionsOrThrow(collectionId); return SideInputReference.of( PipelineNode.pTransform(transformId, transform), localName, PipelineNode.pCollection(collectionId, collection)); }
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))); }