@Test public void testEncodedProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.TestStreamPayload payload = TestStreamTranslation.payloadForTestStream(testStream, components); verifyTestStreamEncoding( testStream, payload, RehydratedComponents.forComponents(components.toComponents())); }
@Test public void testSimpleCoder() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Coder<?> coder = VarIntCoder.of(); String id = sdkComponents.registerCoder(coder); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> rehydratedCoder = rehydratedComponents.getCoder(id); assertThat(rehydratedCoder, equalTo((Coder) coder)); assertThat(rehydratedComponents.getCoder(id), theInstance((Coder) rehydratedCoder)); }
@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); } } }
private static byte[] serializeWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) { try { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); return WindowingStrategyTranslation.toMessageProto(windowingStrategy, sdkComponents) .toByteArray(); } catch (Exception e) { throw new RuntimeException( String.format("Unable to format windowing strategy %s as bytes", windowingStrategy), e); } }
public static SdkComponents create(PipelineOptions options) { SdkComponents sdkComponents = new SdkComponents(); PortablePipelineOptions portablePipelineOptions = options.as(PortablePipelineOptions.class); sdkComponents.registerEnvironment( Environments.createOrGetDefaultEnvironment( portablePipelineOptions.getDefaultEnvironmentType(), portablePipelineOptions.getDefaultEnvironmentConfig())); return sdkComponents; }
public static RunnerApi.MessageWithComponents toProto(Coder<?> coder) throws IOException { SdkComponents components = SdkComponents.create(); RunnerApi.Coder coderProto = toProto(coder, components); return RunnerApi.MessageWithComponents.newBuilder() .setCoder(coderProto) .setComponents(components.toComponents()) .build(); }
public static RunnerApi.Pipeline toProto(Pipeline pipeline, boolean useDeprecatedViewTransforms) { return toProto( pipeline, SdkComponents.create(pipeline.getOptions()), useDeprecatedViewTransforms); }
@Test public void testEncodedProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); components.registerPCollection(testPCollection); AppliedPTransform<?, ?, ?> appliedPTransform = AppliedPTransform.of( "foo", testPCollection.expand(), createViewTransform.getView().expand(), createViewTransform, p); FunctionSpec payload = PTransformTranslation.toProto(appliedPTransform, components).getSpec(); // Checks that the payload is what it should be PCollectionView<?> deserializedView = (PCollectionView<?>) SerializableUtils.deserializeFromByteArray( payload.getPayload().toByteArray(), PCollectionView.class.getSimpleName()); assertThat(deserializedView, Matchers.equalTo(createViewTransform.getView())); }
@Test public void testWindowingStrategy() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy windowingStrategy = WindowingStrategy.of(FixedWindows.of(Duration.millis(1))) .withAllowedLateness(Duration.standardSeconds(4)); String id = sdkComponents.registerWindowingStrategy(windowingStrategy); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); WindowingStrategy<?, ?> rehydratedStrategy = rehydratedComponents.getWindowingStrategy(id); assertThat(rehydratedStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); assertThat( rehydratedComponents.getWindowingStrategy(id), theInstance((WindowingStrategy) rehydratedStrategy)); }
private static String registerCoderOrThrow(SdkComponents components, Coder coder) { try { return components.registerCoder(coder); } catch (IOException exc) { throw new RuntimeException("Failure to register coder", exc); } }
PCollectionTuple output = mainInput.apply(parDo); SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); "foo", inputs, output.expand(), parDo, p), sdkComponents); RunnerApi.Components components = sdkComponents.toComponents(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components); assertThat(restoredView.getCoderInternal(), equalTo(view.getCoderInternal())); String mainInputId = sdkComponents.registerPCollection(mainInput); assertThat( ParDoTranslation.getMainInput(protoTransform, components),
@Test public void toAndFromProto() throws IOException { SdkComponents components = SdkComponents.create(spec.getTransform().getPipeline().getOptions()); RunnerApi.PTransform converted = convert(spec, components); Components protoComponents = components.toComponents(); // Sanity checks assertThat(converted.getInputsCount(), equalTo(spec.getTransform().getInputs().size())); assertThat(converted.getOutputsCount(), equalTo(spec.getTransform().getOutputs().size())); assertThat(converted.getSubtransformsCount(), equalTo(spec.getChildren().size())); assertThat(converted.getUniqueName(), equalTo(spec.getTransform().getFullName())); for (PValue inputValue : spec.getTransform().getInputs().values()) { PCollection<?> inputPc = (PCollection<?>) inputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(inputPc)); } for (PValue outputValue : spec.getTransform().getOutputs().values()) { PCollection<?> outputPc = (PCollection<?>) outputValue; protoComponents.getPcollectionsOrThrow(components.registerPCollection(outputPc)); } }
@Test public void registerCoder() throws IOException { Coder<?> coder = KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(SetCoder.of(ByteArrayCoder.of()))); String id = components.registerCoder(coder); assertThat(components.registerCoder(coder), equalTo(id)); assertThat(id, not(isEmptyOrNullString())); VarLongCoder otherCoder = VarLongCoder.of(); assertThat(components.registerCoder(otherCoder), not(equalTo(id))); components.toComponents().getCodersOrThrow(id); components.toComponents().getCodersOrThrow(components.registerCoder(otherCoder)); }
public static RunnerApi.PCollection toProto(PCollection<?> pCollection, SdkComponents components) throws IOException { String coderId = components.registerCoder(pCollection.getCoder()); String windowingStrategyId = components.registerWindowingStrategy(pCollection.getWindowingStrategy()); // TODO: Display Data return RunnerApi.PCollection.newBuilder() .setUniqueName(pCollection.getName()) .setCoderId(coderId) .setIsBounded(toProto(pCollection.isBounded())) .setWindowingStrategyId(windowingStrategyId) .build(); }
@Test public void registerPCollectionExistingNameCollision() throws IOException { PCollection<Long> pCollection = pipeline.apply("FirstCount", GenerateSequence.from(0)).setName("foo"); String firstId = components.registerPCollection(pCollection); PCollection<Long> duplicate = pipeline.apply("SecondCount", GenerateSequence.from(0)).setName("foo"); String secondId = components.registerPCollection(duplicate); assertThat(firstId, equalTo("foo")); assertThat(secondId, containsString("foo")); assertThat(secondId, not(equalTo("foo"))); components.toComponents().getPcollectionsOrThrow(firstId); components.toComponents().getPcollectionsOrThrow(secondId); }
return RunnerApi.CombinePayload.newBuilder() .setAccumulatorCoderId( components.registerCoder( extractAccumulatorCoder(combineFn, (AppliedPTransform) combine))) .setCombineFn( SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder()
transformBuilder.putInputs( toProto(taggedInput.getKey()), components.registerPCollection((PCollection<?>) taggedInput.getValue())); transformBuilder.putOutputs( toProto(taggedOutput.getKey()), components.registerPCollection((PCollection<?>) taggedOutput.getValue())); transformBuilder.addSubtransforms(components.getExistingPTransformId(subtransform));
@Test public void registerWindowingStrategy() throws IOException { WindowingStrategy<?, ?> strategy = WindowingStrategy.globalDefault().withMode(AccumulationMode.ACCUMULATING_FIRED_PANES); String name = components.registerWindowingStrategy(strategy); assertThat(name, not(isEmptyOrNullString())); components.toComponents().getWindowingStrategiesOrThrow(name); }
@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)); }
@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); } } });