Flink dynamic Key Selector

Main.kt

data class WC(var word: String, var count: Int) : Serializable {
    constructor() : this("", 0)
}

class KS : KeySelector<WC, Int> {
    private var count = 0
    override fun getKey(value: WC): Int {
        println("count=$count, value=$value")
        return if (count++ < 2) 0 else 1
    }
}

class Reducer : ReduceFunction<WC> {
    override fun reduce(value1: WC, value2: WC): WC {
        println("Got $value2")
        return value2.copy(count = value1.count + value2.count)
    }
}

fun main() {
    val values = listOf(
            WC("a", 10),
            WC("a", 10),
            WC("b", 1),
            WC("b", 2)
    )

    val env = StreamExecutionEnvironment.getExecutionEnvironment()
    env.parallelism = 1

    val ans = env.fromCollection(values)
            .keyBy(KS())
            .reduce(Reducer())

    ans.print()

    env.execute()
}

output.txt

count=0, value=WC(word=a, count=10)
count=1, value=WC(word=a, count=10)
count=2, value=WC(word=b, count=1)
count=3, value=WC(word=b, count=2)
count=0, value=WC(word=a, count=10)
WC(word=a, count=10)
count=1, value=WC(word=a, count=10)
Got WC(word=a, count=10)
WC(word=a, count=20)
count=2, value=WC(word=b, count=1)
WC(word=b, count=1)
count=3, value=WC(word=b, count=2)
Got WC(word=b, count=2)
WC(word=b, count=3)