/** * Stages {@link DataflowPipelineOptions#getFilesToStage()}, which defaults to every file on the * classpath unless overridden, as well as {@link * DataflowPipelineDebugOptions#getOverrideWindmillBinary()} if specified. * * @see #stageFiles(List) */ @Override public List<DataflowPackage> stageDefaultFiles() { checkNotNull(options.getStagingLocation()); String windmillBinary = options.as(DataflowPipelineDebugOptions.class).getOverrideWindmillBinary(); String dataflowWorkerJar = options.getDataflowWorkerJar(); List<String> filesToStage = options.getFilesToStage(); if (windmillBinary != null) { filesToStage.add("windmill_main=" + windmillBinary); } if (dataflowWorkerJar != null && !dataflowWorkerJar.isEmpty()) { filesToStage.add("dataflow-worker.jar=" + dataflowWorkerJar); } return stageFiles(filesToStage); }
dataflowOptions.as(DataflowPipelineDebugOptions.class);
@Test public void testSettingOfPipelineOptionsWithCustomUserType() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); options .as(JacksonIncompatibleOptions.class) .setJacksonIncompatible(new JacksonIncompatible("userCustomTypeTest")); Pipeline p = Pipeline.create(options); p.run(); ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture()); Map<String, Object> sdkPipelineOptions = jobCaptor.getValue().getEnvironment().getSdkPipelineOptions(); assertThat(sdkPipelineOptions, hasKey("options")); Map<String, Object> optionsMap = (Map<String, Object>) sdkPipelineOptions.get("options"); assertThat(optionsMap, hasEntry("jacksonIncompatible", (Object) "userCustomTypeTest")); }
/** Tests that all reads are consumed by at least one {@link PTransform}. */ @Test public void testUnconsumedReads() throws IOException { DataflowPipelineOptions dataflowOptions = buildPipelineOptions(); RuntimeTestOptions options = dataflowOptions.as(RuntimeTestOptions.class); Pipeline p = buildDataflowPipeline(dataflowOptions); p.apply(TextIO.read().from(options.getInput())); DataflowRunner.fromOptions(dataflowOptions).replaceTransforms(p); final AtomicBoolean unconsumedSeenAsInput = new AtomicBoolean(); p.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { unconsumedSeenAsInput.set(true); } }); assertThat(unconsumedSeenAsInput.get(), is(true)); }
DataflowPipelineOptions dataflowOptions = options.as(DataflowPipelineOptions.class); maybeRegisterDebuggee(dataflowOptions, requestId);
@Test public void testTextIOWithRuntimeParameters() throws IOException { DataflowPipelineOptions dataflowOptions = buildPipelineOptions(); RuntimeTestOptions options = dataflowOptions.as(RuntimeTestOptions.class); Pipeline p = buildDataflowPipeline(dataflowOptions); p.apply(TextIO.read().from(options.getInput())).apply(TextIO.write().to(options.getOutput())); }
@Test public void testInvalidNumberOfWorkerHarnessThreads() throws IOException { DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); FileSystems.setDefaultPipelineOptions(options); options.setRunner(DataflowRunner.class); options.setProject("foo-12345"); options.setGcpTempLocation(VALID_TEMP_BUCKET); options.setGcsUtil(mockGcsUtil); options.as(DataflowPipelineDebugOptions.class).setNumberOfWorkerHarnessThreads(-1); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Number of worker harness threads"); thrown.expectMessage("Please make sure the value is non-negative."); DataflowRunner.fromOptions(options); }
private <K, V> void groupByKeyHelper( GroupByKey<K, V> transform, TranslationContext context) { StepTranslationContext stepContext = context.addStep(transform, "GroupByKey"); PCollection<KV<K, V>> input = context.getInput(transform); stepContext.addInput(PropertyNames.PARALLEL_INPUT, input); stepContext.addOutput(PropertyNames.OUTPUT, context.getOutput(transform)); WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy(); boolean isStreaming = context.getPipelineOptions().as(StreamingOptions.class).isStreaming(); boolean allowCombinerLifting = windowingStrategy.getWindowFn().isNonMerging() && windowingStrategy.getWindowFn().assignsToOneWindow(); if (isStreaming) { allowCombinerLifting &= transform.fewKeys(); // TODO: Allow combiner lifting on the non-default trigger, as appropriate. allowCombinerLifting &= (windowingStrategy.getTrigger() instanceof DefaultTrigger); } stepContext.addInput(PropertyNames.DISALLOW_COMBINER_LIFTING, !allowCombinerLifting); stepContext.addInput( PropertyNames.SERIALIZED_FN, byteArrayToJsonString(serializeWindowingStrategy(windowingStrategy))); stepContext.addInput( PropertyNames.IS_MERGING_WINDOW_FN, !windowingStrategy.getWindowFn().isNonMerging()); } });