ProcessBundleDescriptors { public static ExecutableProcessBundleDescriptor fromExecutableStage( String id, ExecutableStage stage, ApiServiceDescriptor dataEndpoint, ApiServiceDescriptor stateEndpoint) throws IOException { checkState(id != null, "id must be specified."); checkState(stage != null, "stage must be specified."); checkState(dataEndpoint != null, "dataEndpoint must be specified."); checkState(stateEndpoint != null, "stateEndpoint must be specified."); return fromExecutableStageInternal(id, stage, dataEndpoint, stateEndpoint); }  static ExecutableProcessBundleDescriptor fromExecutableStage(
      String id,
      ExecutableStage stage,
      ApiServiceDescriptor dataEndpoint,
      ApiServiceDescriptor stateEndpoint); static ExecutableProcessBundleDescriptor fromExecutableStage(
      String id, ExecutableStage stage, ApiServiceDescriptor dataEndpoint); static Map<String, Map<String, SideInputSpec>> getSideInputs(ExecutableStage stage);  }
@Test public void testLengthPrefixingOfKeyCoderInStatefulExecutableStage() throws Exception { Pipeline p = Pipeline.create(); Coder<Void> keycoder = VoidCoder.of(); assertThat(ModelCoderRegistrar.isKnownCoder(keycoder), is(false)); p.apply("impulse", Impulse.create()) .apply( "create", ParDo.of( new DoFn<byte[], KV<Void, String>>() { @ProcessElement public void process(ProcessContext ctxt) {} })) .setCoder(KvCoder.of(keycoder, StringUtf8Coder.of())) .apply( "userState", ParDo.of( new DoFn<KV<Void, String>, KV<Void, String>>() { @StateId("stateId") private final StateSpec<BagState<String>> bufferState = StateSpecs.bag(StringUtf8Coder.of()); @TimerId("timerId") private final TimerSpec timerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME); @ProcessElement public void processElement( @Element KV<Void, String> element, @StateId("stateId") BagState<String> state, @TimerId("timerId") Timer timer, OutputReceiver<KV<Void, String>> r) {} @OnTimer("timerId") public void onTimer() {} })) .apply("gbk", GroupByKey.create()); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto); Optional<ExecutableStage> optionalStage = Iterables.tryFind( fused.getFusedStages(), (ExecutableStage stage) -> stage.getUserStates().stream() .anyMatch(spec -> spec.localName().equals("stateId"))); checkState(optionalStage.isPresent(), "Expected a stage with user state."); ExecutableStage stage = optionalStage.get(); PipelineNode.PCollectionNode inputPCollection = stage.getInputPCollection(); Map<String, RunnerApi.Coder> stageCoderMap = stage.getComponents().getCodersMap(); RunnerApi.Coder originalMainInputCoder = stageCoderMap.get(inputPCollection.getPCollection().getCoderId()); String originalKeyCoderId = ModelCoders.getKvCoderComponents(originalMainInputCoder).keyCoderId(); RunnerApi.Coder originalKeyCoder = stageCoderMap.get(originalKeyCoderId); assertThat(originalKeyCoder.getSpec().getUrn(), is(CoderTranslation.JAVA_SERIALIZED_CODER_URN)); BeamFnApi.ProcessBundleDescriptor pbd = ProcessBundleDescriptors.fromExecutableStage( "test_stage", stage, Endpoints.ApiServiceDescriptor.getDefaultInstance()) .getProcessBundleDescriptor(); Map<String, RunnerApi.Coder> pbsCoderMap = pbd.getCodersMap(); RunnerApi.Coder pbsMainInputCoder = pbsCoderMap.get(pbd.getPcollectionsOrThrow(inputPCollection.getId()).getCoderId()); String keyCoderId = ModelCoders.getKvCoderComponents(pbsMainInputCoder).keyCoderId(); RunnerApi.Coder keyCoder = pbsCoderMap.get(keyCoderId); ensureLengthPrefixed(keyCoder, originalKeyCoder, pbsCoderMap); TimerReference timerRef = Iterables.getOnlyElement(stage.getTimers()); String timerTransformId = timerRef.transform().getId(); RunnerApi.ParDoPayload parDoPayload = RunnerApi.ParDoPayload.parseFrom( pbd.getTransformsOrThrow(timerTransformId).getSpec().getPayload()); RunnerApi.TimerFamilySpec timerSpec = parDoPayload.getTimerFamilySpecsOrThrow(timerRef.localName()); RunnerApi.Coder timerCoder = pbsCoderMap.get(timerSpec.getTimerFamilyCoderId()); String timerKeyCoderId = timerCoder.getComponentCoderIds(0); RunnerApi.Coder timerKeyCoder = pbsCoderMap.get(timerKeyCoderId); ensureLengthPrefixed(timerKeyCoder, originalKeyCoder, pbsCoderMap); }
@Test public void testWrapKeyCoderOfStatefulExecutableStageInLengthPrefixCoder() throws Exception { Pipeline p = Pipeline.create(); Coder<Void> keycoder = VoidCoder.of(); assertThat(ModelCoderRegistrar.isKnownCoder(keycoder), is(false)); p.apply("impulse", Impulse.create()) .apply( "create", ParDo.of( new DoFn<byte[], KV<Void, String>>() { @ProcessElement public void process(ProcessContext ctxt) {} })) .setCoder(KvCoder.of(keycoder, StringUtf8Coder.of())) .apply( "userState", ParDo.of( new DoFn<KV<Void, String>, KV<Void, String>>() { @StateId("stateId") private final StateSpec<BagState<String>> bufferState = StateSpecs.bag(StringUtf8Coder.of()); @ProcessElement public void processElement( @Element KV<Void, String> element, @StateId("stateId") BagState<String> state, OutputReceiver<KV<Void, String>> r) { for (String value : state.read()) { r.output(KV.of(element.getKey(), value)); } state.add(element.getValue()); } })) .apply("gbk", GroupByKey.create()); RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p); FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto); Optional<ExecutableStage> optionalStage = Iterables.tryFind( fused.getFusedStages(), (ExecutableStage stage) -> stage.getUserStates().stream() .anyMatch(spec -> spec.localName().equals("stateId"))); checkState(optionalStage.isPresent(), "Expected a stage with user state."); ExecutableStage stage = optionalStage.get(); PipelineNode.PCollectionNode inputPCollection = stage.getInputPCollection(); Map<String, RunnerApi.Coder> stageCoderMap = stage.getComponents().getCodersMap(); RunnerApi.Coder originalCoder = stageCoderMap.get(inputPCollection.getPCollection().getCoderId()); String originalKeyCoderId = ModelCoders.getKvCoderComponents(originalCoder).keyCoderId(); assertThat( stageCoderMap.get(originalKeyCoderId).getSpec().getUrn(), is(CoderTranslation.JAVA_SERIALIZED_CODER_URN)); BeamFnApi.ProcessBundleDescriptor pbDescriptor = ProcessBundleDescriptors.fromExecutableStage( "test_stage", stage, Endpoints.ApiServiceDescriptor.getDefaultInstance()) .getProcessBundleDescriptor(); String inputPCollectionId = inputPCollection.getId(); String inputCoderId = pbDescriptor.getPcollectionsMap().get(inputPCollectionId).getCoderId(); Map<String, RunnerApi.Coder> pbCoderMap = pbDescriptor.getCodersMap(); RunnerApi.Coder coder = pbCoderMap.get(inputCoderId); String keyCoderId = ModelCoders.getKvCoderComponents(coder).keyCoderId(); RunnerApi.Coder keyCoder = pbCoderMap.get(keyCoderId); assertThat(keyCoder.getSpec().getUrn(), is(ModelCoders.LENGTH_PREFIX_CODER_URN)); String lengthPrefixWrappedCoderId = keyCoder.getComponentCoderIds(0); assertThat(lengthPrefixWrappedCoderId, is(originalKeyCoderId)); assertThat( pbCoderMap.get(lengthPrefixWrappedCoderId), is(stageCoderMap.get(originalKeyCoderId))); }