@Override public void visitPrimitiveTransform(Node node) { // TODO: Include DisplayData in the proto children.put(node.getEnclosingNode(), node.toAppliedPTransform(pipeline)); try { components.registerPTransform( node.toAppliedPTransform(pipeline), Collections.emptyList()); } catch (IOException e) { throw new IllegalStateException(e); } } });
@Override public void leaveCompositeTransform(Node node) { if (node.isRootNode()) { for (AppliedPTransform<?, ?, ?> pipelineRoot : children.get(node)) { rootIds.add(components.getExistingPTransformId(pipelineRoot)); } } else { // TODO: Include DisplayData in the proto children.put(node.getEnclosingNode(), node.toAppliedPTransform(pipeline)); try { components.registerPTransform( node.toAppliedPTransform(pipeline), children.get(node)); } catch (IOException e) { throw new RuntimeException(e); } } }
@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)); }
private RunnerApi.PTransform convert(ToAndFromProtoSpec spec, SdkComponents components) throws IOException { List<AppliedPTransform<?, ?, ?>> childTransforms = new ArrayList<>(); for (ToAndFromProtoSpec child : spec.getChildren()) { childTransforms.add(child.getTransform()); System.out.println("Converting child " + child); convert(child, components); // Sanity call components.getExistingPTransformId(child.getTransform()); } RunnerApi.PTransform convert = PTransformTranslation.toProto(spec.getTransform(), childTransforms, components); // Make sure the converted transform is registered. Convert it independently, but if this is a // child spec, the child must be in the components. components.registerPTransform(spec.getTransform(), childTransforms); return convert; }
@Test public void registerTransformAfterChildren() throws IOException { Create.Values<Long> create = Create.of(1L, 2L, 3L); GenerateSequence createChild = GenerateSequence.from(0); PCollection<Long> pt = pipeline.apply(create); String userName = "my_transform"; String childUserName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); AppliedPTransform<?, ?, ?> childTransform = AppliedPTransform.of( childUserName, pipeline.begin().expand(), pt.expand(), createChild, pipeline); String childId = components.registerPTransform(childTransform, Collections.emptyList()); String parentId = components.registerPTransform(transform, Collections.singletonList(childTransform)); Components components = this.components.toComponents(); assertThat(components.getTransformsOrThrow(parentId).getSubtransforms(0), equalTo(childId)); assertThat(components.getTransformsOrThrow(childId).getSubtransformsCount(), equalTo(0)); }
/** Tests that trying to register a transform which has unregistered children throws. */ @Test public void registerTransformWithUnregisteredChildren() throws IOException { Create.Values<Long> create = Create.of(1L, 2L, 3L); GenerateSequence createChild = GenerateSequence.from(0); PCollection<Long> pt = pipeline.apply(create); String userName = "my_transform"; String childUserName = "my_transform/my_nesting"; AppliedPTransform<?, ?, ?> transform = AppliedPTransform.of(userName, pipeline.begin().expand(), pt.expand(), create, pipeline); AppliedPTransform<?, ?, ?> childTransform = AppliedPTransform.of( childUserName, pipeline.begin().expand(), pt.expand(), createChild, pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage(childTransform.toString()); components.registerPTransform(transform, Collections.singletonList(childTransform)); }
@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); }