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); } }
/** Translates a {@link Pipeline} into a {@code JobSpecification}. */ public JobSpecification translate( Pipeline pipeline, DataflowRunner runner, List<DataflowPackage> packages) { // Capture the sdkComponents for look up during step translations SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, sdkComponents, true); LOG.debug("Portable pipeline proto:\n{}", TextFormat.printToString(pipelineProto)); Translator translator = new Translator(pipeline, runner, sdkComponents); Job result = translator.translate(packages); return new JobSpecification( result, pipelineProto, Collections.unmodifiableMap(translator.stepNames)); }
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; }
@Before public void setUp() throws Exception { components = SdkComponents.create(); components.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); }
@Test public void testToFromProtoUnbounded() throws Exception { assumeThat(source, instanceOf(UnboundedSource.class)); UnboundedSource<?, ?> unboundedSource = (UnboundedSource<?, ?>) this.source; Read.Unbounded<?> unboundedRead = Read.from(unboundedSource); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(unboundedRead, components); assertThat(payload.getIsBounded(), equalTo(RunnerApi.IsBounded.Enum.UNBOUNDED)); UnboundedSource<?, ?> deserializedSource = ReadTranslation.unboundedSourceFromProto(payload); assertThat(deserializedSource, equalTo(source)); }
@Test public void testToFromProtoBounded() throws Exception { // TODO: Split into two tests. assumeThat(source, instanceOf(BoundedSource.class)); BoundedSource<?> boundedSource = (BoundedSource<?>) this.source; Read.Bounded<?> boundedRead = Read.from(boundedSource); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(boundedRead, components); assertThat(payload.getIsBounded(), equalTo(RunnerApi.IsBounded.Enum.BOUNDED)); BoundedSource<?> deserializedSource = ReadTranslation.boundedSourceFromProto(payload); assertThat(deserializedSource, equalTo(source)); }
@Test public void testEnvironment() { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Environment env = Environments.createDockerEnvironment("java_test"); String id = sdkComponents.registerEnvironment(env); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Environment rehydratedEnv = rehydratedComponents.getEnvironment(id); assertThat(rehydratedEnv, equalTo(env)); assertThat(rehydratedComponents.getEnvironment(id), theInstance(rehydratedEnv)); } }
@Test public void testToProtoAndBack() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy<?, ?> toProtoAndBackWindowingStrategy = WindowingStrategyTranslation.fromProto( WindowingStrategyTranslation.toMessageProto(windowingStrategy, components)); assertThat( toProtoAndBackWindowingStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); }
@Test public void testStateSpecToFromProto() throws Exception { // Encode SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.StateSpec stateSpecProto = ParDoTranslation.translateStateSpec(stateSpec, sdkComponents); // Decode RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); StateSpec<?> deserializedStateSpec = ParDoTranslation.fromProto(stateSpecProto, rehydratedComponents); assertThat(stateSpec, equalTo(deserializedStateSpec)); } }
@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 testToProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ParDoPayload payload = ParDoTranslation.translateParDo(parDo, p, components); assertThat(ParDoTranslation.getDoFn(payload), equalTo(parDo.getFn())); assertThat(ParDoTranslation.getMainOutputTag(payload), equalTo(parDo.getMainOutputTag())); for (PCollectionView<?> view : parDo.getSideInputs()) { payload.getSideInputsOrThrow(view.getTagInternal().getId()); } }
@Test public void testViewFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new TestViewFn(), PCollectionViewTranslation.viewFnFromProto( ParDoTranslation.translateViewFn(new TestViewFn(), sdkComponents))); }
@Test public void testEncodedProto() throws Exception { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.WriteFilesPayload payload = WriteFilesTranslation.payloadForWriteFiles(writeFiles, components); assertThat( payload.getRunnerDeterminedSharding(), equalTo( writeFiles.getNumShardsProvider() == null && writeFiles.getComputeNumShards() == null)); assertThat(payload.getWindowedWrites(), equalTo(writeFiles.getWindowedWrites())); assertThat( (FileBasedSink<String, Void, String>) WriteFilesTranslation.sinkFromProto(payload.getSink()), equalTo(writeFiles.getSink())); }
@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)); }
@Test public void getEnvironmentUnknownFnType() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN) .build()) .build(); Optional<Environment> env = Environments.getEnvironment(builder, rehydratedComponents); assertThat(env.isPresent(), is(false)); }
@Test public void testEncodeDecodeFields() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.PCollection protoCollection = PCollectionTranslation.toProto(testCollection, sdkComponents); RehydratedComponents protoComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> decodedCoder = protoComponents.getCoder(protoCollection.getCoderId()); WindowingStrategy<?, ?> decodedStrategy = protoComponents.getWindowingStrategy(protoCollection.getWindowingStrategyId()); IsBounded decodedIsBounded = PCollectionTranslation.isBounded(protoCollection); assertThat(decodedCoder, equalTo(testCollection.getCoder())); assertThat(decodedStrategy, equalTo(testCollection.getWindowingStrategy().fixDefaults())); assertThat(decodedIsBounded, equalTo(testCollection.isBounded())); }
@Test public void testWindowMappingFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new GlobalWindows().getDefaultWindowMappingFn(), PCollectionViewTranslation.windowMappingFnFromProto( ParDoTranslation.translateWindowMappingFn( new GlobalWindows().getDefaultWindowMappingFn(), sdkComponents))); }
@Test public void testToProtoAndBackWithComponents() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.WindowingStrategy proto = WindowingStrategyTranslation.toProto(windowingStrategy, components); RehydratedComponents protoComponents = RehydratedComponents.forComponents(components.toComponents()); assertThat( WindowingStrategyTranslation.fromProto(proto, protoComponents).fixDefaults(), equalTo(windowingStrategy.fixDefaults())); protoComponents.getCoder( components.registerCoder(windowingStrategy.getWindowFn().windowCoder())); assertThat( proto.getAssignsToOneWindow(), equalTo(windowingStrategy.getWindowFn().assignsToOneWindow())); } }
@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())); }
@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)); }