PartialGroupByKeyParDoFns { public static <K, InputT, AccumT> ParDoFn create( PipelineOptions options, KvCoder<K, ?> inputElementCoder, @Nullable CloudObject cloudUserFn, @Nullable List<SideInputInfo> sideInputInfos, List<Receiver> receivers, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception { AppliedCombineFn<K, InputT, AccumT, ?> combineFn; SideInputReader sideInputReader; StepContext stepContext; if (cloudUserFn == null) { combineFn = null; sideInputReader = NullSideInputReader.empty(); stepContext = null; } else { Object deserializedFn = SerializableUtils.deserializeFromByteArray( getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN), "serialized combine fn"); @SuppressWarnings("unchecked") AppliedCombineFn<K, InputT, AccumT, ?> combineFnUnchecked = ((AppliedCombineFn<K, InputT, AccumT, ?>) deserializedFn); combineFn = combineFnUnchecked; sideInputReader = executionContext.getSideInputReader( sideInputInfos, combineFn.getSideInputViews(), operationContext); stepContext = executionContext.getStepContext(operationContext); } return create( options, inputElementCoder, combineFn, sideInputReader, receivers.get(0), stepContext); }  static ParDoFn create(
      PipelineOptions options,
      KvCoder<K, ?> inputElementCoder,
      @Nullable CloudObject cloudUserFn,
      @Nullable List<SideInputInfo> sideInputInfos,
      List<Receiver> receivers,
      DataflowExecutionContext<?> executionContext,
      DataflowOperationContext operationContext);  }
@Test public void testPartialGroupByKey() throws Exception { Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder( WindowedValue.getValueOnlyCoder( KvCoder.of(keyCoder, IterableCoder.of(valueCoder)))), counterSet, NameContextsForTests.nameContextForTest()); ParDoFn pgbkParDoFn = new SimplePartialGroupByKeyParDoFn( GroupingTables.buffering( new WindowingCoderGroupingKeyCreator(keyCoder), PairInfo.create(), new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)), new CoderSizeEstimator(valueCoder)), receiver); pgbkParDoFn.startBundle(receiver); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 4))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 5))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 6))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("joe", 7))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 8))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 9))); pgbkParDoFn.finishBundle(); assertThat( receiver.outputElems, IsIterableContainingInAnyOrder.<Object>containsInAnyOrder( WindowedValue.valueInGlobalWindow(KV.of("hi", Arrays.asList(4, 6, 9))), WindowedValue.valueInGlobalWindow(KV.of("there", Arrays.asList(5, 8))), WindowedValue.valueInGlobalWindow(KV.of("joe", Arrays.asList(7))))); CounterUpdateExtractor<?> updateExtractor = Mockito.mock(CounterUpdateExtractor.class); counterSet.extractUpdates(false, updateExtractor); verify(updateExtractor).longSum(getObjectCounterName("test_receiver_out"), false, 3L); verify(updateExtractor) .longMean( getMeanByteCounterName("test_receiver_out"), false, LongCounterMean.ZERO.addValue(49L, 3)); verifyNoMoreInteractions(updateExtractor); }
@Test public void testPartialGroupByKeyWithCombiner() throws Exception { Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder( WindowedValue.getValueOnlyCoder(KvCoder.of(keyCoder, valueCoder))), counterSet, NameContextsForTests.nameContextForTest()); Combiner<WindowedValue<String>, Integer, Integer, Integer> combineFn = new TestCombiner(); ParDoFn pgbkParDoFn = new SimplePartialGroupByKeyParDoFn( GroupingTables.combining( new WindowingCoderGroupingKeyCreator(keyCoder), PairInfo.create(), combineFn, new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)), new CoderSizeEstimator(valueCoder)), receiver); pgbkParDoFn.startBundle(receiver); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 4))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 5))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 6))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("joe", 7))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 8))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 9))); pgbkParDoFn.finishBundle(); assertThat( receiver.outputElems, IsIterableContainingInAnyOrder.<Object>containsInAnyOrder( WindowedValue.valueInGlobalWindow(KV.of("hi", 19)), WindowedValue.valueInGlobalWindow(KV.of("there", 13)), WindowedValue.valueInGlobalWindow(KV.of("joe", 7)))); CounterUpdateExtractor<?> updateExtractor = Mockito.mock(CounterUpdateExtractor.class); counterSet.extractUpdates(false, updateExtractor); verify(updateExtractor).longSum(getObjectCounterName("test_receiver_out"), false, 3L); verify(updateExtractor) .longMean( getMeanByteCounterName("test_receiver_out"), false, LongCounterMean.ZERO.addValue(25L, 3)); verifyNoMoreInteractions(updateExtractor); }
@Test public void testPartialGroupByKeyWithCombinerAndSideInputs() throws Exception { Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder( WindowedValue.getValueOnlyCoder(KvCoder.of(keyCoder, valueCoder))), counterSet, NameContextsForTests.nameContextForTest()); Combiner<WindowedValue<String>, Integer, Integer, Integer> combineFn = new TestCombiner(); ParDoFn pgbkParDoFn = new StreamingSideInputPGBKParDoFn( GroupingTables.combining( new WindowingCoderGroupingKeyCreator(keyCoder), PairInfo.create(), combineFn, new CoderSizeEstimator(WindowedValue.getValueOnlyCoder(keyCoder)), new CoderSizeEstimator(valueCoder)), receiver, mockSideInputFetcher); Set<BoundedWindow> readyWindows = ImmutableSet.<BoundedWindow>of(GlobalWindow.INSTANCE); when(mockSideInputFetcher.getReadyWindows()).thenReturn(readyWindows); when(mockSideInputFetcher.prefetchElements(readyWindows)) .thenReturn(ImmutableList.of(elemsBag)); when(elemsBag.read()) .thenReturn( ImmutableList.of( WindowedValue.valueInGlobalWindow(KV.of("hi", 4)), WindowedValue.valueInGlobalWindow(KV.of("there", 5)))); when(mockSideInputFetcher.storeIfBlocked(Matchers.<WindowedValue<KV<String, Integer>>>any())) .thenReturn(false, false, false, true); pgbkParDoFn.startBundle(receiver); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 6))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("joe", 7))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("there", 8))); pgbkParDoFn.processElement(WindowedValue.valueInGlobalWindow(KV.of("hi", 9))); pgbkParDoFn.finishBundle(); assertThat( receiver.outputElems, IsIterableContainingInAnyOrder.<Object>containsInAnyOrder( WindowedValue.valueInGlobalWindow(KV.of("hi", 10)), WindowedValue.valueInGlobalWindow(KV.of("there", 13)), WindowedValue.valueInGlobalWindow(KV.of("joe", 7)))); CounterUpdateExtractor<?> updateExtractor = Mockito.mock(CounterUpdateExtractor.class); counterSet.extractUpdates(false, updateExtractor); verify(updateExtractor).longSum(getObjectCounterName("test_receiver_out"), false, 3L); verify(updateExtractor) .longMean( getMeanByteCounterName("test_receiver_out"), false, LongCounterMean.ZERO.addValue(25L, 3)); verifyNoMoreInteractions(updateExtractor); }
@Test public void testCreateWithCombinerAndBatchSideInputs() throws Exception { PipelineOptions options = PipelineOptionsFactory.create(); Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)), counterSet, NameContextsForTests.nameContextForTest()); StepContext stepContext = BatchModeExecutionContext.forTesting(options, "testStage") .getStepContext(TestOperationContext.create(counterSet)); when(mockSideInputReader.isEmpty()).thenReturn(false); ParDoFn pgbk = PartialGroupByKeyParDoFns.create( options, kvCoder, AppliedCombineFn.withInputCoder( Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder, ImmutableList.<PCollectionView<?>>of(), WindowingStrategy.globalDefault()), mockSideInputReader, receiver, stepContext); assertTrue(pgbk instanceof BatchSideInputPGBKParDoFn); }
@Test public void testCreateWithCombinerAndStreaming() throws Exception { StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class); options.setStreaming(true); Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)), counterSet, NameContextsForTests.nameContextForTest()); ParDoFn pgbk = PartialGroupByKeyParDoFns.create( options, kvCoder, AppliedCombineFn.withInputCoder( Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder), NullSideInputReader.empty(), receiver, null); assertTrue(pgbk instanceof SimplePartialGroupByKeyParDoFn); }
@Test public void testCreateWithCombinerAndStreamingSideInputs() throws Exception { StreamingOptions options = PipelineOptionsFactory.as(StreamingOptions.class); options.setStreaming(true); Coder keyCoder = StringUtf8Coder.of(); Coder valueCoder = BigEndianIntegerCoder.of(); KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, valueCoder); TestOutputReceiver receiver = new TestOutputReceiver( new ElementByteSizeObservableCoder(WindowedValue.getValueOnlyCoder(kvCoder)), counterSet, NameContextsForTests.nameContextForTest()); when(mockSideInputReader.isEmpty()).thenReturn(false); when(mockStreamingStepContext.stateInternals()).thenReturn((StateInternals) mockStateInternals); when(mockStateInternals.state(Matchers.<StateNamespace>any(), Matchers.<StateTag>any())) .thenReturn(mockState); when(mockState.read()).thenReturn(Maps.newHashMap()); ParDoFn pgbk = PartialGroupByKeyParDoFns.create( options, kvCoder, AppliedCombineFn.withInputCoder( Sum.ofIntegers(), CoderRegistry.createDefault(), kvCoder, ImmutableList.<PCollectionView<?>>of(), WindowingStrategy.globalDefault()), mockSideInputReader, receiver, mockStreamingStepContext); assertTrue(pgbk instanceof StreamingSideInputPGBKParDoFn); }