@Override
public AccumT extractOutput(K key, List<AccumT> accumulator) {
if (accumulator.size() == 0) {
return keyedCombineFn.createAccumulator(key);
} else {
return keyedCombineFn.mergeAccumulators(key, accumulator);
}
}
@Override
public void clear() {
// Even though we're clearing we can't remove this from the in-memory state map, since
// other users may already have a handle on this CombiningValue.
accum = combineFn.createAccumulator(key);
isCleared = true;
}
private void closeWindow(
K key, W w, Map<W, AccumT> accumulators, Map<W, Instant> minTimestamps, ProcessContext c) {
AccumT accum = accumulators.remove(w);
Instant timestamp = minTimestamps.remove(w);
checkState(accum != null && timestamp != null);
c.windowingInternals()
.outputWindowedValue(
KV.of(key, combineFn.extractOutput(key, accum)),
timestamp,
Arrays.asList(w),
PaneInfo.ON_TIME_AND_ONLY_FIRING);
}
@Override
public void populateDisplayData(DisplayData.Builder builder) {
super.populateDisplayData(builder);
CombineFns.populateDisplayData(builder, keyedCombineFns);
}
@Override
public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {
return combineFn.mergeAccumulators(key, accumulators);
}
@Override
public void addAccum(AccumT accum) {
isCleared = false;
this.accum = combineFn.mergeAccumulators(key, Arrays.asList(this.accum, accum));
}
@Override
public void add(InputT input) {
isCleared = false;
accum = combineFn.addInput(key, accum, input);
}
@Override
public OutputT read() {
return combineFn.extractOutput(key, accum);
}
InMemoryCombiningValue(K key, KeyedCombineFn<? super K, InputT, AccumT, OutputT> combineFn) {
this.key = key;
this.combineFn = combineFn;
accum = combineFn.createAccumulator(key);
}
private List<AccumT> mergeToSingleton(K key, Iterable<AccumT> accumulators) {
List<AccumT> singleton = new ArrayList<>();
singleton.add(keyedCombineFn.mergeAccumulators(key, accumulators));
return singleton;
}
@Override
public void processElement(ProcessContext c) throws Exception {
final K key = c.element().getKey();
Iterator<WindowedValue<InputT>> iterator = c.element().getValue().iterator();
final PriorityQueue<W> liveWindows =
new PriorityQueue<>(
11,
new Comparator<BoundedWindow>() {
@Override
public int compare(BoundedWindow w1, BoundedWindow w2) {
return Long.signum(w1.maxTimestamp().getMillis() - w2.maxTimestamp().getMillis());
}
});
final Map<W, AccumT> accumulators = Maps.newHashMap();
final Map<W, Instant> minTimestamps = Maps.newHashMap();
WindowFn<Object, W>.MergeContext mergeContext =
new CombiningMergeContext() {
@Override
public Collection<W> windows() {
return liveWindows;
}
@Override
public void merge(Collection<W> toBeMerged, W mergeResult) throws Exception {
List<AccumT> accumsToBeMerged = new ArrayList<>(toBeMerged.size());
Instant minTimestamp = null;
for (W window : toBeMerged) {
accumsToBeMerged.add(accumulators.remove(window));
Instant timestampToBeMerged = minTimestamps.remove(window);
if (minTimestamp == null
|| (timestampToBeMerged != null && timestampToBeMerged.isBefore(minTimestamp))) {
minTimestamp = timestampToBeMerged;
}
}
liveWindows.removeAll(toBeMerged);
minTimestamps.put(mergeResult, minTimestamp);
liveWindows.add(mergeResult);
accumulators.put(mergeResult, combineFn.mergeAccumulators(key, accumsToBeMerged));
}
};
while (iterator.hasNext()) {
WindowedValue<InputT> e = iterator.next();
@SuppressWarnings("unchecked")
Collection<W> windows = (Collection<W>) e.getWindows();
for (W w : windows) {
Instant timestamp = minTimestamps.get(w);
if (timestamp == null || timestamp.compareTo(e.getTimestamp()) > 0) {
minTimestamps.put(w, e.getTimestamp());
} else {
minTimestamps.put(w, timestamp);
}
AccumT accum = accumulators.get(w);
checkState((timestamp == null && accum == null) || (timestamp != null && accum != null));
if (accum == null) {
accum = combineFn.createAccumulator(key);
liveWindows.add(w);
}
accum = combineFn.addInput(key, accum, e.getValue());
accumulators.put(w, accum);
}
windowFn.mergeWindows(mergeContext);
while (!liveWindows.isEmpty()
&& liveWindows.peek().maxTimestamp().isBefore(e.getTimestamp())) {
closeWindow(key, liveWindows.poll(), accumulators, minTimestamps, c);
}
}
// To have gotten here, we've either not had any elements added, or we've only run merge
// and then closed windows. We don't need to retry merging.
while (!liveWindows.isEmpty()) {
closeWindow(key, liveWindows.poll(), accumulators, minTimestamps, c);
}
}
