private AppliedPTransform<?, ?, ?> appliedWrite(WriteFiles<Integer, Void, Integer> write) { return AppliedPTransform.of( "WriteFiles", Collections.emptyMap(), Collections.emptyMap(), write, p); }
/** Returns the {@link AppliedPTransform} representing this {@link Node}. */ public AppliedPTransform<?, ?, ?> toAppliedPTransform(Pipeline pipeline) { return AppliedPTransform.of( getFullName(), inputs, outputs, (PTransform) getTransform(), pipeline); }
private AppliedPTransform<?, ?, ?> appliedPTransform( String fullName, PTransform<PInput, POutput> transform, Pipeline p) { PInput input = mock(PInput.class); when(input.getPipeline()).thenReturn(p); return AppliedPTransform.of( fullName, Collections.emptyMap(), Collections.emptyMap(), transform, p); }
@Test public void nonAdditionalInputsWithNoInputSucceeds() { AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "input-free", Collections.emptyMap(), Collections.emptyMap(), new TestTransform(), pipeline); assertThat(TransformInputs.nonAdditionalInputs(transform), Matchers.empty()); }
@Test public void nonAdditionalInputsWithOneMainInputSucceeds() { PCollection<Long> input = pipeline.apply(GenerateSequence.from(1L)); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "input-single", Collections.singletonMap(new TupleTag<Long>() {}, input), Collections.emptyMap(), new TestTransform(), pipeline); assertThat(TransformInputs.nonAdditionalInputs(transform), Matchers.containsInAnyOrder(input)); }
@Test public void nonAdditionalInputsWithMultipleNonAdditionalInputsSucceeds() { Map<TupleTag<?>, PValue> allInputs = new HashMap<>(); PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3)); allInputs.put(new TupleTag<Integer>() {}, mainInts); PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of())); allInputs.put(new TupleTag<Void>() {}, voids); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "additional-free", allInputs, Collections.emptyMap(), new TestTransform(), pipeline); assertThat( TransformInputs.nonAdditionalInputs(transform), Matchers.containsInAnyOrder(voids, mainInts)); }
@Test public void getMainInputSingleOutputSingleInput() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", Collections.singletonMap(new TupleTag<Long>(), mainInput), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }
@Test public void getInputEmptySucceeds() { PTransformReplacement<PCollectionList<Long>, PCollection<Long>> replacement = factory.getReplacementTransform( AppliedPTransform.of( "nonEmptyInput", Collections.emptyMap(), Collections.emptyMap(), Flatten.pCollections(), pipeline)); assertThat(replacement.getInput().getAll(), emptyIterable()); }
@Test public void registerTransformNoChildren() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); String userName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); String componentName = components.registerPTransform(transform, Collections.emptyList()); assertThat(componentName, equalTo(userName)); assertThat(components.getExistingPTransformId(transform), equalTo(componentName)); }
@Test public void emptyFlattenWithNonFlatten() { AppliedPTransform application = AppliedPTransform .<PCollection<Iterable<Integer>>, PCollection<Integer>, Flatten.Iterables<Integer>>of( "EmptyFlatten", Collections.emptyMap(), Collections.singletonMap( new TupleTag<Integer>(), PCollection.createPrimitiveOutputInternal( p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())), /* This isn't actually possible to construct, but for the sake of example */ Flatten.iterables(), p); assertThat(PTransformMatchers.emptyFlatten().matches(application), is(false)); }
@Test public void emptyFlattenWithEmptyFlatten() { AppliedPTransform application = AppliedPTransform.of( "EmptyFlatten", Collections.emptyMap(), Collections.singletonMap( new TupleTag<Integer>(), PCollection.createPrimitiveOutputInternal( p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())), Flatten.pCollections(), p); assertThat(PTransformMatchers.emptyFlatten().matches(application), is(true)); }
@Test public void flattenWithDuplicateInputsNonFlatten() { AppliedPTransform application = AppliedPTransform .<PCollection<Iterable<Integer>>, PCollection<Integer>, Flatten.Iterables<Integer>>of( "EmptyFlatten", Collections.emptyMap(), Collections.singletonMap( new TupleTag<Integer>(), PCollection.createPrimitiveOutputInternal( p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())), /* This isn't actually possible to construct, but for the sake of example */ Flatten.iterables(), p); assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false)); }
private static AppliedPTransform<?, ?, ?> generateSequence(Pipeline pipeline) { GenerateSequence sequence = GenerateSequence.from(0); PCollection<Long> pcollection = pipeline.apply(sequence); return AppliedPTransform.of( "Count", pipeline.begin().expand(), pcollection.expand(), sequence, pipeline); }
private static AppliedPTransform<?, ?, ?> read(Pipeline pipeline) { Read.Unbounded<Long> transform = Read.from(CountingSource.unbounded()); PCollection<Long> pcollection = pipeline.apply(transform); return AppliedPTransform.of( "ReadTheCount", pipeline.begin().expand(), pcollection.expand(), transform, pipeline); }
@Test public void registerTransformEmptyFullName() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of("", pipeline.begin().expand(), pt.expand(), create, pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage(transform.toString()); components.getExistingPTransformId(transform); }
@Test public void flattenWithDuplicateInputsWithoutDuplicates() { AppliedPTransform application = AppliedPTransform.of( "Flatten", Collections.singletonMap( new TupleTag<Integer>(), PCollection.createPrimitiveOutputInternal( p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())), Collections.singletonMap( new TupleTag<Integer>(), PCollection.createPrimitiveOutputInternal( p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())), Flatten.pCollections(), p); assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false)); }
@Test public void getMainInputSingleOutputSideInputs() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", ImmutableMap.<TupleTag<?>, PValue>builder() .put(new TupleTag<Long>(), mainInput) .put(sideInput.getTagInternal(), sideInput.getPCollection()) .build(), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()).withSideInputs(sideInput), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }
@Test public void registerTransformNullComponents() throws IOException { Create.Values<Integer> create = Create.of(1, 2, 3); PCollection<Integer> pt = pipeline.apply(create); String userName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); thrown.expect(NullPointerException.class); thrown.expectMessage("child nodes may not be null"); components.registerPTransform(transform, null); }
private PCollection<String> applySplittableParDo( String name, PCollection<Integer> input, DoFn<Integer, String> fn) { ParDo.MultiOutput<Integer, String> multiOutput = ParDo.of(fn).withOutputTags(MAIN_OUTPUT_TAG, TupleTagList.empty()); PCollectionTuple output = multiOutput.expand(input); output.get(MAIN_OUTPUT_TAG).setName("main"); AppliedPTransform<PCollection<Integer>, PCollectionTuple, ?> transform = AppliedPTransform.of("ParDo", input.expand(), output.expand(), multiOutput, pipeline); return input.apply(name, SplittableParDo.forAppliedParDo(transform)).get(MAIN_OUTPUT_TAG); }
@Test public void testRegistrarEncodedProto() throws Exception { PCollection<String> output = p.apply(testStream); AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream = AppliedPTransform.of("fakeName", PBegin.in(p).expand(), output.expand(), testStream, p); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.FunctionSpec spec = PTransformTranslation.toProto(appliedTestStream, components).getSpec(); assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN)); RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload()); verifyTestStreamEncoding( testStream, payload, RehydratedComponents.forComponents(components.toComponents())); }