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 void onNext(ArtifactChunk value) { try { target.write(value.getData().toByteArray()); } catch (IOException e) { // This should never happen throw new AssertionError(e); } }
private static Environment createEmbeddedEnvironment(String config) { return Environment.newBuilder() .setUrn(ENVIRONMENT_EMBEDDED) .setPayload(ByteString.copyFromUtf8(MoreObjects.firstNonNull(config, ""))) .build(); }
channel.read(readBuffer); readBuffer.flip(); ByteString chunk = ByteString.copyFrom(readBuffer); hasher.putBytes(chunk.toByteArray()); readBuffer.rewind(); PutArtifactRequest request =
@Override public void onNext(ArtifactApi.PutArtifactRequest value) { try { if (value.getData() == null) { StatusRuntimeException e = Status.INVALID_ARGUMENT .withDescription( String.format( "Expected all chunks in the current stream state to contain data, got %s", value.getContentCase())) .asRuntimeException(); throw e; } value.getData().getData().writeTo(target); } catch (Exception e) { cleanedUp(e); } }
private static SdkFunctionSpec toProto(String urn, Serializable serializable) { return SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(urn) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(serializable))) .build()) .build(); }
@VisibleForTesting static FileBasedSink<?, ?, ?> sinkFromProto(SdkFunctionSpec sinkProto) throws IOException { checkArgument( sinkProto.getSpec().getUrn().equals(CUSTOM_JAVA_FILE_BASED_SINK_URN), "Cannot extract %s instance from %s with URN %s", FileBasedSink.class.getSimpleName(), FunctionSpec.class.getSimpleName(), sinkProto.getSpec().getUrn()); byte[] serializedSink = sinkProto.getSpec().getPayload().toByteArray(); return (FileBasedSink<?, ?, ?>) SerializableUtils.deserializeFromByteArray( serializedSink, FileBasedSink.class.getSimpleName()); }
@Override public void getArtifact( ArtifactApi.GetArtifactRequest request, StreamObserver<ArtifactApi.ArtifactChunk> responseObserver) { try { ByteBuffer artifact = getArtifact(request.getName()); do { responseObserver.onNext( ArtifactChunk.newBuilder() .setData( ByteString.copyFrom( artifact, Math.min(artifact.remaining(), DEFAULT_CHUNK_SIZE))) .build()); } while (artifact.hasRemaining()); responseObserver.onCompleted(); } catch (FileNotFoundException e) { responseObserver.onError( Status.INVALID_ARGUMENT .withDescription(String.format("No such artifact %s", request.getName())) .withCause(e) .asException()); } catch (Exception e) { responseObserver.onError( Status.INTERNAL .withDescription( String.format("Could not retrieve artifact with name %s", request.getName())) .withCause(e) .asException()); } }
public static DoFnAndMainOutput doFnAndMainOutputTagFromProto(SdkFunctionSpec fnSpec) { checkArgument( fnSpec.getSpec().getUrn().equals(CUSTOM_JAVA_DO_FN_URN), "Expected %s to be %s with URN %s, but URN was %s", DoFn.class.getSimpleName(), FunctionSpec.class.getSimpleName(), CUSTOM_JAVA_DO_FN_URN, fnSpec.getSpec().getUrn()); byte[] serializedFn = fnSpec.getSpec().getPayload().toByteArray(); return (DoFnAndMainOutput) SerializableUtils.deserializeFromByteArray(serializedFn, "Custom DoFn And Main Output tag"); }
@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(); } }
private static Coder<?> fromCustomCoder(RunnerApi.Coder protoCoder) throws IOException { return (Coder<?>) SerializableUtils.deserializeFromByteArray( protoCoder.getSpec().getSpec().getPayload().toByteArray(), "Custom Coder Bytes"); } }
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 void onNext(PutArtifactRequest value) { try { stream.write(value.getData().getData().toByteArray()); } catch (IOException e) { throw new RuntimeException(e); } }
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(); }
/** * Converts a {@link org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec} into a {@link * ViewFn} using the URN. */ public static ViewFn<?, ?> viewFnFromProto(RunnerApi.SdkFunctionSpec viewFn) throws InvalidProtocolBufferException { RunnerApi.FunctionSpec spec = viewFn.getSpec(); checkArgument( spec.getUrn().equals(ParDoTranslation.CUSTOM_JAVA_VIEW_FN_URN), "Can't deserialize unknown %s type %s", ViewFn.class.getSimpleName(), spec.getUrn()); return (ViewFn<?, ?>) SerializableUtils.deserializeFromByteArray( spec.getPayload().toByteArray(), "Custom ViewFn"); }
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(); } }
/** * Converts a {@link org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec} into a {@link * WindowMappingFn} using the URN. */ public static WindowMappingFn<?> windowMappingFnFromProto( RunnerApi.SdkFunctionSpec windowMappingFn) throws InvalidProtocolBufferException { RunnerApi.FunctionSpec spec = windowMappingFn.getSpec(); checkArgument( spec.getUrn().equals(ParDoTranslation.CUSTOM_JAVA_WINDOW_MAPPING_FN_URN), "Can't deserialize unknown %s type %s", WindowMappingFn.class.getSimpleName(), spec.getUrn()); return (WindowMappingFn<?>) SerializableUtils.deserializeFromByteArray( spec.getPayload().toByteArray(), "Custom WinodwMappingFn"); } }
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(); }
private void translateStreamingImpulse( String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext context) { RunnerApi.PTransform pTransform = pipeline.getComponents().getTransformsOrThrow(id); TypeInformation<WindowedValue<byte[]>> typeInfo = new CoderTypeInformation<>( WindowedValue.getFullCoder(ByteArrayCoder.of(), GlobalWindow.Coder.INSTANCE)); ObjectMapper objectMapper = new ObjectMapper(); final int intervalMillis; final int messageCount; try { JsonNode config = objectMapper.readTree(pTransform.getSpec().getPayload().toByteArray()); intervalMillis = config.path("interval_ms").asInt(100); messageCount = config.path("message_count").asInt(0); } catch (IOException e) { throw new RuntimeException("Failed to parse configuration for streaming impulse", e); } SingleOutputStreamOperator<WindowedValue<byte[]>> source = context .getExecutionEnvironment() .addSource(new StreamingImpulseSource(intervalMillis, messageCount)) .returns(typeInfo); context.addDataStream(Iterables.getOnlyElement(pTransform.getOutputsMap().values()), source); }
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(); }