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(); }
private static SdkFunctionSpec toProto(String urn, Serializable serializable) { return SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(urn) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(serializable))) .build()) .build(); }
@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()); } }
@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(); } }
channel.read(readBuffer); readBuffer.flip(); ByteString chunk = ByteString.copyFrom(readBuffer);
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(); } }
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(); }
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 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(); }
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(); }
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(); }
@Test public void putArtifactBeforeNameFails() { byte[] data = "foo-".getBytes(UTF_8); RecordingStreamObserver<ArtifactApi.PutArtifactResponse> responseObserver = new RecordingStreamObserver<>(); StreamObserver<ArtifactApi.PutArtifactRequest> requestObserver = stub.putArtifact(responseObserver); requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(data)).build()) .build()); responseObserver.awaitTerminalState(); assertThat(responseObserver.error, Matchers.not(Matchers.nullValue())); }
.setTimestamp(element.getTimestamp().getMillis()) .setEncodedElement( ByteString.copyFrom( CoderUtils.encodeToByteArray(coder, element.getValue()))));
@Test public void singleDataPutArtifactSucceeds() throws Exception { byte[] data = "foo-bar-baz".getBytes(UTF_8); RecordingStreamObserver<ArtifactApi.PutArtifactResponse> responseObserver = new RecordingStreamObserver<>(); StreamObserver<ArtifactApi.PutArtifactRequest> requestObserver = stub.putArtifact(responseObserver); String name = "my-artifact"; requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setMetadata( ArtifactApi.PutArtifactMetadata.newBuilder() .setMetadata(ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build()) .setStagingSessionToken("token") .build()) .build()); requestObserver.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(data)).build()) .build()); requestObserver.onCompleted(); responseObserver.awaitTerminalState(); File staged = stager.getLocation().getArtifactFile(name); assertThat(staged.exists(), is(true)); ByteBuffer buf = ByteBuffer.allocate(data.length); new FileInputStream(staged).getChannel().read(buf); Assert.assertArrayEquals(data, buf.array()); }
.setData( ArtifactApi.ArtifactChunk.newBuilder() .setData(ByteString.copyFrom(partOne)) .build()) .build()); .setData( ArtifactApi.ArtifactChunk.newBuilder() .setData(ByteString.copyFrom(partTwo)) .build()) .build()); .setData( ArtifactApi.ArtifactChunk.newBuilder() .setData(ByteString.copyFrom(partThree)) .build()) .build());
.setUrn(JAVA_SERIALIZED_COMBINE_FN_URN) .setPayload( ByteString.copyFrom( SerializableUtils.serializeToByteArray( combine.getTransform().getFn())))
ByteString serializedFn = ByteString.copyFrom(SerializableUtils.serializeToByteArray(windowFn)); if (windowFn instanceof GlobalWindows) { return SdkFunctionSpec.newBuilder()
.build(); ByteString newPayload = ByteString.copyFrom("foo-bar-baz".getBytes(StandardCharsets.UTF_8)); Pipeline updated = ProtoOverrides.updateTransform(
private ArtifactApi.ArtifactMetadata stageBytes(String name, byte[] bytes) { StreamObserver<ArtifactApi.PutArtifactRequest> requests = stub.putArtifact(new RecordingStreamObserver<>()); requests.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setMetadata( ArtifactApi.PutArtifactMetadata.newBuilder() .setMetadata(ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build()) .setStagingSessionToken("token") .build()) .build()); requests.onNext( ArtifactApi.PutArtifactRequest.newBuilder() .setData( ArtifactApi.ArtifactChunk.newBuilder().setData(ByteString.copyFrom(bytes)).build()) .build()); requests.onCompleted(); return ArtifactApi.ArtifactMetadata.newBuilder().setName(name).build(); }