@Test
public void testCombineValuesFnMerge() throws Exception {
TestReceiver receiver = new TestReceiver();
MeanInts mean = new MeanInts();
Combine.KeyedCombineFn<String, Integer, MeanInts.CountSum, String> combiner = mean.asKeyedFn();
ParDoFn combineParDoFn = createCombineValuesFn(CombineValuesFn.CombinePhase.MERGE, combiner);
combineParDoFn.startBundle(receiver);
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(
KV.of(
"a",
Arrays.asList(
mean.new CountSum(3, 6), mean.new CountSum(2, 9), mean.new CountSum(1, 12)))));
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(
KV.of("b", Arrays.asList(mean.new CountSum(2, 20), mean.new CountSum(1, 1)))));
combineParDoFn.finishBundle();
Object[] expectedReceivedElems = {
WindowedValue.valueInGlobalWindow(KV.of("a", mean.new CountSum(6, 27))),
WindowedValue.valueInGlobalWindow(KV.of("b", mean.new CountSum(3, 21))),
};
assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}
@Test
public void testCombineValuesFnAll() throws Exception {
TestReceiver receiver = new TestReceiver();
Combine.KeyedCombineFn<String, Integer, MeanInts.CountSum, String> combiner =
(new MeanInts()).asKeyedFn();
ParDoFn combineParDoFn = createCombineValuesFn(CombineValuesFn.CombinePhase.ALL, combiner);
combineParDoFn.startBundle(receiver);
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(KV.of("a", Arrays.asList(5, 6, 7))));
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(KV.of("b", Arrays.asList(1, 3, 7))));
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(KV.of("c", Arrays.asList(3, 6, 8, 9))));
combineParDoFn.finishBundle();
Object[] expectedReceivedElems = {
WindowedValue.valueInGlobalWindow(KV.of("a", String.format("%.1f", 6.0))),
WindowedValue.valueInGlobalWindow(KV.of("b", String.format("%.1f", 3.7))),
WindowedValue.valueInGlobalWindow(KV.of("c", String.format("%.1f", 6.5))),
};
assertArrayEquals(expectedReceivedElems, receiver.receivedElems.toArray());
}
@Test
public void testCombineValuesFnExtract() throws Exception {
TestReceiver receiver = new TestReceiver();
MeanInts mean = new MeanInts();
Combine.KeyedCombineFn<String, Integer, MeanInts.CountSum, String> combiner = mean.asKeyedFn();
ParDoFn combineParDoFn = createCombineValuesFn(CombineValuesFn.CombinePhase.EXTRACT, combiner);
combineParDoFn.startBundle(receiver);
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(KV.of("a", mean.new CountSum(6, 27))));
combineParDoFn.processElement(
WindowedValue.valueInGlobalWindow(KV.of("b", mean.new CountSum(3, 21))));
combineParDoFn.finishBundle();
assertArrayEquals(
new Object[] {
WindowedValue.valueInGlobalWindow(KV.of("a", String.format("%.1f", 4.5))),
WindowedValue.valueInGlobalWindow(KV.of("b", String.format("%.1f", 7.0)))
},
receiver.receivedElems.toArray());
}
