GreedyPipelineFuser { public static FusedPipeline fuse(Pipeline p) { return new GreedyPipelineFuser(p).fusedPipeline; } private  GreedyPipelineFuser(Pipeline p); static FusedPipeline fuse(Pipeline p);  }
@Test public void singleEnvironmentBecomesASingleStage() { String name = "read.out"; Components components = partialComponents .toBuilder() .putTransforms( "read", 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("py") .build()) .putPcollections("read.out", pc(name)) .putTransforms( "parDo", PTransform.newBuilder() .setUniqueName("ParDo") .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("py") .build()) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms( "window", PTransform.newBuilder() .setUniqueName("Window") .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("py") .build()) .putPcollections("window.out", pc("window.out")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), contains(PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), contains( ExecutableStageMatcher.withInput("impulse.out") .withNoOutputs() .withTransforms("read", "parDo", "window"))); }
@Test public void transformsWithNoEnvironmentBecomeRunnerExecuted() { Components components = partialComponents .toBuilder() .putTransforms( "mystery", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .setUniqueName("Mystery") .putInputs("input", "impulse.out") .putOutputs("output", "mystery.out") .build()) .putPcollections("mystery.out", pc("mystery.out")) .putTransforms( "enigma", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .setUniqueName("Enigma") .putInputs("input", "impulse.out") .putOutputs("output", "enigma.out") .build()) .putPcollections("enigma.out", pc("enigma.out")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")), PipelineNode.pTransform("mystery", components.getTransformsOrThrow("mystery")), PipelineNode.pTransform("enigma", components.getTransformsOrThrow("enigma")))); assertThat(fused.getFusedStages(), emptyIterable()); }
@Test public void singleEnvironmentAcrossGroupByKeyMultipleStages() { Components components = partialComponents .toBuilder() .putTransforms( "read", 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("py") .build()) .putPcollections("read.out", pc("read.out")) .putTransforms( "groupByKey", PTransform.newBuilder() .setUniqueName("GroupByKey") .putInputs("input", "read.out") .putOutputs("output", "groupByKey.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .build()) .putPcollections("groupByKey.out", pc("groupByKey.out")) .putTransforms( "parDo", PTransform.newBuilder() .setUniqueName("ParDo") .putInputs("input", "groupByKey.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("parDo.out", pc("parDo.out")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")), PipelineNode.pTransform("groupByKey", components.getTransformsOrThrow("groupByKey")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("read.out") .withTransforms("read"), ExecutableStageMatcher.withInput("groupByKey.out") .withNoOutputs() .withTransforms("parDo"))); }
@Test public void multipleEnvironmentsBecomesMultipleStages() { Components components = partialComponents .toBuilder() .putTransforms( "read", 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("py") .build()) .putPcollections("read.out", pc("read.out")) .putTransforms( "goTransform", PTransform.newBuilder() .setUniqueName("GoTransform") .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", pc("go.out")) .putTransforms( "pyTransform", PTransform.newBuilder() .setUniqueName("PyTransform") .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", pc("py.out")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat(fused.getRunnerExecutedTransforms(), hasSize(1)); assertThat(fused.getFusedStages(), hasSize(3)); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("read.out") .withTransforms("read"), ExecutableStageMatcher.withInput("read.out") .withNoOutputs() .withTransforms("pyTransform"), ExecutableStageMatcher.withInput("read.out") .withNoOutputs() .withTransforms("goTransform"))); }
@Test public void flattenWithHeterogenousInputsAndOutputsEntirelyMaterialized() { Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .putTransforms( "pyImpulse", PTransform.newBuilder() .setUniqueName("PyImpulse") .putOutputs("output", "pyImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("pyImpulse.out", pc("pyImpulse.out")) .putTransforms( "pyRead", PTransform.newBuilder() .setUniqueName("PyRead") .putInputs("input", "pyImpulse.out") .putOutputs("output", "pyRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("pyRead.out", pc("pyRead.out")) .putTransforms( "goImpulse", PTransform.newBuilder() .setUniqueName("GoImpulse") .putOutputs("output", "goImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("goImpulse.out", pc("goImpulse.out")) .putTransforms( "goRead", PTransform.newBuilder() .setUniqueName("GoRead") .putInputs("input", "goImpulse.out") .putOutputs("output", "goRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("go") .build()) .putPcollections("goRead.out", pc("goRead.out")) .putTransforms( "flatten", PTransform.newBuilder() .setUniqueName("Flatten") .putInputs("goReadInput", "goRead.out") .putInputs("pyReadInput", "pyRead.out") .putOutputs("output", "flatten.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)) .build()) .putPcollections("flatten.out", pc("flatten.out")) .putTransforms( "pyParDo", PTransform.newBuilder() .setUniqueName("PyParDo") .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())) .setEnvironmentId("py") .build()) .putPcollections("pyParDo.out", pc("pyParDo.out")) .putTransforms( "goParDo", PTransform.newBuilder() .setUniqueName("GoParDo") .putInputs("input", "flatten.out") .putOutputs("output", "goParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("go") .build()) .putPcollections("goParDo.out", pc("goParDo.out")) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat(fused.getRunnerExecutedTransforms(), hasSize(3)); assertThat( "The runner should include the impulses for both languages, plus an introduced flatten", fused.getRunnerExecutedTransforms(), hasItems( PipelineNode.pTransform("pyImpulse", components.getTransformsOrThrow("pyImpulse")), PipelineNode.pTransform("goImpulse", components.getTransformsOrThrow("goImpulse")))); PTransformNode flattenNode = null; for (PTransformNode runnerTransform : fused.getRunnerExecutedTransforms()) { if (getOnlyElement(runnerTransform.getTransform().getOutputsMap().values()) .equals("flatten.out")) { flattenNode = runnerTransform; } } assertThat(flattenNode, not(nullValue())); assertThat( flattenNode.getTransform().getSpec().getUrn(), equalTo(PTransformTranslation.FLATTEN_TRANSFORM_URN)); assertThat(new HashSet<>(flattenNode.getTransform().getInputsMap().values()), hasSize(2)); Collection<String> introducedOutputs = flattenNode.getTransform().getInputsMap().values(); AnyOf<String> anyIntroducedPCollection = anyOf(introducedOutputs.stream().map(Matchers::equalTo).collect(Collectors.toSet())); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("goImpulse.out") .withOutputs(anyIntroducedPCollection) .withTransforms("goRead", "flatten"), ExecutableStageMatcher.withInput("pyImpulse.out") .withOutputs(anyIntroducedPCollection) .withTransforms("pyRead", "flatten"), ExecutableStageMatcher.withInput("flatten.out") .withNoOutputs() .withTransforms("goParDo"), ExecutableStageMatcher.withInput("flatten.out") .withNoOutputs() .withTransforms("pyParDo"))); Set<String> materializedStageOutputs = fused.getFusedStages().stream() .flatMap(executableStage -> executableStage.getOutputPCollections().stream()) .map(PCollectionNode::getId) .collect(Collectors.toSet()); assertThat( "All materialized stage outputs should be flattened, and no more", materializedStageOutputs, containsInAnyOrder( flattenNode.getTransform().getInputsMap().values().toArray(new String[0]))); }
@Test public void flattenWithHeterogeneousInputsSingleEnvOutputPartiallyMaterialized() { Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .putTransforms( "pyImpulse", PTransform.newBuilder() .setUniqueName("PyImpulse") .putOutputs("output", "pyImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("pyImpulse.out", pc("pyImpulse.out")) .putTransforms( "pyRead", PTransform.newBuilder() .setUniqueName("PyRead") .putInputs("input", "pyImpulse.out") .putOutputs("output", "pyRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("pyRead.out", pc("pyRead.out")) .putTransforms( "goImpulse", PTransform.newBuilder() .setUniqueName("GoImpulse") .putOutputs("output", "goImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("goImpulse.out", pc("goImpulse.out")) .putTransforms( "goRead", PTransform.newBuilder() .setUniqueName("GoRead") .putInputs("input", "goImpulse.out") .putOutputs("output", "goRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("go") .build()) .putPcollections("goRead.out", pc("goRead.out")) .putTransforms( "flatten", PTransform.newBuilder() .setUniqueName("Flatten") .putInputs("goReadInput", "goRead.out") .putInputs("pyReadInput", "pyRead.out") .putOutputs("output", "flatten.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)) .build()) .putPcollections("flatten.out", pc("flatten.out")) .putTransforms( "goParDo", PTransform.newBuilder() .setUniqueName("GoParDo") .putInputs("input", "flatten.out") .putOutputs("output", "goParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("go") .build()) .putPcollections("goParDo.out", pc("goParDo.out")) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("pyImpulse", components.getTransformsOrThrow("pyImpulse")), PipelineNode.pTransform("goImpulse", components.getTransformsOrThrow("goImpulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("goImpulse.out") .withNoOutputs() .withTransforms("goRead", "flatten", "goParDo"), ExecutableStageMatcher.withInput("pyImpulse.out") .withOutputs("flatten.out") .withTransforms("pyRead", "flatten"), ExecutableStageMatcher.withInput("flatten.out") .withNoOutputs() .withTransforms("goParDo"))); }
@Test public void flattenAfterNoEnvDoesNotFuse() { Components components = partialComponents .toBuilder() .putTransforms( "flatten", PTransform.newBuilder() .setUniqueName("Flatten") .putInputs("impulseInput", "impulse.out") .putOutputs("output", "flatten.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN) .build()) .build()) .putPcollections("flatten.out", pc("flatten.out")) .putTransforms( "read", PTransform.newBuilder() .setUniqueName("Read") .putInputs("input", "flatten.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("read.out", pc("read.out")) .putTransforms( "parDo", PTransform.newBuilder() .setUniqueName("ParDo") .putInputs("input", "read.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder().build()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("parDo.out", pc("parDo.out")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")), PipelineNode.pTransform("flatten", components.getTransformsOrThrow("flatten")))); assertThat( fused.getFusedStages(), contains( ExecutableStageMatcher.withInput("flatten.out") .withNoOutputs() .withTransforms("read", "parDo"))); }
@Test public void sideInputRootsNewStage() { Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .putTransforms( "mainImpulse", PTransform.newBuilder() .setUniqueName("MainImpulse") .putOutputs("output", "mainImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("mainImpulse.out", pc("mainImpulse.out")) .putTransforms( "read", PTransform.newBuilder() .setUniqueName("Read") .putInputs("input", "mainImpulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("read.out", pc("read.out")) .putTransforms( "sideImpulse", PTransform.newBuilder() .setUniqueName("SideImpulse") .putOutputs("output", "sideImpulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("sideImpulse.out", pc("sideImpulse.out")) .putTransforms( "sideRead", PTransform.newBuilder() .setUniqueName("SideRead") .putInputs("input", "sideImpulse.out") .putOutputs("output", "sideRead.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build()) .putPcollections("sideRead.out", pc("sideRead.out")) .putTransforms( "leftParDo", PTransform.newBuilder() .setUniqueName("LeftParDo") .putInputs("main", "read.out") .putOutputs("output", "leftParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("py") .build()) .putPcollections("leftParDo.out", pc("leftParDo.out")) .putTransforms( "rightParDo", PTransform.newBuilder() .setUniqueName("RightParDo") .putInputs("main", "read.out") .putOutputs("output", "rightParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .build() .toByteString()) .build()) .setEnvironmentId("py") .build()) .putPcollections("rightParDo.out", pc("rightParDo.out")) .putTransforms( "sideParDo", PTransform.newBuilder() .setUniqueName("SideParDo") .putInputs("main", "read.out") .putInputs("side", "sideRead.out") .putOutputs("output", "sideParDo.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putSideInputs("side", SideInput.getDefaultInstance()) .build() .toByteString()) .build()) .setEnvironmentId("py") .build()) .putPcollections("sideParDo.out", pc("sideParDo.out")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("mainImpulse", components.getTransformsOrThrow("mainImpulse")), PipelineNode.pTransform( "sideImpulse", components.getTransformsOrThrow("sideImpulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("mainImpulse.out") .withOutputs("read.out") .withTransforms("read"), ExecutableStageMatcher.withInput("read.out") .withNoOutputs() .withTransforms("leftParDo", "rightParDo"), ExecutableStageMatcher.withInput("read.out") .withSideInputs( RunnerApi.ExecutableStagePayload.SideInputId.newBuilder() .setTransformId("sideParDo") .setLocalName("side") .build()) .withNoOutputs() .withTransforms("sideParDo"), ExecutableStageMatcher.withInput("sideImpulse.out") .withOutputs("sideRead.out") .withTransforms("sideRead"))); }
@Test public void statefulParDoRootsStage() { PTransform parDoTransform = PTransform.newBuilder() .setUniqueName("ParDo") .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() .setUniqueName("StatefulParDo") .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(); Components components = partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms("stateful", statefulTransform) .putPcollections("stateful.out", pc("stateful.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse( Pipeline.newBuilder() .setComponents(components) .addRequirements(ParDoTranslation.REQUIRES_STATEFUL_PROCESSING_URN) .build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("parDo.out") .withTransforms("parDo"), ExecutableStageMatcher.withInput("parDo.out") .withNoOutputs() .withTransforms("stateful"))); }
@Test public void parDoWithTimerRootsStage() { PTransform parDoTransform = PTransform.newBuilder() .setUniqueName("ParDo") .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() .setUniqueName("TimerParDo") .putInputs("input", "parDo.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.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(); Components components = partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse( Pipeline.newBuilder() .setComponents(components) .addRequirements(ParDoTranslation.REQUIRES_STATEFUL_PROCESSING_URN) .build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("parDo.out") .withTransforms("parDo"), ExecutableStageMatcher.withInput("parDo.out").withNoOutputs().withTransforms("timer"))); }
@Test public void parDoWithStateAndTimerRootsStage() { PTransform timerTransform = PTransform.newBuilder() .setUniqueName("TimerParDo") .putInputs("input", "impulse.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putStateSpecs("state", StateSpec.getDefaultInstance()) .putTimerFamilySpecs("timer", TimerFamilySpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); Components components = partialComponents .toBuilder() .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse( Pipeline.newBuilder() .setComponents(components) .addRequirements(ParDoTranslation.REQUIRES_STATEFUL_PROCESSING_URN) .build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), contains( ExecutableStageMatcher.withInput("impulse.out") .withNoOutputs() .withTransforms("timer"))); }
@Test public void compositesIgnored() { Components components = partialComponents .toBuilder() .putTransforms( "read", 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("py") .build()) .putPcollections("read.out", pc("read.out")) .putTransforms( "goTransform", PTransform.newBuilder() .setUniqueName("GoTransform") .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", pc("go.out")) .putTransforms( "pyTransform", PTransform.newBuilder() .setUniqueName("PyTransform") .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", pc("py.out")) .putTransforms( "compositeMultiLang", PTransform.newBuilder() .setUniqueName("CompositeMultiLang") .putInputs("input", "impulse.out") .putOutputs("pyOut", "py.out") .putOutputs("goOut", "go.out") .addSubtransforms("read") .addSubtransforms("goTransform") .addSubtransforms("pyTransform") .build()) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse( Pipeline.newBuilder() .addRootTransformIds("impulse") .addRootTransformIds("compositeMultiLang") .setComponents(components) .build()); assertThat(fused.getRunnerExecutedTransforms(), hasSize(1)); assertThat(fused.getFusedStages(), hasSize(3)); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("read.out") .withTransforms("read"), ExecutableStageMatcher.withInput("read.out") .withNoOutputs() .withTransforms("pyTransform"), ExecutableStageMatcher.withInput("read.out") .withNoOutputs() .withTransforms("goTransform"))); }
@Test public void sanitizedTransforms() throws Exception { PCollection flattenOutput = pc("flatten.out"); PCollection read1Output = pc("read1.out"); PCollection read2Output = pc("read2.out"); PCollection impulse1Output = pc("impulse1.out"); PCollection impulse2Output = pc("impulse2.out"); PTransform flattenTransform = PTransform.newBuilder() .setUniqueName("Flatten") .putInputs(read1Output.getUniqueName(), read1Output.getUniqueName()) .putInputs(read2Output.getUniqueName(), read2Output.getUniqueName()) .putOutputs(flattenOutput.getUniqueName(), flattenOutput.getUniqueName()) .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build(); PTransform read1Transform = PTransform.newBuilder() .setUniqueName("read1") .putInputs(impulse1Output.getUniqueName(), impulse1Output.getUniqueName()) .putOutputs(read1Output.getUniqueName(), read1Output.getUniqueName()) .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build(); PTransform read2Transform = PTransform.newBuilder() .setUniqueName("read2") .putInputs(impulse2Output.getUniqueName(), impulse2Output.getUniqueName()) .putOutputs(read2Output.getUniqueName(), read2Output.getUniqueName()) .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .setEnvironmentId("py") .build(); PTransform impulse1Transform = PTransform.newBuilder() .setUniqueName("impulse1") .putOutputs(impulse1Output.getUniqueName(), impulse1Output.getUniqueName()) .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .build(); PTransform impulse2Transform = PTransform.newBuilder() .setUniqueName("impulse2") .putOutputs(impulse2Output.getUniqueName(), impulse2Output.getUniqueName()) .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN) .setPayload( WindowIntoPayload.newBuilder() .setWindowFn(FunctionSpec.newBuilder()) .build() .toByteString())) .build(); Pipeline impulse = Pipeline.newBuilder() .addRootTransformIds(impulse1Transform.getUniqueName()) .addRootTransformIds(impulse2Transform.getUniqueName()) .addRootTransformIds(flattenTransform.getUniqueName()) .setComponents( Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .putEnvironments("py", Environments.createDockerEnvironment("py")) .putPcollections(flattenOutput.getUniqueName(), flattenOutput) .putTransforms(flattenTransform.getUniqueName(), flattenTransform) .putPcollections(read1Output.getUniqueName(), read1Output) .putTransforms(read1Transform.getUniqueName(), read1Transform) .putPcollections(read2Output.getUniqueName(), read2Output) .putTransforms(read2Transform.getUniqueName(), read2Transform) .putPcollections(impulse1Output.getUniqueName(), impulse1Output) .putTransforms(impulse1Transform.getUniqueName(), impulse1Transform) .putPcollections(impulse2Output.getUniqueName(), impulse2Output) .putTransforms(impulse2Transform.getUniqueName(), impulse2Transform) .build()) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(impulse); assertThat(fused.getRunnerExecutedTransforms(), hasSize(2)); assertThat(fused.getFusedStages(), hasSize(2)); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput(impulse1Output.getUniqueName()) .withTransforms(flattenTransform.getUniqueName(), read1Transform.getUniqueName()), ExecutableStageMatcher.withInput(impulse2Output.getUniqueName()) .withTransforms(flattenTransform.getUniqueName(), read2Transform.getUniqueName()))); assertThat( fused.getFusedStages().stream() .flatMap( s -> s.getComponents().getTransformsOrThrow(flattenTransform.getUniqueName()) .getInputsMap().values().stream()) .collect(Collectors.toList()), containsInAnyOrder(read1Output.getUniqueName(), read2Output.getUniqueName())); }
@Test public void statefulParDoRootsStage() { PTransform parDoTransform = PTransform.newBuilder() .setUniqueName("ParDo") .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() .setUniqueName("StatefulParDo") .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(); Components components = partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms("stateful", statefulTransform) .putPcollections("stateful.out", pc("stateful.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("parDo.out") .withTransforms("parDo"), ExecutableStageMatcher.withInput("parDo.out") .withNoOutputs() .withTransforms("stateful"))); }
@Test public void parDoWithTimerRootsStage() { PTransform parDoTransform = PTransform.newBuilder() .setUniqueName("ParDo") .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() .setUniqueName("TimerParDo") .putInputs("input", "parDo.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.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(); Components components = partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), containsInAnyOrder( ExecutableStageMatcher.withInput("impulse.out") .withOutputs("parDo.out") .withTransforms("parDo"), ExecutableStageMatcher.withInput("parDo.out").withNoOutputs().withTransforms("timer"))); }
@Test public void parDoWithStateAndTimerRootsStage() { PTransform timerTransform = PTransform.newBuilder() .setUniqueName("TimerParDo") .putInputs("input", "impulse.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .setPayload( ParDoPayload.newBuilder() .setDoFn(FunctionSpec.newBuilder()) .putStateSpecs("state", StateSpec.getDefaultInstance()) .putTimerSpecs("timer", TimerSpec.getDefaultInstance()) .build() .toByteString())) .setEnvironmentId("common") .build(); Components components = partialComponents .toBuilder() .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build(); FusedPipeline fused = GreedyPipelineFuser.fuse(Pipeline.newBuilder().setComponents(components).build()); assertThat( fused.getRunnerExecutedTransforms(), containsInAnyOrder( PipelineNode.pTransform("impulse", components.getTransformsOrThrow("impulse")))); assertThat( fused.getFusedStages(), contains( ExecutableStageMatcher.withInput("impulse.out") .withNoOutputs() .withTransforms("timer"))); }