Code review comment: i'll have a look at what the flink runner does as well.

Code before review:
public void translate(GroupByKey<K, V> transform, Context cxt) {
WindowingStrategy<?, ?> windowing = cxt.getInput().getWindowingStrategy();
TimestampCombiner tsCombiner = windowing.getTimestampCombiner();
Dataset<WindowedValue<KV<K, V>>> input = cxt.getDataset(cxt.getInput());
KvCoder<K, V> inputCoder = (KvCoder<K, V>) cxt.getInput().getCoder();
KvCoder<K, Iterable<V>> outputCoder = (KvCoder<K, Iterable<V>>) cxt.getOutput().getCoder();
Encoder<V> valueEnc = cxt.valueEncoderOf(inputCoder);
Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder);
final Dataset<WindowedValue<KV<K, Iterable<V>>>> result;
if (useCollectList && eligibleForGlobalGroupBy(windowing, false)) {
result =
input
.groupBy(col("value.key").as("key"))
.agg(collect_list(col("value.value")).as("values"), timestampAggregator(tsCombiner))
.select(
inGlobalWindow(
keyValue(col("key").as(keyEnc), col("values").as(iterableEnc(valueEnc))),
windowTimestamp(tsCombiner)));
} else if (eligibleForGlobalGroupBy(windowing, true)) {
result =
cxt.getDataset(cxt.getInput())
.groupByKey(valueKey(), keyEnc)
.mapValues(valueValue(), cxt.valueEncoderOf(inputCoder))
.mapGroups(fun2((k, it) -> KV.of(k, iterableOnce(it))), cxt.kvEncoderOf(outputCoder))
.map(fun1(WindowedValue::valueInGlobalWindow), cxt.windowedEncoder(outputCoder));
} else if (useCollectList
&& eligibleForGroupByWindow(windowing, false)
&& (windowing.getWindowFn().assignsToOneWindow() || transform.fewKeys())) {
result =
input
.select(explode(col("windows")).as("window"), col("value"), col("timestamp"))
.groupBy(col("value.key"), col("window"))
.agg(collect_list(col("value.value")).as("values"), timestampAggregator(tsCombiner))
.select(
inSingleWindow(
keyValue(col("key").as(keyEnc), col("values").as(iterableEnc(valueEnc))),
col("window").as(cxt.windowEncoder()),
windowTimestamp(tsCombiner)));
} else if (eligibleForGroupByWindow(windowing, true)
&& (windowing.getWindowFn().assignsToOneWindow() || transform.fewKeys())) {
Encoder<Tuple2<BoundedWindow, K>> windowedKeyEnc = windowedKeyEnc(keyEnc, cxt);
result =
cxt.getDataset(cxt.getInput())
.flatMap(explodeWindowedKey(valueValue()), cxt.tupleEncoder(windowedKeyEnc, valueEnc))
.groupByKey(fun1(Tuple2::_1), windowedKeyEnc)
.mapValues(fun1(Tuple2::_2), valueEnc)
.mapGroups(
fun2((wKey, it) -> windowedKV(wKey, iterableOnce(it))),
cxt.windowedEncoder(outputCoder));
} else {
result =
input
.groupByKey(valueKey(), keyEnc)
.flatMapGroups(
new GroupAlsoByWindowViaOutputBufferFn<>(
windowing,
(SerStateInternalsFactory) key -> InMemoryStateInternals.forKey(key),
SystemReduceFn.buffering(inputCoder.getValueCoder()),
cxt.getSerializableOptions()),
cxt.windowedEncoder(outputCoder));
}
cxt.putDataset(cxt.getOutput(), result);
}


Code after addressing the review:
public void translate(GroupByKey<K, V> transform, Context cxt) {
WindowingStrategy<?, ?> windowing = cxt.getInput().getWindowingStrategy();
TimestampCombiner tsCombiner = windowing.getTimestampCombiner();
Dataset<WindowedValue<KV<K, V>>> input = cxt.getDataset(cxt.getInput());
KvCoder<K, V> inputCoder = (KvCoder<K, V>) cxt.getInput().getCoder();
KvCoder<K, Iterable<V>> outputCoder = (KvCoder<K, Iterable<V>>) cxt.getOutput().getCoder();
Encoder<V> valueEnc = cxt.valueEncoderOf(inputCoder);
Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder);
final Dataset<WindowedValue<KV<K, Iterable<V>>>> result;
if (useCollectList && eligibleForGlobalGroupBy(windowing, false)) {
result =
input
.groupBy(col("value.key").as("key"))
.agg(collect_list(col("value.value")).as("values"), timestampAggregator(tsCombiner))
.select(
inGlobalWindow(
keyValue(col("key").as(keyEnc), col("values").as(iterableEnc(valueEnc))),
windowTimestamp(tsCombiner)));
} else if (eligibleForGlobalGroupBy(windowing, true)) {
result =
cxt.getDataset(cxt.getInput())
.groupByKey(valueKey(), keyEnc)
.mapValues(valueValue(), cxt.valueEncoderOf(inputCoder))
.mapGroups(fun2((k, it) -> KV.of(k, iterableOnce(it))), cxt.kvEncoderOf(outputCoder))
.map(fun1(WindowedValue::valueInGlobalWindow), cxt.windowedEncoder(outputCoder));
} else if (useCollectList
&& eligibleForGroupByWindow(windowing, false)
&& (windowing.getWindowFn().assignsToOneWindow() || transform.fewKeys())) {
result =
input
.select(explode(col("windows")).as("window"), col("value"), col("timestamp"))
.groupBy(col("value.key"), col("window"))
.agg(collect_list(col("value.value")).as("values"), timestampAggregator(tsCombiner))
.select(
inSingleWindow(
keyValue(col("key").as(keyEnc), col("values").as(iterableEnc(valueEnc))),
col("window").as(cxt.windowEncoder()),
windowTimestamp(tsCombiner)));
} else if (eligibleForGroupByWindow(windowing, true)
&& (windowing.getWindowFn().assignsToOneWindow() || transform.fewKeys())) {
Encoder<Tuple2<BoundedWindow, K>> windowedKeyEnc =
cxt.tupleEncoder(cxt.windowEncoder(), keyEnc);
result =
cxt.getDataset(cxt.getInput())
.flatMap(explodeWindowedKey(valueValue()), cxt.tupleEncoder(windowedKeyEnc, valueEnc))
.groupByKey(fun1(Tuple2::_1), windowedKeyEnc)
.mapValues(fun1(Tuple2::_2), valueEnc)
.mapGroups(
fun2((wKey, it) -> windowedKV(wKey, iterableOnce(it))),
cxt.windowedEncoder(outputCoder));
} else {
result =
input
.groupByKey(valueKey(), keyEnc)
.flatMapGroups(
new GroupAlsoByWindowViaOutputBufferFn<>(
windowing,
(SerStateInternalsFactory) key -> InMemoryStateInternals.forKey(key),
SystemReduceFn.buffering(inputCoder.getValueCoder()),
cxt.getSerializableOptions()),
cxt.windowedEncoder(outputCoder));
}
cxt.putDataset(cxt.getOutput(), result);
}