public static ImmutableExecutableStage ofFullComponents( Components components, Environment environment, PCollectionNode input, Collection<SideInputReference> sideInputs, Collection<UserStateReference> userStates, Collection<TimerReference> timers, Collection<PTransformNode> transforms, Collection<PCollectionNode> outputs) { Components prunedComponents = components .toBuilder() .clearTransforms() .putAllTransforms( transforms .stream() .collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform))) .build(); return of( prunedComponents, environment, input, sideInputs, userStates, timers, transforms, outputs); }
getComponents().toBuilder().putAllTransforms(executableStageTransforms).build(); List<String> rootTransformIds = StreamSupport.stream(
/** * Update all composites present in the {@code originalPipeline} with an URN equal to the provided * {@code urn} using the provided {@link TransformReplacement}. */ public static Pipeline updateTransform( String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) { Components.Builder resultComponents = originalPipeline.getComponents().toBuilder(); for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) { if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) { MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents()); checkArgument( updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName()); removeSubtransforms(pt.getValue(), resultComponents); resultComponents .mergeFrom(updated.getComponents()) .putTransforms(pt.getKey(), updated.getPtransform()); } } return originalPipeline.toBuilder().setComponents(resultComponents).build(); }
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections(
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections(
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections(
Components components = partialComponents .toBuilder() .putTransforms( "read",
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections("read.out", userStateMainInputPCollection)
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms( "read",
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections(
.toBuilder() .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out"))
FnDataReceiver<WindowedValue<byte[]>> fooReceiver = factory.create(fooPC.getId()); Components.Builder builder = baseComponents.toBuilder(); String sdkWireCoderId = WireCoders.addSdkWireCoder(fooPC, builder); String barSdkWireCoderId = WireCoders.addSdkWireCoder(barPC, builder);
QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections(
FnDataReceiver<WindowedValue<byte[]>> secondReceiver = factory.create(fooPC.getId()); Components.Builder builder = baseComponents.toBuilder(); String sdkWireCoderId = WireCoders.addSdkWireCoder(fooPC, builder); Components components = builder.build();
FnDataReceiver<WindowedValue<byte[]>> receiver = factory.create(fooPC.getId()); Builder builder = baseComponents.toBuilder(); String sdkWireCoderId = WireCoders.addSdkWireCoder(fooPC, builder); Components components = builder.build();
Components components = partialComponents .toBuilder() .putTransforms( "mystery",
stage .getComponents() .toBuilder() .clearTransforms() .putAllTransforms(
/** This method doesn't do any pruning for reachability, but this may not require a test. */ @Test public void retainOnlyPrimitivesIgnoresUnreachableNodes() { Pipeline p = Pipeline.create(); p.apply( new org.apache.beam.sdk.transforms.PTransform<PBegin, PCollection<Long>>() { @Override public PCollection<Long> expand(PBegin input) { return input .apply(GenerateSequence.from(2L)) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(5L)))) .apply(MapElements.into(TypeDescriptors.longs()).via(l -> l + 1)); } }); Components augmentedComponents = PipelineTranslation.toProto(p) .getComponents() .toBuilder() .putCoders("extra-coder", RunnerApi.Coder.getDefaultInstance()) .putWindowingStrategies( "extra-windowing-strategy", RunnerApi.WindowingStrategy.getDefaultInstance()) .putEnvironments("extra-env", RunnerApi.Environment.getDefaultInstance()) .putPcollections("extra-pc", RunnerApi.PCollection.getDefaultInstance()) .build(); Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(augmentedComponents); } }
private Coder<K> getKeyCoder(PTransformNode application) { PCollectionNode inputPCollection = getOnlyElement(graph.getPerElementInputs(application)); try { // We know the type restrictions on the input PCollection, and the restrictions on the // Wire coder Builder builder = GroupByKeyOnlyEvaluatorFactory.this.components.toBuilder(); String wireCoderId = WireCoders.addRunnerWireCoder(inputPCollection, builder); Coder<WindowedValue<KV<K, V>>> wireCoder = (Coder<WindowedValue<KV<K, V>>>) RehydratedComponents.forComponents(builder.build()).getCoder(wireCoderId); checkArgument( wireCoder instanceof WindowedValue.WindowedValueCoder, "Wire %s must be a %s", Coder.class.getSimpleName(), WindowedValueCoder.class.getSimpleName()); WindowedValueCoder<KV<K, V>> windowedValueCoder = (WindowedValueCoder<KV<K, V>>) wireCoder; checkArgument( windowedValueCoder.getValueCoder() instanceof KvCoder, "Input elements to %s must be encoded with a %s", DirectGroupByKey.DirectGroupByKeyOnly.class.getSimpleName(), KvCoder.class.getSimpleName()); KvCoder<K, V> kvCoder = (KvCoder<K, V>) windowedValueCoder.getValueCoder(); return kvCoder.getKeyCoder(); } catch (IOException e) { throw new RuntimeException(e); } }
@Test public void noEnvironmentThrows() { // (impulse.out) -> runnerTransform -> gbk.out // runnerTransform can't be executed in an environment, so trying to construct it should fail PTransform gbkTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .setSpec( FunctionSpec.newBuilder().setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .putOutputs("output", "gbk.out") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("runnerTransform", gbkTransform) .putPcollections( "gbk.out", PCollection.newBuilder().setUniqueName("gbk.out").build()) .build()); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Environment must be populated"); GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("runnerTransform", gbkTransform))); }