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()); }
Components.Builder componentBuilder = stage.getComponents().toBuilder(); stage .getComponents() .getPcollectionsMap() .entrySet() .stream()
public static WindowingStrategy getWindowingStrategy( String pCollectionId, RunnerApi.Components components) { RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow( components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId()); final WindowingStrategy<?, ?> windowingStrategy; try { return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.forComponents(components)); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException( String.format( "Unable to hydrate windowing strategy %s for %s.", windowingStrategyProto, pCollectionId), e); } } }
.build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introducedIds.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap().entrySet(), hasItems( result
@Test public void flattenWithHeterogeneousInputsSingleEnvOutputPartiallyMaterialized() { Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("pyImpulse", components.getTransformsOrThrow("pyImpulse")), PipelineNode.pTransform("goImpulse", components.getTransformsOrThrow("goImpulse"))));
.build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introduced.getTransform().getInputsMap().values().toArray(new String[0]))); result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(stageOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap(), hasEntry(introduced.getId(), introduced.getTransform()));
spk.getSpec().getUrn()); Components.Builder newComponents = Components.newBuilder(); newComponents.putAllCoders(components.getCodersMap());
.addRootTransformIds(flattenTransform.getUniqueName()) .setComponents( Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) s -> s.getComponents() .getTransformsOrThrow(flattenTransform.getUniqueName()) .getInputsMap() .values()
.build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut) result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introduced.getTransform().getInputsMap().values().toArray(new String[0]))); containsInAnyOrder(introducedOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getPcollectionsMap().keySet(), hasItems(introducedOutputs.toArray(new String[0]))); assertThat( result.getDeduplicatedComponents().getTransformsMap(), hasEntry(introduced.getId(), introduced.getTransform()));
.build(); RunnerApi.Components components = Components.newBuilder() .putTransforms("one", one) .putPcollections(oneOut.getUniqueName(), oneOut)
PCollectionNode inputCollection = PipelineNode.pCollection( inputPCollectionId, components.getPcollectionsOrThrow(inputPCollectionId)); DataSet<WindowedValue<KV<K, V>>> inputDataSet = context.getDataSetOrThrow(inputPCollectionId); RunnerApi.WindowingStrategy windowingStrategyProto = pipeline .getComponents() .getWindowingStrategiesOrThrow( pipeline .getComponents() .getPcollectionsOrThrow(inputPCollectionId) .getWindowingStrategyId());
pipeline .getComponents() .getTransformsMap() .forEach( (transformId, transform) -> { pipeline .getComponents() .getTransformsMap() .forEach( (transformId, transform) -> {
Components.newBuilder() .putTransforms("pt", pt) .putPcollections("input.out", input)
PCollectionNode inputCollection = PipelineNode.pCollection( inputPCollectionId, components.getPcollectionsOrThrow(inputPCollectionId)); DataSet<WindowedValue<KV<K, V>>> inputDataSet = context.getDataSetOrThrow(inputPCollectionId); RunnerApi.WindowingStrategy windowingStrategyProto = pipeline .getComponents() .getWindowingStrategiesOrThrow( pipeline .getComponents() .getPcollectionsOrThrow(inputPCollectionId) .getWindowingStrategyId());
.addRootTransformIds(flatten.getId()) .setComponents( RunnerApi.Components.newBuilder() .putTransforms(source.getId(), source.getTransform()) .putPcollections(left.getId(), left.getPCollection())
for (String transformId : components.getTransformsMap().keySet()) { PTransform transform = components.getTransformsOrThrow(transformId); String previousId = uniqueNamesById.put(transform.getUniqueName(), transformId); for (String pcollectionId : components.getPcollectionsMap().keySet()) { PCollection pc = components.getPcollectionsOrThrow(pcollectionId); checkArgument( !pc.getUniqueName().isEmpty(), pc.getUniqueName()); checkArgument( components.containsCoders(pc.getCoderId()), "%s: PCollection %s uses unknown coder %s", context, pc.getCoderId()); checkArgument( components.containsWindowingStrategies(pc.getWindowingStrategyId()), "%s: PCollection %s uses unknown windowing strategy %s", context, for (String strategyId : components.getWindowingStrategiesMap().keySet()) { WindowingStrategy strategy = components.getWindowingStrategiesOrThrow(strategyId); checkArgument( components.containsCoders(strategy.getWindowCoderId()), "%s: WindowingStrategy %s uses unknown coder %s", context, for (String coderId : components.getCodersMap().keySet()) {
Components.newBuilder() .putTransforms( "impulse",
.addAllRootTransformIds(ImmutableList.of("first", "second")) .setComponents( Components.newBuilder() .putTransforms( "first", WindowingStrategy.newBuilder().setAccumulationMode(Enum.ACCUMULATING).build(); RunnerApi.Components extraComponents = Components.newBuilder() .putPcollections( "intermediatePc", ProtoOverrides.updateTransform( "beam:second", p, new TestReplacer(secondReplacement, extraComponents)); PTransform updatedSecond = updated.getComponents().getTransformsOrThrow("second"); updated.getComponents().getWindowingStrategiesOrThrow("new_ws"), equalTo(introducedWS)); assertThat( updated.getComponents().getTransformsOrThrow("second_sub"), equalTo(PTransform.getDefaultInstance())); updated.getComponents().getPcollectionsOrThrow("intermediatePc").getUniqueName(), equalTo("intermediate_replacement")); updated.getComponents().getTransformsOrThrow("first"), equalTo(p.getComponents().getTransformsOrThrow("first"))); assertThat( updated.getComponents().getCodersOrThrow("coder"), equalTo(p.getComponents().getCodersOrThrow("coder")));
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections(
.toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(),