@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()); }
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"); }
/** * 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 WindowFn<?, ?> windowFnFromProto(SdkFunctionSpec windowFnSpec) { try { String s = windowFnSpec.getSpec().getUrn(); if (s.equals(getUrn(GlobalWindowsPayload.Enum.PROPERTIES))) { return new GlobalWindows(); } else if (s.equals(getUrn(FixedWindowsPayload.Enum.PROPERTIES))) { FixedWindowsPayload fixedParams = FixedWindowsPayload.parseFrom(windowFnSpec.getSpec().getPayload()); return FixedWindows.of(Duration.millis(Durations.toMillis(fixedParams.getSize()))) .withOffset(Duration.millis(Timestamps.toMillis(fixedParams.getOffset()))); } else if (s.equals(getUrn(SlidingWindowsPayload.Enum.PROPERTIES))) { SlidingWindowsPayload slidingParams = SlidingWindowsPayload.parseFrom(windowFnSpec.getSpec().getPayload()); return SlidingWindows.of(Duration.millis(Durations.toMillis(slidingParams.getSize()))) .every(Duration.millis(Durations.toMillis(slidingParams.getPeriod()))) } else if (s.equals(getUrn(SessionsPayload.Enum.PROPERTIES))) { SessionsPayload sessionParams = SessionsPayload.parseFrom(windowFnSpec.getSpec().getPayload()); return Sessions.withGapDuration( Duration.millis(Durations.toMillis(sessionParams.getGapSize()))); return (WindowFn<?, ?>) SerializableUtils.deserializeFromByteArray( windowFnSpec.getSpec().getPayload().toByteArray(), "WindowFn"); } else { throw new IllegalArgumentException( "Unknown or unsupported WindowFn: " + windowFnSpec.getSpec().getUrn()); FunctionSpec.class.getSimpleName(),
return StateSpecs.bag(components.getCoder(stateSpec.getBagSpec().getElementCoderId())); case COMBINING_SPEC: FunctionSpec combineFnSpec = stateSpec.getCombiningSpec().getCombineFn().getSpec();
combineFn, SerializableUtils.deserializeFromByteArray( combineProto.getCombineFn().getSpec().getPayload().toByteArray(), "CombineFn"));
combineFn, SerializableUtils.deserializeFromByteArray( combineProto.getCombineFn().getSpec().getPayload().toByteArray(), "CombineFn"));
private static Coder<?> fromKnownCoder(RunnerApi.Coder coder, RehydratedComponents components) throws IOException { String coderUrn = coder.getSpec().getSpec().getUrn(); List<Coder<?>> coderComponents = new ArrayList<>(); for (String componentId : coder.getComponentCoderIdsList()) { Coder<?> innerCoder = components.getCoder(componentId); coderComponents.add(innerCoder); } Class<? extends Coder> coderType = KNOWN_CODER_URNS.inverse().get(coderUrn); CoderTranslator<?> translator = KNOWN_TRANSLATORS.get(coderType); checkArgument( translator != null, "Unknown Coder URN %s. Known URNs: %s", coderUrn, KNOWN_CODER_URNS.values()); return translator.fromComponents( coderComponents, coder.getSpec().getSpec().getPayload().toByteArray()); }
@Test public void toAndFromProto() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.Coder coderProto = CoderTranslation.toProto(coder, sdkComponents); Components encodedComponents = sdkComponents.toComponents(); Coder<?> decodedCoder = CoderTranslation.fromProto( coderProto, RehydratedComponents.forComponents(encodedComponents)); assertThat(decodedCoder, equalTo(coder)); if (KNOWN_CODERS.contains(coder)) { for (RunnerApi.Coder encodedCoder : encodedComponents.getCodersMap().values()) { assertThat( encodedCoder.getSpec().getSpec().getUrn(), not(equalTo(CoderTranslation.JAVA_SERIALIZED_CODER_URN))); } } }
/** * 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 BoundedSource<?> boundedSourceFromProto(ReadPayload payload) throws InvalidProtocolBufferException { checkArgument(payload.getIsBounded().equals(IsBounded.Enum.BOUNDED)); return (BoundedSource<?>) SerializableUtils.deserializeFromByteArray( payload.getSource().getSpec().getPayload().toByteArray(), "BoundedSource"); }
public static Coder<?> fromProto(RunnerApi.Coder protoCoder, RehydratedComponents components) throws IOException { String coderSpecUrn = protoCoder.getSpec().getSpec().getUrn(); if (coderSpecUrn.equals(JAVA_SERIALIZED_CODER_URN)) { return fromCustomCoder(protoCoder); } return fromKnownCoder(protoCoder, components); }
public static WindowedValueCoderComponents getWindowedValueCoderComponents(Coder coder) { checkArgument(WINDOWED_VALUE_CODER_URN.equals(coder.getSpec().getSpec().getUrn())); return new AutoValue_ModelCoders_WindowedValueCoderComponents( coder.getComponentCoderIds(0), coder.getComponentCoderIds(1)); }
private static Coder<?> fromCustomCoder(RunnerApi.Coder protoCoder) throws IOException { return (Coder<?>) SerializableUtils.deserializeFromByteArray( protoCoder.getSpec().getSpec().getPayload().toByteArray(), "Custom Coder Bytes"); } }
public static KvCoderComponents getKvCoderComponents(Coder coder) { checkArgument(KV_CODER_URN.equals(coder.getSpec().getSpec().getUrn())); return new AutoValue_ModelCoders_KvCoderComponents( coder.getComponentCoderIds(0), coder.getComponentCoderIds(1)); }
public static UnboundedSource<?, ?> unboundedSourceFromProto(ReadPayload payload) throws InvalidProtocolBufferException { checkArgument(payload.getIsBounded().equals(IsBounded.Enum.UNBOUNDED)); return (UnboundedSource<?, ?>) SerializableUtils.deserializeFromByteArray( payload.getSource().getSpec().getPayload().toByteArray(), "UnboundedSource"); }