Examples with CombineFnWithContext org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext used on opensource projects

Example 1 with CombineFnWithContext

use of org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext in project beam by apache.

the class CombineFnUtilTest method testToFnWithContext.

@Test
public void testToFnWithContext() throws Exception {
    CombineFnWithContext<Integer, int[], Integer> fnWithContext = CombineFnUtil.toFnWithContext(Sum.ofIntegers());
    List<Integer> inputs = ImmutableList.of(1, 2, 3, 4);
    Context nullContext = CombineContextFactory.nullContext();
    int[] accum = fnWithContext.createAccumulator(nullContext);
    for (Integer i : inputs) {
        accum = fnWithContext.addInput(accum, i, nullContext);
    }
    assertEquals(10, fnWithContext.extractOutput(accum, nullContext).intValue());
}

Example 2 with CombineFnWithContext

use of org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext in project beam by apache.

the class BatchGroupAlsoByWindowAndCombineFn method processElement.

@Override
public void processElement(KV<K, Iterable<WindowedValue<InputT>>> element, PipelineOptions options, StepContext stepContext, SideInputReader sideInputReader, OutputWindowedValue<KV<K, OutputT>> output) throws Exception {
    final PerKeyCombineFnRunner<K, InputT, AccumT, OutputT> perKeyCombineFnRunner;
    if (perKeyCombineFn instanceof CombineFn) {
        perKeyCombineFnRunner = new KeyedCombineFnRunner((CombineFn<InputT, AccumT, OutputT>) perKeyCombineFn);
    } else {
        perKeyCombineFnRunner = new KeyedCombineFnWithContextRunner(options, (CombineFnWithContext<InputT, AccumT, OutputT>) perKeyCombineFn, sideInputReader);
    }
    final K key = element.getKey();
    Iterator<WindowedValue<InputT>> iterator = element.getValue().iterator();
    final PriorityQueue<W> liveWindows = new PriorityQueue<>(11, (w1, w2) -> Long.signum(w1.maxTimestamp().getMillis() - w2.maxTimestamp().getMillis()));
    final Map<W, AccumT> accumulators = Maps.newHashMap();
    final Map<W, Instant> accumulatorOutputTimestamps = Maps.newHashMap();
    WindowFn<Object, W>.MergeContext mergeContext = windowingStrategy.getWindowFn().new MergeContext() {

        @Override
        public Collection<W> windows() {
            return liveWindows;
        }

        @Override
        public void merge(Collection<W> toBeMerged, W mergeResult) throws Exception {
            List<AccumT> accumsToBeMerged = Lists.newArrayListWithCapacity(toBeMerged.size());
            List<Instant> timestampsToBeMerged = Lists.newArrayListWithCapacity(toBeMerged.size());
            for (W window : toBeMerged) {
                accumsToBeMerged.add(accumulators.remove(window));
                timestampsToBeMerged.add(accumulatorOutputTimestamps.remove(window));
            }
            liveWindows.removeAll(toBeMerged);
            Instant mergedOutputTimestamp = windowingStrategy.getTimestampCombiner().merge(mergeResult, timestampsToBeMerged);
            accumulatorOutputTimestamps.put(mergeResult, mergedOutputTimestamp);
            liveWindows.add(mergeResult);
            AccumT accum = perKeyCombineFnRunner.mergeAccumulators(accumsToBeMerged, mergeResult);
            accumulators.put(mergeResult, accum);
        }
    };
    while (iterator.hasNext()) {
        WindowedValue<InputT> e = iterator.next();
        @SuppressWarnings("unchecked") Collection<W> windows = (Collection<W>) e.getWindows();
        for (W window : windows) {
            Instant outputTime = windowingStrategy.getTimestampCombiner().assign(window, e.getTimestamp());
            Instant accumulatorOutputTime = accumulatorOutputTimestamps.get(window);
            if (accumulatorOutputTime == null) {
                accumulatorOutputTimestamps.put(window, outputTime);
            } else {
                accumulatorOutputTimestamps.put(window, windowingStrategy.getTimestampCombiner().combine(outputTime, accumulatorOutputTime));
            }
            AccumT accum = accumulators.get(window);
            checkState((accumulatorOutputTime == null && accum == null) || (accumulatorOutputTime != null && accum != null), "accumulator and accumulatorOutputTime should both be null or both be non-null");
            if (accum == null) {
                accum = perKeyCombineFnRunner.createAccumulator(window);
                liveWindows.add(window);
            }
            accum = perKeyCombineFnRunner.addInput(accum, e.getValue(), window);
            accumulators.put(window, accum);
        }
        windowingStrategy.getWindowFn().mergeWindows(mergeContext);
        while (!liveWindows.isEmpty() && liveWindows.peek().maxTimestamp().isBefore(e.getTimestamp())) {
            closeWindow(perKeyCombineFnRunner, key, liveWindows.poll(), accumulators, accumulatorOutputTimestamps, output);
        }
    }
    // and then closed windows. We don't need to retry merging.
    while (!liveWindows.isEmpty()) {
        closeWindow(perKeyCombineFnRunner, key, liveWindows.poll(), accumulators, accumulatorOutputTimestamps, output);
    }
}