Examples with CountSum org.apache.beam.sdk.transforms.Mean.CountSum used on opensource projects

Example 1 with CountSum

use of org.apache.beam.sdk.transforms.Mean.CountSum in project beam by apache.

the class ParDoTest method testCombiningState.

@Test
@Category({ ValidatesRunner.class, UsesStatefulParDo.class })
public void testCombiningState() {
    final String stateId = "foo";
    DoFn<KV<String, Double>, String> fn = new DoFn<KV<String, Double>, String>() {

        private static final double EPSILON = 0.0001;

        @StateId(stateId)
        private final StateSpec<CombiningState<Double, CountSum<Double>, Double>> combiningState = StateSpecs.combining(new Mean.CountSumCoder<Double>(), Mean.<Double>of());

        @ProcessElement
        public void processElement(ProcessContext c, @StateId(stateId) CombiningState<Double, CountSum<Double>, Double> state) {
            state.add(c.element().getValue());
            Double currentValue = state.read();
            if (Math.abs(currentValue - 0.5) < EPSILON) {
                c.output("right on");
            }
        }
    };
    PCollection<String> output = pipeline.apply(Create.of(KV.of("hello", 0.3), KV.of("hello", 0.6), KV.of("hello", 0.6))).apply(ParDo.of(fn));
    // There should only be one moment at which the average is exactly 0.5
    PAssert.that(output).containsInAnyOrder("right on");
    pipeline.run();
}