| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.core.construction.graph; |
| |
| import static org.apache.beam.vendor.guava.v20_0.com.google.common.collect.Iterables.getOnlyElement; |
| import static org.hamcrest.Matchers.contains; |
| import static org.hamcrest.Matchers.emptyIterable; |
| import static org.hamcrest.Matchers.equalTo; |
| import static org.hamcrest.Matchers.hasItem; |
| import static org.hamcrest.Matchers.not; |
| import static org.junit.Assert.assertThat; |
| |
| import java.util.Collections; |
| import java.util.Set; |
| import java.util.stream.Collectors; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Components; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Environment; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.SideInput; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.StateSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.TimerSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.WindowIntoPayload; |
| import org.apache.beam.runners.core.construction.Environments; |
| import org.apache.beam.runners.core.construction.PTransformTranslation; |
| import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode; |
| import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode; |
| import org.apache.beam.vendor.guava.v20_0.com.google.common.collect.ImmutableSet; |
| import org.hamcrest.Description; |
| import org.hamcrest.TypeSafeMatcher; |
| import org.junit.Before; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.ExpectedException; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** Tests for {@link GreedyStageFuser}. */ |
| @RunWith(JUnit4.class) |
| public class GreedyStageFuserTest { |
| @Rule public ExpectedException thrown = ExpectedException.none(); |
| |
| private final PCollection impulseDotOut = |
| PCollection.newBuilder().setUniqueName("impulse.out").build(); |
| private final PCollectionNode impulseOutputNode = |
| PipelineNode.pCollection("impulse.out", impulseDotOut); |
| |
| private Components partialComponents; |
| |
| @Before |
| public void setup() { |
| partialComponents = |
| Components.newBuilder() |
| .putTransforms( |
| "impulse", |
| PTransform.newBuilder() |
| .putOutputs("output", "impulse.out") |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) |
| .build()) |
| .putPcollections("impulse.out", impulseDotOut) |
| .build(); |
| } |
| |
| @Test |
| public void noInitialConsumersThrows() { |
| // (impulse.out) -> () is not a meaningful stage, so it should never be called |
| 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() { |
| // (impulse.out) -> read -> read.out --> go -> go.out |
| // \ |
| // -> py -> py.out |
| // read.out can't be fused with both 'go' and 'py', so we should refuse to create this stage |
| 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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("go")) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("py")) |
| .build() |
| .toByteString())) |
| .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() { |
| // (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))); |
| } |
| |
| @Test |
| public void fusesCompatibleEnvironments() { |
| // (impulse.out) -> parDo -> parDo.out -> window -> window.out |
| // parDo and window both have the environment "common" and can be fused together |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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))); |
| // Nothing consumes the outputs of ParDo or Window, so they don't have to be materialized |
| assertThat(subgraph.getOutputPCollections(), emptyIterable()); |
| assertThat(subgraph, hasSubtransforms("parDo", "window")); |
| } |
| |
| @Test |
| public void materializesWithStatefulConsumer() { |
| // (impulse.out) -> parDo -> (parDo.out) |
| // (parDo.out) -> stateful -> stateful.out |
| // stateful has a state spec which prevents it from fusing with an upstream ParDo |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putStateSpecs("state", StateSpec.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> parDo -> (parDo.out) |
| // (parDo.out) -> timer -> timer.out |
| // timer has a timer spec which prevents it from fusing with an upstream ParDo |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putTimerSpecs("timer", TimerSpec.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> parDo -> parDo.out --> flatten -> flatten.out -> window -> window.out |
| // \ / |
| // -> read -> read.out - |
| // The flatten can be executed within the same environment as any transform; the window can |
| // execute in the same environment as the rest of the transforms, and can fuse with the stage |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> read -> read.out \ -> window -> window.out |
| // -------> flatten -> flatten.out / |
| // (impulse.out) -> envRead -> envRead.out / |
| // fuses into |
| // read -> read.out -> flatten -> flatten.out -> window -> window.out |
| // envRead -> envRead.out -> flatten -> (flatten.out) |
| // (flatten.out) -> window -> window.out |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("rare")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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")); |
| |
| // Flatten shows up in both of these subgraphs, but elements only go through a path to the |
| // flatten once. |
| 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() { |
| // (impulse.out) -> pyRead -> pyRead.out \ -> pyParDo -> pyParDo.out |
| // (impulse.out) -> -> flatten -> flatten.out | |
| // (impulse.out) -> goRead -> goRead.out / -> goWindow -> goWindow.out |
| // fuses into |
| // (impulse.out) -> pyRead -> pyRead.out -> flatten -> (flatten.out) |
| // (impulse.out) -> goRead -> goRead.out -> flatten -> (flatten.out) |
| // (flatten.out) -> pyParDo -> pyParDo.out |
| // (flatten.out) -> goWindow -> goWindow.out |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("py")) |
| .build() |
| .toByteString()) |
| .build()) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("go")) |
| .build() |
| .toByteString()) |
| .build()) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("py")) |
| .build() |
| .toByteString()) |
| .build()) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("go")) |
| .build() |
| .toByteString()) |
| .build()) |
| .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))); |
| // The stage must materialize the flatten, so the `go` stage can read it; this means that this |
| // parDo can't be in the stage, as it'll be a reader of that materialized PCollection. The same |
| // is true for the go window. |
| 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() { |
| // (impulse.out) -> parDo -> parDo.out -> window -> window.out |
| // Fuses into |
| // (impulse.out) -> parDo -> (parDo.out) |
| // (parDo.out) -> window -> window.out |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("rare")) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> read -> read.out -> parDo -> parDo.out |
| // \ |
| // -> window -> window.out |
| // Fuses into |
| // (impulse.out) -> read -> (read.out) |
| // (read.out) -> parDo -> parDo.out |
| // (read.out) -> window -> window.out |
| // The window can't be fused into the stage, which forces the PCollection to be materialized. |
| // ParDo in this case _could_ be fused into the stage, but is not for simplicity of |
| // implementation |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("rare")) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> read -> read.out -----------> parDo -> parDo.out -> window -> window.out |
| // (impulse.out) -> side_read -> side_read.out / |
| // Where parDo takes side_read as a side input, fuses into |
| // (impulse.out) -> read -> (read.out) |
| // (impulse.out) -> side_read -> (side_read.out) |
| // (read.out) -> parDo -> parDo.out -> window -> window.out |
| // parDo doesn't have a per-element consumer from side_read.out, so it can't root a stage |
| // which consumes from that materialized collection. Nodes with side inputs must root a stage, |
| // but do not restrict fusion of consumers. |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putSideInputs("side_input", SideInput.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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( |
| SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putSideInputs("side_input", SideInput.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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() { |
| // impulse -- ParDo(createSide) |
| // \_ ParDo(processMain) with side input from createSide |
| // The ExecutableStage executing createSide must have an output. |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putSideInputs("side", SideInput.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .putStateSpecs("state_spec", StateSpec.getDefaultInstance()) |
| .build() |
| .toByteString())) |
| .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() { |
| // (impulse.out) -> read -> read.out -> gbk -> gbk.out |
| // Fuses to |
| // (impulse.out) -> read -> (read.out) |
| // GBK is the responsibility of the runner, so it is not included in a stage. |
| 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(SdkFunctionSpec.newBuilder().setEnvironmentId("common")) |
| .build() |
| .toByteString())) |
| .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())); |
| } |
| |
| private static TypeSafeMatcher<ExecutableStage> hasSubtransforms(String id, String... ids) { |
| Set<String> expectedTransforms = ImmutableSet.<String>builder().add(id).add(ids).build(); |
| return new TypeSafeMatcher<ExecutableStage>() { |
| @Override |
| protected boolean matchesSafely(ExecutableStage executableStage) { |
| // NOTE: Transform names must be unique, so it's fine to throw here if this does not hold. |
| Set<String> stageTransforms = |
| executableStage.getTransforms().stream() |
| .map(PTransformNode::getId) |
| .collect(Collectors.toSet()); |
| return stageTransforms.containsAll(expectedTransforms) |
| && expectedTransforms.containsAll(stageTransforms); |
| } |
| |
| @Override |
| public void describeTo(Description description) { |
| description.appendText("ExecutableStage with subtransform ids: " + expectedTransforms); |
| } |
| }; |
| } |
| } |
