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)); }
private static <T, DestinationT> WriteFilesPayload getWriteFilesPayload( AppliedPTransform< PCollection<T>, WriteFilesResult<DestinationT>, ? extends PTransform<PCollection<T>, WriteFilesResult<DestinationT>>> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return WriteFilesPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()); }
private static <T> ReadPayload getReadPayload( AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return ReadPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()); }
private static <T> ParDoPayload getParDoPayload(AppliedPTransform<?, ?, ?> transform) throws IOException { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); RunnerApi.PTransform parDoPTransform = PTransformTranslation.toProto(transform, Collections.emptyList(), components); return ParDoPayload.parseFrom(parDoPTransform.getSpec().getPayload()); }
public static PCollection.IsBounded sourceIsBounded(AppliedPTransform<?, ?, ?> transform) { try { SdkComponents components = SdkComponents.create(transform.getPipeline().getOptions()); return PCollectionTranslation.fromProto( ReadPayload.parseFrom( PTransformTranslation.toProto(transform, Collections.emptyList(), components) .getSpec() .getPayload()) .getIsBounded()); } catch (IOException e) { throw new RuntimeException("Internal error determining boundedness of Read", e); } }
@Before public void setUp() throws Exception { components = SdkComponents.create(); components.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); }
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(); }
/** Extracts an {@link ExecutableStagePayload} from the given transform. */ public static ExecutableStagePayload getExecutableStagePayload( AppliedPTransform<?, ?, ?> appliedTransform) throws IOException { RunnerApi.PTransform transform = PTransformTranslation.toProto( appliedTransform, SdkComponents.create(appliedTransform.getPipeline().getOptions())); checkArgument(ExecutableStage.URN.equals(transform.getSpec().getUrn())); return ExecutableStagePayload.parseFrom(transform.getSpec().getPayload()); } }
private static Coder<?> getAccumulatorCoder(AppliedPTransform<?, ?, ?> transform) throws IOException { SdkComponents sdkComponents = SdkComponents.create(transform.getPipeline().getOptions()); String id = getCombinePayload(transform, sdkComponents) .map(CombinePayload::getAccumulatorCoderId) .orElseThrow(() -> new IOException("Transform does not contain an AccumulatorCoder")); Components components = sdkComponents.toComponents(); return CoderTranslation.fromProto( components.getCodersOrThrow(id), RehydratedComponents.forComponents(components)); }
@Test public void toAndFromCloudObjectWithSdkComponents() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); CloudObject cloudObject = CloudObjects.asCloudObject(coder, sdkComponents); Coder<?> fromCloudObject = CloudObjects.coderFromCloudObject(cloudObject); assertEquals(coder.getClass(), fromCloudObject.getClass()); assertEquals(coder, fromCloudObject); checkPipelineProtoCoderIds(coder, cloudObject, sdkComponents); }
@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 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 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 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 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 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 testWindowMappingFnTranslation() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); assertEquals( new GlobalWindows().getDefaultWindowMappingFn(), PCollectionViewTranslation.windowMappingFnFromProto( ParDoTranslation.translateWindowMappingFn( new GlobalWindows().getDefaultWindowMappingFn(), sdkComponents))); }