BufferedOutputManager( Output<StreamRecord<WindowedValue<OutputT>>> output, TupleTag<OutputT> mainTag, Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags, final Map<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders, Map<TupleTag<?>, Integer> tagsToIds, StateInternals stateInternals) { this.output = output; this.mainTag = mainTag; this.tagsToOutputTags = tagsToOutputTags; this.tagsToIds = tagsToIds; this.idsToTags = new HashMap<>(); for (Map.Entry<TupleTag<?>, Integer> entry : tagsToIds.entrySet()) { idsToTags.put(entry.getValue(), entry.getKey()); } ImmutableMap.Builder<Integer, Coder<WindowedValue<?>>> idsToCodersBuilder = ImmutableMap.builder(); for (Map.Entry<TupleTag<?>, Integer> entry : tagsToIds.entrySet()) { idsToCodersBuilder.put(entry.getValue(), tagsToCoders.get(entry.getKey())); } StateTag<BagState<KV<Integer, WindowedValue<?>>>> bufferTag = StateTags.bag("bundle-buffer-tag", new TaggedKvCoder(idsToCodersBuilder.build())); bufferState = stateInternals.state(StateNamespaces.global(), bufferTag); }
@Before public void setUpMocks() throws Exception { MockitoAnnotations.initMocks(this); when(runtimeContext.getDistributedCache()).thenReturn(distributedCache); when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory); RemoteBundle remoteBundle = Mockito.mock(RemoteBundle.class); when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(remoteBundle); ImmutableMap input = ImmutableMap.builder().put("input", Mockito.mock(FnDataReceiver.class)).build(); when(remoteBundle.getInputReceivers()).thenReturn(input); when(processBundleDescriptor.getTimerSpecs()).thenReturn(Collections.emptyMap()); }
@Test public void testStageBundleClosed() throws Exception { TupleTag<Integer> mainOutput = new TupleTag<>("main-output"); DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory = new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of()); ExecutableStageDoFnOperator<Integer, Integer> operator = getOperator(mainOutput, Collections.emptyList(), outputManagerFactory); OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness = new OneInputStreamOperatorTestHarness<>(operator); RemoteBundle bundle = Mockito.mock(RemoteBundle.class); when(bundle.getInputReceivers()) .thenReturn( ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder() .put("input", Mockito.mock(FnDataReceiver.class)) .build()); when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle); testHarness.open(); testHarness.close(); verify(stageBundleFactory).getProcessBundleDescriptor(); verify(stageBundleFactory).close(); verify(stageContext).close(); // DoFnOperator generates a final watermark, which triggers a new bundle.. verify(stageBundleFactory).getBundle(any(), any(), any()); verifyNoMoreInteractions(stageBundleFactory); }
.put(additionalOutput1, new OutputTag<String>(additionalOutput1.getId()) {}) .put(additionalOutput2, new OutputTag<String>(additionalOutput2.getId()) {}) .build(); ImmutableMap<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders = ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder() .put(additionalOutput1, coder) .put(additionalOutput2, coder) .build(); ImmutableMap<TupleTag<?>, Integer> tagsToIds = ImmutableMap.<TupleTag<?>, Integer>builder() .put(additionalOutput1, 1) .put(additionalOutput2, 2) .build();
.put(additionalOutput1, new OutputTag<String>(additionalOutput1.getId()) {}) .put(additionalOutput2, new OutputTag<String>(additionalOutput2.getId()) {}) .build(); ImmutableMap<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders = ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder() .put(additionalOutput1, coder) .put(additionalOutput2, coder) .build(); ImmutableMap<TupleTag<?>, Integer> tagsToIds = ImmutableMap.<TupleTag<?>, Integer>builder() .put(additionalOutput1, 1) .put(additionalOutput2, 2) .build();
ImmutableMap.<Integer, PCollectionView<?>>builder().put(1, view1).put(2, view2).build();
additionalOutput, new OutputTag<>(additionalOutput.getId(), TypeInformation.of(Integer.class))) .build(); ImmutableMap<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders = ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder() .put(mainOutput, (Coder) coder) .put(additionalOutput, coder) .build(); ImmutableMap<TupleTag<?>, Integer> tagsToIds = ImmutableMap.<TupleTag<?>, Integer>builder() .put(mainOutput, 0) .put(additionalOutput, 1) .build();
.put(1, view1) .put(2, view2) .build();
.put(1, view1) .put(2, view2) .build();
.put(1, view1) .put(2, view2) .build();
.put(1, view1) .put(2, view2) .build();
/** * Creates a new state handler for the given stage. Note that this requires a traversal of the * stage itself, so this should only be called once per stage rather than once per bundle. */ public static FlinkStreamingSideInputHandlerFactory forStage( ExecutableStage stage, Map<SideInputId, PCollectionView<?>> viewMapping, org.apache.beam.runners.core.SideInputHandler runnerHandler) { ImmutableMap.Builder<SideInputId, PCollectionView<?>> sideInputBuilder = ImmutableMap.builder(); for (SideInputReference sideInput : stage.getSideInputs()) { SideInputId sideInputId = SideInputId.newBuilder() .setTransformId(sideInput.transform().getId()) .setLocalName(sideInput.localName()) .build(); sideInputBuilder.put( sideInputId, checkNotNull( viewMapping.get(sideInputId), "No side input for %s/%s", sideInputId.getTransformId(), sideInputId.getLocalName())); } FlinkStreamingSideInputHandlerFactory factory = new FlinkStreamingSideInputHandlerFactory(sideInputBuilder.build(), runnerHandler); return factory; }
/** Creates a batch translator. */ public static FlinkBatchPortablePipelineTranslator createTranslator() { ImmutableMap.Builder<String, PTransformTranslator> translatorMap = ImmutableMap.builder(); translatorMap.put( PTransformTranslation.FLATTEN_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateFlatten); translatorMap.put( PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateGroupByKey); translatorMap.put( PTransformTranslation.IMPULSE_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateImpulse); translatorMap.put( PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, FlinkBatchPortablePipelineTranslator::translateAssignWindows); translatorMap.put( ExecutableStage.URN, FlinkBatchPortablePipelineTranslator::translateExecutableStage); translatorMap.put( PTransformTranslation.RESHUFFLE_URN, FlinkBatchPortablePipelineTranslator::translateReshuffle); return new FlinkBatchPortablePipelineTranslator(translatorMap.build()); }
@Override public Map< ? extends Class<? extends PTransform>, ? extends PTransformTranslation.TransformPayloadTranslator> getTransformPayloadTranslators() { return ImmutableMap .<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder() .put( CreateStreamingFlinkView.CreateFlinkPCollectionView.class, new CreateStreamingFlinkViewPayloadTranslator()) .put( SplittableParDoViaKeyedWorkItems.ProcessElements.class, new SplittableParDoProcessElementsTranslator()) .put( SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class, new SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator()) .build(); } }
/** * Creates a new state handler for the given stage. Note that this requires a traversal of the * stage itself, so this should only be called once per stage rather than once per bundle. */ static FlinkBatchSideInputHandlerFactory forStage( ExecutableStage stage, RuntimeContext runtimeContext) { ImmutableMap.Builder<SideInputId, PCollectionNode> sideInputBuilder = ImmutableMap.builder(); for (SideInputReference sideInput : stage.getSideInputs()) { sideInputBuilder.put( SideInputId.newBuilder() .setTransformId(sideInput.transform().getId()) .setLocalName(sideInput.localName()) .build(), sideInput.collection()); } return new FlinkBatchSideInputHandlerFactory(sideInputBuilder.build(), runtimeContext); }
FlinkStreamingPortablePipelineTranslator() { ImmutableMap.Builder<String, PTransformTranslator<StreamingTranslationContext>> translatorMap = ImmutableMap.builder(); translatorMap.put(PTransformTranslation.FLATTEN_TRANSFORM_URN, this::translateFlatten); translatorMap.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, this::translateGroupByKey); translatorMap.put(PTransformTranslation.IMPULSE_TRANSFORM_URN, this::translateImpulse); translatorMap.put(STREAMING_IMPULSE_TRANSFORM_URN, this::translateStreamingImpulse); translatorMap.put( PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, this::translateAssignWindows); translatorMap.put(ExecutableStage.URN, this::translateExecutableStage); translatorMap.put(PTransformTranslation.RESHUFFLE_URN, this::translateReshuffle); this.urnToTransformTranslator = translatorMap.build(); }
@Test public void testOutputMapCreation() { List<String> outputs = Arrays.asList("output1", "output2", "output3"); BiMap<String, Integer> outputMap = FlinkPipelineTranslatorUtils.createOutputMap(outputs); Map<Object, Object> expected = ImmutableMap.builder().put("output1", 0).put("output2", 1).put("output3", 2).build(); assertThat(outputMap, is(expected)); } }
@SuppressWarnings("unchecked") public MultiOutputOutputManagerFactory( TupleTag<OutputT> mainTag, Coder<WindowedValue<OutputT>> mainCoder) { this( mainTag, new HashMap<>(), ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder() .put(mainTag, (Coder) mainCoder) .build(), ImmutableMap.<TupleTag<?>, Integer>builder().put(mainTag, 0).build()); }