GreedyStageFuser { public static ExecutableStage forGrpcPortRead( QueryablePipeline pipeline, PCollectionNode inputPCollection, Set<PTransformNode> initialNodes) { checkArgument( !initialNodes.isEmpty(), "%s must contain at least one %s.", GreedyStageFuser.class.getSimpleName(), PTransformNode.class.getSimpleName()); Environment environment = getStageEnvironment(pipeline, initialNodes); ImmutableSet.Builder<PTransformNode> fusedTransforms = ImmutableSet.builder(); fusedTransforms.addAll(initialNodes); Set<SideInputReference> sideInputs = new LinkedHashSet<>(); Set<UserStateReference> userStates = new LinkedHashSet<>(); Set<TimerReference> timers = new LinkedHashSet<>(); Set<PCollectionNode> fusedCollections = new LinkedHashSet<>(); Set<PCollectionNode> materializedPCollections = new LinkedHashSet<>(); Queue<PCollectionNode> fusionCandidates = new ArrayDeque<>(); for (PTransformNode initialConsumer : initialNodes) { fusionCandidates.addAll(pipeline.getOutputPCollections(initialConsumer)); sideInputs.addAll(pipeline.getSideInputs(initialConsumer)); userStates.addAll(pipeline.getUserStates(initialConsumer)); timers.addAll(pipeline.getTimers(initialConsumer)); } while (!fusionCandidates.isEmpty()) { PCollectionNode candidate = fusionCandidates.poll(); if (fusedCollections.contains(candidate) || materializedPCollections.contains(candidate)) { LOG.debug( "Skipping fusion candidate {} because it is {} in this {}", candidate, fusedCollections.contains(candidate) ? "fused" : "materialized", ExecutableStage.class.getSimpleName()); continue; } PCollectionFusibility fusibility = canFuse(pipeline, candidate, environment, fusedCollections); switch (fusibility) { case MATERIALIZE: materializedPCollections.add(candidate); break; case FUSE: fusedCollections.add(candidate); fusedTransforms.addAll(pipeline.getPerElementConsumers(candidate)); for (PTransformNode consumer : pipeline.getPerElementConsumers(candidate)) { fusionCandidates.addAll(pipeline.getOutputPCollections(consumer)); sideInputs.addAll(pipeline.getSideInputs(consumer)); } break; default: throw new IllegalStateException( String.format( "Unknown type of %s %s", PCollectionFusibility.class.getSimpleName(), fusibility)); } } return ImmutableExecutableStage.ofFullComponents( pipeline.getComponents(), environment, inputPCollection, sideInputs, userStates, timers, fusedTransforms.build(), materializedPCollections, DEFAULT_WIRE_CODER_SETTINGS); } private  GreedyStageFuser(); static ExecutableStage forGrpcPortRead(
      QueryablePipeline pipeline,
      PCollectionNode inputPCollection,
      Set<PTransformNode> initialNodes);  }
@Test public void noInitialConsumersThrows() { QueryablePipeline p = QueryablePipeline.forPrimitivesIn(partialComponents); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("at least one PTransform"); GreedyStageFuser.forGrpcPortRead(p, impulseOutputNode, Collections.emptySet()); }
@Test public void differentEnvironmentsThrows() { QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms( "read", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .build()) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "goTransform", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "go.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("go") .build()) .putPcollections("go.out", PCollection.newBuilder().setUniqueName("go.out").build()) .putTransforms( "pyTransform", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "py.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("py.out", PCollection.newBuilder().setUniqueName("py.out").build()) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .build()); Set<PTransformNode> differentEnvironments = p.getPerElementConsumers( PipelineNode.pCollection( "read.out", PCollection.newBuilder().setUniqueName("read.out").build())); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("go"); thrown.expectMessage("py"); thrown.expectMessage("same"); GreedyStageFuser.forGrpcPortRead( p, PipelineNode.pCollection( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()), differentEnvironments); }
@Test public void noEnvironmentThrows() { 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))); }
@Test public void fusesCompatibleEnvironments() { PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform windowTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("window", windowTransform) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of( PipelineNode.pTransform("parDo", parDoTransform), PipelineNode.pTransform("window", windowTransform))); assertThat(subgraph.getOutputPCollections(), emptyIterable()); assertThat(subgraph, hasSubtransforms("parDo", "window")); }
@Test public void materializesWithStatefulConsumer() { PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform statefulTransform = PTransform.newBuilder() .putInputs("input", "parDo.out") .putOutputs("output", "stateful.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putStateSpecs("state", StateSpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("stateful", statefulTransform) .putPcollections( "stateful.out", PCollection.newBuilder().setUniqueName("stateful.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform))); assertThat( subgraph.getOutputPCollections(), contains( PipelineNode.pCollection( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()))); assertThat(subgraph, hasSubtransforms("parDo")); }
@Test public void materializesWithConsumerWithTimer() { PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform timerTransform = PTransform.newBuilder() .putInputs("input", "parDo.out") .putOutputs("output", "timer.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putTimerFamilySpecs("timer", TimerFamilySpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("timer", timerTransform) .putPcollections( "timer.out", PCollection.newBuilder().setUniqueName("timer.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform))); assertThat( subgraph.getOutputPCollections(), contains( PipelineNode.pCollection( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()))); assertThat(subgraph, hasSubtransforms("parDo")); }
@Test public void fusesFlatten() { PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform flattenTransform = PTransform.newBuilder() .putInputs("readInput", "read.out") .putInputs("parDoInput", "parDo.out") .putOutputs("output", "flatten.out") .setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)) .build(); PTransform windowTransform = PTransform.newBuilder() .putInputs("input", "flatten.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("flatten", flattenTransform) .putPcollections( "flatten.out", PCollection.newBuilder().setUniqueName("flatten.out").build()) .putTransforms("window", windowTransform) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, p.getPerElementConsumers(impulseOutputNode)); assertThat(subgraph.getOutputPCollections(), emptyIterable()); assertThat(subgraph, hasSubtransforms("read", "parDo", "flatten", "window")); }
@Test public void fusesFlattenWithDifferentEnvironmentInputs() { PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform otherEnvRead = PTransform.newBuilder() .putInputs("impulse", "impulse.out") .putOutputs("output", "envRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("rare") .build(); PTransform flattenTransform = PTransform.newBuilder() .putInputs("readInput", "read.out") .putInputs("otherEnvInput", "envRead.out") .putOutputs("output", "flatten.out") .setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)) .build(); PTransform windowTransform = PTransform.newBuilder() .putInputs("input", "flatten.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); Components components = partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections("read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms("envRead", otherEnvRead) .putPcollections( "envRead.out", PCollection.newBuilder().setUniqueName("envRead.out").build()) .putTransforms("flatten", flattenTransform) .putPcollections( "flatten.out", PCollection.newBuilder().setUniqueName("flatten.out").build()) .putTransforms("window", windowTransform) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .putEnvironments("rare", Environments.createDockerEnvironment("rare")) .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn(components); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("read", readTransform))); assertThat(subgraph.getOutputPCollections(), emptyIterable()); assertThat(subgraph, hasSubtransforms("read", "flatten", "window")); ExecutableStage readFromOtherEnv = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("envRead", otherEnvRead))); assertThat( readFromOtherEnv.getOutputPCollections(), contains( PipelineNode.pCollection( "flatten.out", components.getPcollectionsOrThrow("flatten.out")))); assertThat(readFromOtherEnv, hasSubtransforms("envRead", "flatten")); }
@Test public void flattenWithHeterogeneousInputsAndOutputs() { PTransform pyRead = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "pyRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("py") .build(); PTransform goRead = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "goRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("go") .build(); PTransform pyParDo = PTransform.newBuilder() .putInputs("input", "flatten.out") .putOutputs("output", "pyParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("py") .build(); PTransform goWindow = PTransform.newBuilder() .putInputs("input", "flatten.out") .putOutputs("output", "goWindow.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("go") .build(); PCollection flattenPc = PCollection.newBuilder().setUniqueName("flatten.out").build(); Components components = partialComponents .toBuilder() .putTransforms("pyRead", pyRead) .putPcollections( "pyRead.out", PCollection.newBuilder().setUniqueName("pyRead.out").build()) .putTransforms("goRead", goRead) .putPcollections( "goRead.out", PCollection.newBuilder().setUniqueName("goRead.out").build()) .putTransforms( "flatten", PTransform.newBuilder() .putInputs("py_input", "pyRead.out") .putInputs("go_input", "goRead.out") .putOutputs("output", "flatten.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN) .build()) .build()) .putPcollections("flatten.out", flattenPc) .putTransforms("pyParDo", pyParDo) .putPcollections( "pyParDo.out", PCollection.newBuilder().setUniqueName("pyParDo.out").build()) .putTransforms("goWindow", goWindow) .putPcollections( "goWindow.out", PCollection.newBuilder().setUniqueName("goWindow.out").build()) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn(components); ExecutableStage readFromPy = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("pyRead", pyRead))); ExecutableStage readFromGo = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("goRead", goRead))); assertThat( readFromPy.getOutputPCollections(), contains(PipelineNode.pCollection("flatten.out", flattenPc))); assertThat( readFromPy.getTransforms(), not(hasItem(PipelineNode.pTransform("pyParDo", pyParDo)))); assertThat( readFromGo.getOutputPCollections(), contains(PipelineNode.pCollection("flatten.out", flattenPc))); assertThat( readFromGo.getTransforms(), not(hasItem(PipelineNode.pTransform("goWindow", goWindow)))); }
@Test public void materializesWithDifferentEnvConsumer() { Environment env = Environments.createDockerEnvironment("common"); PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("out", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PCollection parDoOutput = PCollection.newBuilder().setUniqueName("parDo.out").build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", parDoOutput) .putTransforms( "window", PTransform.newBuilder() .putInputs("input", "parDo.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("rare") .build()) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("rare", Environments.createDockerEnvironment("rare")) .putEnvironments("common", env) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, p.getPerElementConsumers(impulseOutputNode)); assertThat( subgraph.getOutputPCollections(), contains(PipelineNode.pCollection("parDo.out", parDoOutput))); assertThat(subgraph.getInputPCollection(), equalTo(impulseOutputNode)); assertThat(subgraph.getEnvironment(), equalTo(env)); assertThat( subgraph.getTransforms(), contains(PipelineNode.pTransform("parDo", parDoTransform))); }
@Test public void materializesWithDifferentEnvSibling() { Environment env = Environments.createDockerEnvironment("common"); PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "parDo", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build()) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms( "window", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("rare") .build()) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("rare", Environments.createDockerEnvironment("rare")) .putEnvironments("common", env) .build()); PTransformNode readNode = PipelineNode.pTransform("read", readTransform); PCollectionNode readOutput = getOnlyElement(p.getOutputPCollections(readNode)); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("read", readTransform))); assertThat(subgraph.getOutputPCollections(), contains(readOutput)); assertThat(subgraph.getTransforms(), contains(readNode)); }
@Test public void materializesWithSideInputConsumer() { Environment env = Environments.createDockerEnvironment("common"); PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "side_read", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .putInputs("input", "impulse.out") .putOutputs("output", "side_read.out") .build()) .putPcollections( "side_read.out", PCollection.newBuilder().setUniqueName("side_read.out").build()) .putTransforms( "parDo", PTransform.newBuilder() .putInputs("input", "read.out") .putInputs("side_input", "side_read.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putSideInputs("side_input", SideInput.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build()) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms( "window", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build()) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("common", env) .build()); PTransformNode readNode = PipelineNode.pTransform("read", readTransform); PCollectionNode readOutput = getOnlyElement(p.getOutputPCollections(readNode)); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead(p, impulseOutputNode, ImmutableSet.of(readNode)); assertThat(subgraph.getOutputPCollections(), contains(readOutput)); assertThat(subgraph, hasSubtransforms(readNode.getId())); }
@Test public void sideInputIncludedInStage() { Environment env = Environments.createDockerEnvironment("common"); PTransform readTransform = PTransform.newBuilder() .setUniqueName("read") .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform parDoTransform = PTransform.newBuilder() .setUniqueName("parDo") .putInputs("input", "read.out") .putInputs("side_input", "side_read.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putSideInputs("side_input", SideInput.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); PCollection sideInputPCollection = PCollection.newBuilder().setUniqueName("side_read.out").build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "side_read", PTransform.newBuilder() .setUniqueName("side_read") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .putInputs("input", "impulse.out") .putOutputs("output", "side_read.out") .build()) .putPcollections("side_read.out", sideInputPCollection) .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putEnvironments("common", env) .build()); PCollectionNode readOutput = getOnlyElement(p.getOutputPCollections(PipelineNode.pTransform("read", readTransform))); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, readOutput, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform))); PTransformNode parDoNode = PipelineNode.pTransform("parDo", parDoTransform); SideInputReference sideInputRef = SideInputReference.of( parDoNode, "side_input", PipelineNode.pCollection("side_read.out", sideInputPCollection)); assertThat(subgraph.getSideInputs(), contains(sideInputRef)); assertThat(subgraph.getOutputPCollections(), emptyIterable()); }
@Test public void executableStageProducingSideInputMaterializesIt() { Environment env = Environments.createDockerEnvironment("common"); PTransform impulse = PTransform.newBuilder() .setUniqueName("impulse") .putOutputs("output", "impulsePC") .setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build(); PTransform createSide = PTransform.newBuilder() .setUniqueName("createSide") .putInputs("input", "impulsePC") .putOutputs("output", "sidePC") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform processMain = PTransform.newBuilder() .setUniqueName("processMain") .putInputs("main", "impulsePC") .putInputs("side", "sidePC") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putSideInputs("side", SideInput.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); PCollection sidePC = PCollection.newBuilder().setUniqueName("sidePC").build(); PCollection impulsePC = PCollection.newBuilder().setUniqueName("impulsePC").build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("impulse", impulse) .putTransforms("createSide", createSide) .putTransforms("processMain", processMain) .putPcollections("impulsePC", impulsePC) .putPcollections("sidePC", sidePC) .putEnvironments("common", env) .build()); PCollectionNode impulseOutput = getOnlyElement(p.getOutputPCollections(PipelineNode.pTransform("impulse", impulse))); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutput, ImmutableSet.of(PipelineNode.pTransform("createSide", createSide))); assertThat( subgraph.getOutputPCollections(), contains(PipelineNode.pCollection("sidePC", sidePC))); }
@Test public void userStateIncludedInStage() { Environment env = Environments.createDockerEnvironment("common"); PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putStateSpecs("state_spec", StateSpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); PCollection userStateMainInputPCollection = PCollection.newBuilder().setUniqueName("read.out").build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections("read.out", userStateMainInputPCollection) .putTransforms( "user_state", PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "user_state.out") .build()) .putPcollections( "user_state.out", PCollection.newBuilder().setUniqueName("user_state.out").build()) .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putEnvironments("common", env) .build()); PCollectionNode readOutput = getOnlyElement(p.getOutputPCollections(PipelineNode.pTransform("read", readTransform))); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, readOutput, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform))); PTransformNode parDoNode = PipelineNode.pTransform("parDo", parDoTransform); UserStateReference userStateRef = UserStateReference.of( parDoNode, "state_spec", PipelineNode.pCollection("read.out", userStateMainInputPCollection)); assertThat(subgraph.getUserStates(), contains(userStateRef)); assertThat(subgraph.getOutputPCollections(), emptyIterable()); }
@Test public void materializesWithGroupByKeyConsumer() { Environment env = Environments.createDockerEnvironment("common"); PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "gbk", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "gbk.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .build()) .putPcollections( "gbk.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putEnvironments("common", env) .build()); PTransformNode readNode = PipelineNode.pTransform("read", readTransform); PCollectionNode readOutput = getOnlyElement(p.getOutputPCollections(readNode)); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead(p, impulseOutputNode, ImmutableSet.of(readNode)); assertThat(subgraph.getOutputPCollections(), contains(readOutput)); assertThat(subgraph, hasSubtransforms(readNode.getId())); }
@Test public void materializesWithConsumerWithTimer() { PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("common") .build(); PTransform timerTransform = PTransform.newBuilder() .putInputs("input", "parDo.out") .putOutputs("output", "timer.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putTimerSpecs("timer", TimerSpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("timer", timerTransform) .putPcollections( "timer.out", PCollection.newBuilder().setUniqueName("timer.out").build()) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build()); ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform))); assertThat( subgraph.getOutputPCollections(), contains( PipelineNode.pCollection( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()))); assertThat(subgraph, hasSubtransforms("parDo")); }