private static RunnerApi.Coder toKnownCoder(Coder<?> coder, SdkComponents components) throws IOException { CoderTranslator translator = KNOWN_TRANSLATORS.get(coder.getClass()); List<String> componentIds = registerComponents(coder, translator, components); return RunnerApi.Coder.newBuilder() .addAllComponentCoderIds(componentIds) .setSpec( SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(KNOWN_CODER_URNS.get(coder.getClass())) .setPayload(ByteString.copyFrom(translator.getPayload(coder))))) .build(); }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>> transform, SdkComponents components) throws IOException { if (transform.getTransform().getSideInputs().isEmpty()) { return FunctionSpec.newBuilder() .setUrn(COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payloadForCombine((AppliedPTransform) transform, components).toByteString()) .build(); } else { // Combines with side inputs are translated as generic composites, which have a blank // FunctionSpec. return null; } }
@Test public void getEnvironmentRead() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(Read.from(CountingSource.unbounded()), components); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getSource().getEnvironmentId()))); }
public static SdkFunctionSpec translateDoFn( DoFn<?, ?> fn, TupleTag<?> tag, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(CUSTOM_JAVA_DO_FN_URN) .setPayload( ByteString.copyFrom( SerializableUtils.serializeToByteArray(DoFnAndMainOutput.of(fn, tag)))) .build()) .build(); }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Read.Bounded<?>> transform, SdkComponents components) { ReadPayload payload = toProto(transform.getTransform(), components); return RunnerApi.FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload(payload.toByteString()) .build(); } }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Read.Unbounded<?>> transform, SdkComponents components) { ReadPayload payload = toProto(transform.getTransform(), components); return RunnerApi.FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload(payload.toByteString()) .build(); } }
private static SdkFunctionSpec toProto(UnboundedSource<?, ?> source, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(JAVA_SERIALIZED_UNBOUNDED_SOURCE) .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(source))) .build()) .build(); }
@Override public FunctionSpec migrate(SdkComponents components) throws IOException { return FunctionSpec.newBuilder() .setUrn(WRITE_FILES_TRANSFORM_URN) .setPayload(payloadForWriteFilesLike(this, components).toByteString()) .build(); }
@Override public RunnerApi.FunctionSpec translate( AppliedPTransform<?, ?, ParDoSingle<?, ?>> transform, SdkComponents components) throws IOException { RunnerApi.ParDoPayload payload = payloadForParDoSingle(transform.getTransform(), components); return RunnerApi.FunctionSpec.newBuilder() .setUrn(PAR_DO_TRANSFORM_URN) .setPayload(payload.toByteString()) .build(); }
private static RunnerApi.Coder toCustomCoder(Coder<?> coder) throws IOException { RunnerApi.Coder.Builder coderBuilder = RunnerApi.Coder.newBuilder(); return coderBuilder .setSpec( SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(JAVA_SERIALIZED_CODER_URN) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(coder))) .build())) .build(); }
private <T> RunnerApi.FunctionSpec translateTyped( final TestStream<T> testStream, SdkComponents components) throws IOException { return RunnerApi.FunctionSpec.newBuilder() .setUrn(TEST_STREAM_TRANSFORM_URN) .setPayload(payloadForTestStream(testStream, components).toByteString()) .build(); }
private static SdkFunctionSpec toProto(BoundedSource<?> source, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(JAVA_SERIALIZED_BOUNDED_SOURCE) .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(source))) .build()) .build(); }
public static SdkFunctionSpec toProto( GlobalCombineFn<?, ?, ?> combineFn, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(JAVA_SERIALIZED_COMBINE_FN_URN) .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(combineFn))) .build()) .build(); } }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, WriteFiles<?, ?, ?>> transform, SdkComponents components) throws IOException { return FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload(payloadForWriteFiles(transform.getTransform(), components).toByteString()) .build(); } }
public static SdkFunctionSpec translateViewFn(ViewFn<?, ?> viewFn, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(CUSTOM_JAVA_VIEW_FN_URN) .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(viewFn))) .build()) .build(); }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Window.Assign<?>> transform, SdkComponents components) { WindowIntoPayload payload = toProto(transform.getTransform(), components); return RunnerApi.FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload(payload.toByteString()) .build(); } }
private static SdkFunctionSpec toProto(String urn, Serializable serializable) { return SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(urn) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(serializable))) .build()) .build(); }
public static SdkFunctionSpec translateWindowMappingFn( WindowMappingFn<?> windowMappingFn, SdkComponents components) { return SdkFunctionSpec.newBuilder() .setEnvironmentId(components.getOnlyEnvironmentId()) .setSpec( FunctionSpec.newBuilder() .setUrn(CUSTOM_JAVA_WINDOW_MAPPING_FN_URN) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(windowMappingFn))) .build()) .build(); }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, View.CreatePCollectionView<?, ?>> transform, SdkComponents components) { return FunctionSpec.newBuilder() .setUrn(getUrn(transform.getTransform())) .setPayload( ByteString.copyFrom( SerializableUtils.serializeToByteArray(transform.getTransform().getView()))) .build(); } }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Window.Assign<?>> transform, SdkComponents components) { return FunctionSpec.newBuilder() .setUrn("urn:beam:transform:window:v1") .setPayload( WindowIntoTranslation.toProto(transform.getTransform(), components).toByteString()) .build(); } }