| /* |
| * 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.dataflow; |
| |
| import static org.apache.beam.runners.dataflow.util.Structs.getString; |
| import static org.apache.beam.sdk.util.StringUtils.jsonStringToByteArray; |
| import static org.hamcrest.Matchers.equalTo; |
| import static org.hamcrest.Matchers.hasKey; |
| import static org.hamcrest.Matchers.instanceOf; |
| import static org.hamcrest.Matchers.not; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertNotNull; |
| import static org.junit.Assert.assertNull; |
| import static org.junit.Assert.assertThat; |
| import static org.junit.Assert.assertTrue; |
| import static org.mockito.Matchers.any; |
| import static org.mockito.Matchers.argThat; |
| import static org.mockito.Matchers.eq; |
| import static org.mockito.Mockito.mock; |
| import static org.mockito.Mockito.when; |
| |
| import com.google.api.services.dataflow.Dataflow; |
| import com.google.api.services.dataflow.model.Job; |
| import com.google.api.services.dataflow.model.Step; |
| import com.google.api.services.dataflow.model.WorkerPool; |
| import java.io.IOException; |
| import java.io.Serializable; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collection; |
| import java.util.Collections; |
| import java.util.HashSet; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Components; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload; |
| import org.apache.beam.runners.core.construction.PTransformTranslation; |
| import org.apache.beam.runners.core.construction.ParDoTranslation; |
| import org.apache.beam.runners.core.construction.RehydratedComponents; |
| import org.apache.beam.runners.dataflow.DataflowPipelineTranslator.JobSpecification; |
| import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; |
| import org.apache.beam.runners.dataflow.options.DataflowPipelineWorkerPoolOptions; |
| import org.apache.beam.runners.dataflow.util.CloudObject; |
| import org.apache.beam.runners.dataflow.util.CloudObjects; |
| import org.apache.beam.runners.dataflow.util.PropertyNames; |
| import org.apache.beam.runners.dataflow.util.Structs; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.SerializableCoder; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.coders.VarIntCoder; |
| import org.apache.beam.sdk.coders.VoidCoder; |
| import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; |
| import org.apache.beam.sdk.extensions.gcp.util.GcsUtil; |
| import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; |
| import org.apache.beam.sdk.io.FileSystems; |
| import org.apache.beam.sdk.io.TextIO; |
| import org.apache.beam.sdk.io.range.OffsetRange; |
| import org.apache.beam.sdk.options.PipelineOptionsFactory; |
| import org.apache.beam.sdk.options.ValueProvider; |
| import org.apache.beam.sdk.state.StateSpec; |
| import org.apache.beam.sdk.state.StateSpecs; |
| import org.apache.beam.sdk.state.ValueState; |
| import org.apache.beam.sdk.transforms.Count; |
| import org.apache.beam.sdk.transforms.Create; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.PTransform; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.Sum; |
| import org.apache.beam.sdk.transforms.View; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker; |
| import org.apache.beam.sdk.transforms.windowing.FixedWindows; |
| import org.apache.beam.sdk.transforms.windowing.Window; |
| import org.apache.beam.sdk.transforms.windowing.WindowFn; |
| import org.apache.beam.sdk.util.DoFnInfo; |
| import org.apache.beam.sdk.util.SerializableUtils; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollectionTuple; |
| import org.apache.beam.sdk.values.PDone; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.TupleTagList; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableSet; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.hamcrest.Matchers; |
| import org.joda.time.Duration; |
| import org.junit.Assert; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.ExpectedException; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| import org.mockito.ArgumentMatcher; |
| |
| /** Tests for DataflowPipelineTranslator. */ |
| @RunWith(JUnit4.class) |
| public class DataflowPipelineTranslatorTest implements Serializable { |
| @Rule public transient ExpectedException thrown = ExpectedException.none(); |
| |
| // A Custom Mockito matcher for an initial Job that checks that all |
| // expected fields are set. |
| private static class IsValidCreateRequest implements ArgumentMatcher<Job> { |
| @Override |
| public boolean matches(Job o) { |
| Job job = (Job) o; |
| return job.getId() == null |
| && job.getProjectId() == null |
| && job.getName() != null |
| && job.getType() != null |
| && job.getEnvironment() != null |
| && job.getSteps() != null |
| && job.getCurrentState() == null |
| && job.getCurrentStateTime() == null |
| && job.getExecutionInfo() == null |
| && job.getCreateTime() == null; |
| } |
| } |
| |
| private Pipeline buildPipeline(DataflowPipelineOptions options) { |
| options.setRunner(DataflowRunner.class); |
| Pipeline p = Pipeline.create(options); |
| |
| p.apply("ReadMyFile", TextIO.read().from("gs://bucket/object")) |
| .apply("WriteMyFile", TextIO.write().to("gs://bucket/object")); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| runner.replaceTransforms(p); |
| |
| return p; |
| } |
| |
| private static Dataflow buildMockDataflow(ArgumentMatcher<Job> jobMatcher) throws IOException { |
| Dataflow mockDataflowClient = mock(Dataflow.class); |
| Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); |
| Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class); |
| Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); |
| |
| when(mockDataflowClient.projects()).thenReturn(mockProjects); |
| when(mockProjects.jobs()).thenReturn(mockJobs); |
| when(mockJobs.create(eq("someProject"), argThat(jobMatcher))).thenReturn(mockRequest); |
| |
| Job resultJob = new Job(); |
| resultJob.setId("newid"); |
| when(mockRequest.execute()).thenReturn(resultJob); |
| return mockDataflowClient; |
| } |
| |
| private static DataflowPipelineOptions buildPipelineOptions() throws IOException { |
| GcsUtil mockGcsUtil = mock(GcsUtil.class); |
| when(mockGcsUtil.expand(any(GcsPath.class))) |
| .then(invocation -> ImmutableList.of((GcsPath) invocation.getArguments()[0])); |
| when(mockGcsUtil.bucketAccessible(any(GcsPath.class))).thenReturn(true); |
| |
| DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); |
| options.setRunner(DataflowRunner.class); |
| options.setGcpCredential(new TestCredential()); |
| options.setJobName("some-job-name"); |
| options.setProject("some-project"); |
| options.setRegion("some-region"); |
| options.setTempLocation(GcsPath.fromComponents("somebucket", "some/path").toString()); |
| options.setFilesToStage(new ArrayList<>()); |
| options.setDataflowClient(buildMockDataflow(new IsValidCreateRequest())); |
| options.setGcsUtil(mockGcsUtil); |
| |
| // Enable the FileSystems API to know about gs:// URIs in this test. |
| FileSystems.setDefaultPipelineOptions(options); |
| |
| return options; |
| } |
| |
| @Test |
| public void testNetworkConfig() throws IOException { |
| final String testNetwork = "test-network"; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setNetwork(testNetwork); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertEquals(testNetwork, job.getEnvironment().getWorkerPools().get(0).getNetwork()); |
| } |
| |
| @Test |
| public void testNetworkConfigMissing() throws IOException { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertNull(job.getEnvironment().getWorkerPools().get(0).getNetwork()); |
| } |
| |
| @Test |
| public void testSubnetworkConfig() throws IOException { |
| final String testSubnetwork = "regions/REGION/subnetworks/SUBNETWORK"; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setSubnetwork(testSubnetwork); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertEquals(testSubnetwork, job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); |
| } |
| |
| @Test |
| public void testSubnetworkConfigMissing() throws IOException { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertNull(job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); |
| } |
| |
| @Test |
| public void testScalingAlgorithmMissing() throws IOException { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| // Autoscaling settings are always set. |
| assertNull( |
| job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); |
| assertEquals( |
| 0, |
| job.getEnvironment() |
| .getWorkerPools() |
| .get(0) |
| .getAutoscalingSettings() |
| .getMaxNumWorkers() |
| .intValue()); |
| } |
| |
| @Test |
| public void testScalingAlgorithmNone() throws IOException { |
| final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = |
| DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType.NONE; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setAutoscalingAlgorithm(noScaling); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertEquals( |
| "AUTOSCALING_ALGORITHM_NONE", |
| job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); |
| assertEquals( |
| 0, |
| job.getEnvironment() |
| .getWorkerPools() |
| .get(0) |
| .getAutoscalingSettings() |
| .getMaxNumWorkers() |
| .intValue()); |
| } |
| |
| @Test |
| public void testMaxNumWorkersIsPassedWhenNoAlgorithmIsSet() throws IOException { |
| final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = null; |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setMaxNumWorkers(42); |
| options.setAutoscalingAlgorithm(noScaling); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertNull( |
| job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); |
| assertEquals( |
| 42, |
| job.getEnvironment() |
| .getWorkerPools() |
| .get(0) |
| .getAutoscalingSettings() |
| .getMaxNumWorkers() |
| .intValue()); |
| } |
| |
| @Test |
| public void testZoneConfig() throws IOException { |
| final String testZone = "test-zone-1"; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setZone(testZone); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertEquals(testZone, job.getEnvironment().getWorkerPools().get(0).getZone()); |
| } |
| |
| @Test |
| public void testWorkerMachineTypeConfig() throws IOException { |
| final String testMachineType = "test-machine-type"; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setWorkerMachineType(testMachineType); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| |
| WorkerPool workerPool = job.getEnvironment().getWorkerPools().get(0); |
| assertEquals(testMachineType, workerPool.getMachineType()); |
| } |
| |
| @Test |
| public void testDiskSizeGbConfig() throws IOException { |
| final Integer diskSizeGb = 1234; |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| options.setDiskSizeGb(diskSizeGb); |
| |
| Pipeline p = buildPipeline(options); |
| p.traverseTopologically(new RecordingPipelineVisitor()); |
| Job job = |
| DataflowPipelineTranslator.fromOptions(options) |
| .translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()) |
| .getJob(); |
| |
| assertEquals(1, job.getEnvironment().getWorkerPools().size()); |
| assertEquals(diskSizeGb, job.getEnvironment().getWorkerPools().get(0).getDiskSizeGb()); |
| } |
| |
| /** A composite transform that returns an output that is unrelated to the input. */ |
| private static class UnrelatedOutputCreator |
| extends PTransform<PCollection<Integer>, PCollection<Integer>> { |
| |
| @Override |
| public PCollection<Integer> expand(PCollection<Integer> input) { |
| // Apply an operation so that this is a composite transform. |
| input.apply(Count.perElement()); |
| |
| // Return a value unrelated to the input. |
| return input.getPipeline().apply(Create.of(1, 2, 3, 4)); |
| } |
| } |
| |
| /** A composite transform that returns an output that is unbound. */ |
| private static class UnboundOutputCreator extends PTransform<PCollection<Integer>, PDone> { |
| |
| @Override |
| public PDone expand(PCollection<Integer> input) { |
| // Apply an operation so that this is a composite transform. |
| input.apply(Count.perElement()); |
| |
| return PDone.in(input.getPipeline()); |
| } |
| } |
| |
| /** |
| * A composite transform that returns a partially bound output. |
| * |
| * <p>This is not allowed and will result in a failure. |
| */ |
| private static class PartiallyBoundOutputCreator |
| extends PTransform<PCollection<Integer>, PCollectionTuple> { |
| |
| public final TupleTag<Integer> sumTag = new TupleTag<>("sum"); |
| public final TupleTag<Void> doneTag = new TupleTag<>("done"); |
| |
| @Override |
| public PCollectionTuple expand(PCollection<Integer> input) { |
| PCollection<Integer> sum = input.apply(Sum.integersGlobally()); |
| |
| // Fails here when attempting to construct a tuple with an unbound object. |
| return PCollectionTuple.of(sumTag, sum) |
| .and( |
| doneTag, |
| PCollection.createPrimitiveOutputInternal( |
| input.getPipeline(), |
| WindowingStrategy.globalDefault(), |
| input.isBounded(), |
| VoidCoder.of())); |
| } |
| } |
| |
| @Test |
| public void testMultiGraphPipelineSerialization() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| Pipeline p = Pipeline.create(options); |
| |
| PCollection<Integer> input = p.begin().apply(Create.of(1, 2, 3)); |
| |
| input.apply(new UnrelatedOutputCreator()); |
| input.apply(new UnboundOutputCreator()); |
| |
| DataflowPipelineTranslator t = |
| DataflowPipelineTranslator.fromOptions( |
| PipelineOptionsFactory.as(DataflowPipelineOptions.class)); |
| |
| // Check that translation doesn't fail. |
| JobSpecification jobSpecification = |
| t.translate(p, DataflowRunner.fromOptions(options), Collections.emptyList()); |
| assertAllStepOutputsHaveUniqueIds(jobSpecification.getJob()); |
| } |
| |
| @Test |
| public void testPartiallyBoundFailure() throws IOException { |
| Pipeline p = Pipeline.create(buildPipelineOptions()); |
| |
| PCollection<Integer> input = p.begin().apply(Create.of(1, 2, 3)); |
| |
| thrown.expect(IllegalArgumentException.class); |
| input.apply(new PartiallyBoundOutputCreator()); |
| |
| Assert.fail("Failure expected from use of partially bound output"); |
| } |
| |
| /** This tests a few corner cases that should not crash. */ |
| @Test |
| public void testGoodWildcards() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| Pipeline pipeline = Pipeline.create(options); |
| DataflowPipelineTranslator t = DataflowPipelineTranslator.fromOptions(options); |
| |
| applyRead(pipeline, "gs://bucket/foo"); |
| applyRead(pipeline, "gs://bucket/foo/"); |
| applyRead(pipeline, "gs://bucket/foo/*"); |
| applyRead(pipeline, "gs://bucket/foo/?"); |
| applyRead(pipeline, "gs://bucket/foo/[0-9]"); |
| applyRead(pipeline, "gs://bucket/foo/*baz*"); |
| applyRead(pipeline, "gs://bucket/foo/*baz?"); |
| applyRead(pipeline, "gs://bucket/foo/[0-9]baz?"); |
| applyRead(pipeline, "gs://bucket/foo/baz/*"); |
| applyRead(pipeline, "gs://bucket/foo/baz/*wonka*"); |
| applyRead(pipeline, "gs://bucket/foo/*baz/wonka*"); |
| applyRead(pipeline, "gs://bucket/foo*/baz"); |
| applyRead(pipeline, "gs://bucket/foo?/baz"); |
| applyRead(pipeline, "gs://bucket/foo[0-9]/baz"); |
| |
| // Check that translation doesn't fail. |
| JobSpecification jobSpecification = |
| t.translate(pipeline, DataflowRunner.fromOptions(options), Collections.emptyList()); |
| assertAllStepOutputsHaveUniqueIds(jobSpecification.getJob()); |
| } |
| |
| private void applyRead(Pipeline pipeline, String path) { |
| pipeline.apply("Read(" + path + ")", TextIO.read().from(path)); |
| } |
| |
| private static class TestValueProvider implements ValueProvider<String>, Serializable { |
| @Override |
| public boolean isAccessible() { |
| return false; |
| } |
| |
| @Override |
| public String get() { |
| throw new RuntimeException("Should not be called."); |
| } |
| } |
| |
| @Test |
| public void testInaccessibleProvider() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| Pipeline pipeline = Pipeline.create(options); |
| DataflowPipelineTranslator t = DataflowPipelineTranslator.fromOptions(options); |
| |
| pipeline.apply(TextIO.read().from(new TestValueProvider())); |
| |
| // Check that translation does not fail. |
| t.translate(pipeline, DataflowRunner.fromOptions(options), Collections.emptyList()); |
| } |
| |
| /** |
| * Test that in translation the name for a collection (in this case just a Create output) is |
| * overridden to be what the Dataflow service expects. |
| */ |
| @Test |
| public void testNamesOverridden() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| options.setStreaming(false); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| |
| pipeline.apply("Jazzy", Create.of(3)).setName("foobizzle"); |
| |
| runner.replaceTransforms(pipeline); |
| |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| |
| // The Create step |
| Step step = job.getSteps().get(0); |
| |
| // This is the name that is "set by the user" that the Dataflow translator must override |
| String userSpecifiedName = |
| getString( |
| Structs.getListOfMaps(step.getProperties(), PropertyNames.OUTPUT_INFO, null).get(0), |
| PropertyNames.USER_NAME); |
| |
| // This is the calculated name that must actually be used |
| String calculatedName = getString(step.getProperties(), PropertyNames.USER_NAME) + ".out0"; |
| |
| assertThat(userSpecifiedName, equalTo(calculatedName)); |
| } |
| |
| /** |
| * Test that in translation the name for collections of a multi-output ParDo - a special case |
| * because the user can name tags - are overridden to be what the Dataflow service expects. |
| */ |
| @Test |
| public void testTaggedNamesOverridden() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| options.setStreaming(false); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| |
| TupleTag<Integer> tag1 = new TupleTag<Integer>("frazzle") {}; |
| TupleTag<Integer> tag2 = new TupleTag<Integer>("bazzle") {}; |
| TupleTag<Integer> tag3 = new TupleTag<Integer>() {}; |
| |
| PCollectionTuple outputs = |
| pipeline |
| .apply(Create.of(3)) |
| .apply( |
| ParDo.of( |
| new DoFn<Integer, Integer>() { |
| @ProcessElement |
| public void drop() {} |
| }) |
| .withOutputTags(tag1, TupleTagList.of(tag2).and(tag3))); |
| |
| outputs.get(tag1).setName("bizbazzle"); |
| outputs.get(tag2).setName("gonzaggle"); |
| outputs.get(tag3).setName("froonazzle"); |
| |
| runner.replaceTransforms(pipeline); |
| |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| |
| // The ParDo step |
| Step step = job.getSteps().get(1); |
| String stepName = getString(step.getProperties(), PropertyNames.USER_NAME); |
| |
| List<Map<String, Object>> outputInfos = |
| Structs.getListOfMaps(step.getProperties(), PropertyNames.OUTPUT_INFO, null); |
| |
| assertThat(outputInfos.size(), equalTo(3)); |
| |
| // The names set by the user _and_ the tags _must_ be ignored, or metrics will not show up. |
| for (int i = 0; i < outputInfos.size(); ++i) { |
| assertThat( |
| getString(outputInfos.get(i), PropertyNames.USER_NAME), |
| equalTo(String.format("%s.out%s", stepName, i))); |
| } |
| } |
| |
| /** Smoke test to fail fast if translation of a stateful ParDo in batch breaks. */ |
| @Test |
| public void testBatchStatefulParDoTranslation() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| options.setStreaming(false); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| |
| TupleTag<Integer> mainOutputTag = new TupleTag<Integer>() {}; |
| |
| pipeline |
| .apply(Create.of(KV.of(1, 1))) |
| .apply( |
| ParDo.of( |
| new DoFn<KV<Integer, Integer>, Integer>() { |
| @StateId("unused") |
| final StateSpec<ValueState<Integer>> stateSpec = |
| StateSpecs.value(VarIntCoder.of()); |
| |
| @ProcessElement |
| public void process(ProcessContext c) { |
| // noop |
| } |
| }) |
| .withOutputTags(mainOutputTag, TupleTagList.empty())); |
| |
| runner.replaceTransforms(pipeline); |
| |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| |
| // The job should look like: |
| // 0. ParallelRead (Create) |
| // 1. ParDo(ReifyWVs) |
| // 2. GroupByKeyAndSortValuesONly |
| // 3. A ParDo over grouped and sorted KVs that is executed via ungrouping service-side |
| |
| List<Step> steps = job.getSteps(); |
| assertEquals(4, steps.size()); |
| |
| Step createStep = steps.get(0); |
| assertEquals("ParallelRead", createStep.getKind()); |
| |
| Step reifyWindowedValueStep = steps.get(1); |
| assertEquals("ParallelDo", reifyWindowedValueStep.getKind()); |
| |
| Step gbkStep = steps.get(2); |
| assertEquals("GroupByKey", gbkStep.getKind()); |
| |
| Step statefulParDoStep = steps.get(3); |
| assertEquals("ParallelDo", statefulParDoStep.getKind()); |
| assertThat( |
| (String) statefulParDoStep.getProperties().get(PropertyNames.USES_KEYED_STATE), |
| not(equalTo("true"))); |
| } |
| |
| /** Smoke test to fail fast if translation of a splittable ParDo in streaming breaks. */ |
| @Test |
| public void testStreamingSplittableParDoTranslation() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| options.setStreaming(true); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| |
| PCollection<String> windowedInput = |
| pipeline |
| .apply(Create.of("a")) |
| .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1)))); |
| windowedInput.apply(ParDo.of(new TestSplittableFn())); |
| |
| runner.replaceTransforms(pipeline); |
| |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| |
| // The job should contain a SplittableParDo.ProcessKeyedElements step, translated as |
| // "SplittableProcessKeyed". |
| |
| List<Step> steps = job.getSteps(); |
| Step processKeyedStep = null; |
| for (Step step : steps) { |
| if ("SplittableProcessKeyed".equals(step.getKind())) { |
| assertNull(processKeyedStep); |
| processKeyedStep = step; |
| } |
| } |
| assertNotNull(processKeyedStep); |
| |
| @SuppressWarnings({"unchecked", "rawtypes"}) |
| DoFnInfo<String, Integer> fnInfo = |
| (DoFnInfo<String, Integer>) |
| SerializableUtils.deserializeFromByteArray( |
| jsonStringToByteArray( |
| getString(processKeyedStep.getProperties(), PropertyNames.SERIALIZED_FN)), |
| "DoFnInfo"); |
| assertThat(fnInfo.getDoFn(), instanceOf(TestSplittableFn.class)); |
| assertThat( |
| fnInfo.getWindowingStrategy().getWindowFn(), |
| Matchers.<WindowFn>equalTo(FixedWindows.of(Duration.standardMinutes(1)))); |
| assertThat(fnInfo.getInputCoder(), instanceOf(StringUtf8Coder.class)); |
| Coder<?> restrictionCoder = |
| CloudObjects.coderFromCloudObject( |
| (CloudObject) |
| Structs.getObject( |
| processKeyedStep.getProperties(), PropertyNames.RESTRICTION_CODER)); |
| |
| assertEquals(SerializableCoder.of(OffsetRange.class), restrictionCoder); |
| } |
| |
| /** Smoke test to fail fast if translation of a splittable ParDo in FnAPI. */ |
| @Test |
| public void testSplittableParDoTranslationFnApi() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| options.setExperiments(Arrays.asList("beam_fn_api")); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| |
| PCollection<String> windowedInput = |
| pipeline |
| .apply(Create.of("a")) |
| .apply(Window.into(FixedWindows.of(Duration.standardMinutes(1)))); |
| windowedInput.apply(ParDo.of(new TestSplittableFn())); |
| |
| runner.replaceTransforms(pipeline); |
| |
| JobSpecification result = translator.translate(pipeline, runner, Collections.emptyList()); |
| |
| Job job = result.getJob(); |
| |
| // The job should contain a ParDo step, containing a "restriction_encoding". |
| |
| List<Step> steps = job.getSteps(); |
| Step splittableParDo = null; |
| for (Step step : steps) { |
| if ("ParallelDo".equals(step.getKind()) |
| && step.getProperties().containsKey(PropertyNames.RESTRICTION_ENCODING)) { |
| assertNull(splittableParDo); |
| splittableParDo = step; |
| } |
| } |
| assertNotNull(splittableParDo); |
| |
| String fn = Structs.getString(splittableParDo.getProperties(), PropertyNames.SERIALIZED_FN); |
| |
| Components componentsProto = result.getPipelineProto().getComponents(); |
| RehydratedComponents components = RehydratedComponents.forComponents(componentsProto); |
| RunnerApi.PTransform splittableTransform = componentsProto.getTransformsOrThrow(fn); |
| assertEquals( |
| PTransformTranslation.PAR_DO_TRANSFORM_URN, splittableTransform.getSpec().getUrn()); |
| ParDoPayload payload = ParDoPayload.parseFrom(splittableTransform.getSpec().getPayload()); |
| assertThat( |
| ParDoTranslation.doFnWithExecutionInformationFromProto(payload.getDoFn()).getDoFn(), |
| instanceOf(TestSplittableFn.class)); |
| assertThat( |
| components.getCoder(payload.getRestrictionCoderId()), instanceOf(SerializableCoder.class)); |
| |
| // In the Fn API case, we still translate the restriction coder into the RESTRICTION_CODER |
| // property as a CloudObject, and it gets passed through the Dataflow backend, but in the end |
| // the Dataflow worker will end up fetching it from the SPK transform payload instead. |
| Coder<?> restrictionCoder = |
| CloudObjects.coderFromCloudObject( |
| (CloudObject) |
| Structs.getObject( |
| splittableParDo.getProperties(), PropertyNames.RESTRICTION_ENCODING)); |
| assertEquals(SerializableCoder.of(OffsetRange.class), restrictionCoder); |
| } |
| |
| @Test |
| public void testToSingletonTranslationWithIsmSideInput() throws Exception { |
| // A "change detector" test that makes sure the translation |
| // of getting a PCollectionView<T> does not change |
| // in bad ways during refactor |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| pipeline.apply(Create.of(1)).apply(View.asSingleton()); |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| runner.replaceTransforms(pipeline); |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| assertAllStepOutputsHaveUniqueIds(job); |
| |
| List<Step> steps = job.getSteps(); |
| assertEquals(9, steps.size()); |
| |
| @SuppressWarnings("unchecked") |
| List<Map<String, Object>> toIsmRecordOutputs = |
| (List<Map<String, Object>>) steps.get(7).getProperties().get(PropertyNames.OUTPUT_INFO); |
| assertTrue( |
| Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); |
| |
| Step collectionToSingletonStep = steps.get(8); |
| assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); |
| } |
| |
| @Test |
| public void testToIterableTranslationWithIsmSideInput() throws Exception { |
| // A "change detector" test that makes sure the translation |
| // of getting a PCollectionView<Iterable<T>> does not change |
| // in bad ways during refactor |
| |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| |
| Pipeline pipeline = Pipeline.create(options); |
| pipeline.apply(Create.of(1, 2, 3)).apply(View.asIterable()); |
| |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| runner.replaceTransforms(pipeline); |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| assertAllStepOutputsHaveUniqueIds(job); |
| |
| List<Step> steps = job.getSteps(); |
| assertEquals(3, steps.size()); |
| |
| @SuppressWarnings("unchecked") |
| List<Map<String, Object>> toIsmRecordOutputs = |
| (List<Map<String, Object>>) steps.get(1).getProperties().get(PropertyNames.OUTPUT_INFO); |
| assertTrue( |
| Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); |
| |
| Step collectionToSingletonStep = steps.get(2); |
| assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); |
| } |
| |
| @Test |
| public void testStepDisplayData() throws Exception { |
| DataflowPipelineOptions options = buildPipelineOptions(); |
| DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); |
| Pipeline pipeline = Pipeline.create(options); |
| |
| DoFn<Integer, Integer> fn1 = |
| new DoFn<Integer, Integer>() { |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| c.output(c.element()); |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder |
| .add(DisplayData.item("foo", "bar")) |
| .add( |
| DisplayData.item("foo2", DataflowPipelineTranslatorTest.class) |
| .withLabel("Test Class") |
| .withLinkUrl("http://www.google.com")); |
| } |
| }; |
| |
| DoFn<Integer, Integer> fn2 = |
| new DoFn<Integer, Integer>() { |
| @ProcessElement |
| public void processElement(ProcessContext c) throws Exception { |
| c.output(c.element()); |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder.add(DisplayData.item("foo3", 1234)); |
| } |
| }; |
| |
| ParDo.SingleOutput<Integer, Integer> parDo1 = ParDo.of(fn1); |
| ParDo.SingleOutput<Integer, Integer> parDo2 = ParDo.of(fn2); |
| pipeline.apply(Create.of(1, 2, 3)).apply(parDo1).apply(parDo2); |
| |
| DataflowRunner runner = DataflowRunner.fromOptions(options); |
| runner.replaceTransforms(pipeline); |
| Job job = translator.translate(pipeline, runner, Collections.emptyList()).getJob(); |
| assertAllStepOutputsHaveUniqueIds(job); |
| |
| List<Step> steps = job.getSteps(); |
| assertEquals(3, steps.size()); |
| |
| Map<String, Object> parDo1Properties = steps.get(1).getProperties(); |
| Map<String, Object> parDo2Properties = steps.get(2).getProperties(); |
| assertThat(parDo1Properties, hasKey("display_data")); |
| |
| @SuppressWarnings("unchecked") |
| Collection<Map<String, String>> fn1displayData = |
| (Collection<Map<String, String>>) parDo1Properties.get("display_data"); |
| @SuppressWarnings("unchecked") |
| Collection<Map<String, String>> fn2displayData = |
| (Collection<Map<String, String>>) parDo2Properties.get("display_data"); |
| |
| ImmutableSet<ImmutableMap<String, Object>> expectedFn1DisplayData = |
| ImmutableSet.of( |
| ImmutableMap.<String, Object>builder() |
| .put("key", "foo") |
| .put("type", "STRING") |
| .put("value", "bar") |
| .put("namespace", fn1.getClass().getName()) |
| .build(), |
| ImmutableMap.<String, Object>builder() |
| .put("key", "fn") |
| .put("label", "Transform Function") |
| .put("type", "JAVA_CLASS") |
| .put("value", fn1.getClass().getName()) |
| .put("shortValue", fn1.getClass().getSimpleName()) |
| .put("namespace", parDo1.getClass().getName()) |
| .build(), |
| ImmutableMap.<String, Object>builder() |
| .put("key", "foo2") |
| .put("type", "JAVA_CLASS") |
| .put("value", DataflowPipelineTranslatorTest.class.getName()) |
| .put("shortValue", DataflowPipelineTranslatorTest.class.getSimpleName()) |
| .put("namespace", fn1.getClass().getName()) |
| .put("label", "Test Class") |
| .put("linkUrl", "http://www.google.com") |
| .build()); |
| |
| ImmutableSet<ImmutableMap<String, Object>> expectedFn2DisplayData = |
| ImmutableSet.of( |
| ImmutableMap.<String, Object>builder() |
| .put("key", "fn") |
| .put("label", "Transform Function") |
| .put("type", "JAVA_CLASS") |
| .put("value", fn2.getClass().getName()) |
| .put("shortValue", fn2.getClass().getSimpleName()) |
| .put("namespace", parDo2.getClass().getName()) |
| .build(), |
| ImmutableMap.<String, Object>builder() |
| .put("key", "foo3") |
| .put("type", "INTEGER") |
| .put("value", 1234L) |
| .put("namespace", fn2.getClass().getName()) |
| .build()); |
| |
| assertEquals(expectedFn1DisplayData, ImmutableSet.copyOf(fn1displayData)); |
| assertEquals(expectedFn2DisplayData, ImmutableSet.copyOf(fn2displayData)); |
| } |
| |
| private static void assertAllStepOutputsHaveUniqueIds(Job job) throws Exception { |
| List<String> outputIds = new ArrayList<>(); |
| for (Step step : job.getSteps()) { |
| List<Map<String, Object>> outputInfoList = |
| (List<Map<String, Object>>) step.getProperties().get(PropertyNames.OUTPUT_INFO); |
| if (outputInfoList != null) { |
| for (Map<String, Object> outputInfo : outputInfoList) { |
| outputIds.add(getString(outputInfo, PropertyNames.OUTPUT_NAME)); |
| } |
| } |
| } |
| Set<String> uniqueOutputNames = new HashSet<>(outputIds); |
| outputIds.removeAll(uniqueOutputNames); |
| assertTrue(String.format("Found duplicate output ids %s", outputIds), outputIds.isEmpty()); |
| } |
| |
| private static class TestSplittableFn extends DoFn<String, Integer> { |
| @ProcessElement |
| public void process(ProcessContext c, RestrictionTracker<OffsetRange, Long> tracker) { |
| // noop |
| } |
| |
| @GetInitialRestriction |
| public OffsetRange getInitialRange(String element) { |
| return null; |
| } |
| } |
| } |
