Kotlin Flow 的 GroupBy 运算符

Question

我正在尝试从 RxJava 切换到 Kotlin Flow。流量确实令人印象深刻。但是现在kotlin Flow中是否有类似RxJava的“GroupBy”的操作符？

Answer 1

从 Kotlin Coroutines 1.3 开始，标准库似乎没有提供此运算符。但是，由于Flow 的设计使得所有运算符都是扩展函数，因此提供它的标准库与您编写自己的标准库之间没有根本区别。

考虑到这一点，以下是我关于如何处理它的一些想法。

1。将每个组收集到一个列表中

如果您只需要每个键的所有项目的列表，请使用这个简单的实现，它会发出 (K, List<T>) 对：

fun <T, K> Flow<T>.groupToList(getKey: (T) -> K): Flow<Pair<K, List<T>>> = flow {
    val storage = mutableMapOf<K, MutableList<T>>()
    collect { t -> storage.getOrPut(getKey(t)) { mutableListOf() } += t }
    storage.forEach { (k, ts) -> emit(k to ts) }
}

对于这个例子：

suspend fun main() {
    val input = 1..10
    input.asFlow()
            .groupToList { it % 2 }
            .collect { println(it) }
}

打印出来

(1, [1, 3, 5, 7, 9])
(0, [2, 4, 6, 8, 10])

2.a 为每个组发出一个流

如果您需要完整的 RxJava 语义，将输入流转换为多个输出流（每个不同的键一个），事情就会变得更加复杂。

每当您在输入中看到一个新键时，您必须向下游发出一个新的内部流，然后在再次遇到相同的键时以异步方式将更多数据推入其中。

这是一个执行此操作的实现：

fun <T, K> Flow<T>.groupBy(getKey: (T) -> K): Flow<Pair<K, Flow<T>>> = flow {
    val storage = mutableMapOf<K, SendChannel<T>>()
    try {
        collect { t ->
            val key = getKey(t)
            storage.getOrPut(key) {
                Channel<T>(32).also { emit(key to it.consumeAsFlow()) }
            }.send(t)
        }
    } finally {
        storage.values.forEach { chan -> chan.close() }
    }
}

它为每个键设置一个Channel，并将通道作为流公开给下游。

2.b 并发收集和减少分组流

由于groupBy 在将流本身发送到下游后，不断将数据发送到内部流，因此您必须非常小心如何收集它们。

您必须同时收集所有内部流，并发级别没有上限。否则，排队等待稍后收集的流的通道最终会阻塞发送方，您最终会陷入死锁。

这是一个正确执行此操作的函数：

fun <T, K, R> Flow<Pair<K, Flow<T>>>.reducePerKey(
        reduce: suspend Flow<T>.() -> R
): Flow<Pair<K, R>> = flow {
    coroutineScope {
        this@reducePerKey
                .map { (key, flow) -> key to async { flow.reduce() } }
                .toList()
                .forEach { (key, deferred) -> emit(key to deferred.await()) }
    }
}

map 阶段为它接收到的每个内部流启动一个协程。协程将其简化为最终结果。

toList() 是一个终端操作，收集整个上游流，启动进程中的所有async 协程。即使我们仍在收集主要流程，协程也开始消耗内部流程。这对于防止死锁至关重要。

最后，在所有协程启动后，我们启动一个forEach 循环，等待并在最终结果可用时发出。

您可以在flatMapMerge 方面实现几乎相同的行为：

fun <T, K, R> Flow<Pair<K, Flow<T>>>.reducePerKey(
        reduce: suspend Flow<T>.() -> R
): Flow<Pair<K, R>> = flatMapMerge(Int.MAX_VALUE) { (key, flow) ->
    flow { emit(key to flow.reduce()) }
}

区别在于顺序：第一个实现尊重输入中键的出现顺序，而这个不尊重。两者的表现相似。

3。示例

此示例对 4000 万个整数进行分组和求和：

suspend fun main() {
    val input = 1..40_000_000
    input.asFlow()
            .groupBy { it % 100 }
            .reducePerKey { sum { it.toLong() } }
            .collect { println(it) }
}

suspend fun <T> Flow<T>.sum(toLong: suspend (T) -> Long): Long {
    var sum = 0L
    collect { sum += toLong(it) }
    return sum
}

我可以使用-Xmx64m 成功运行它。在我的 4 核笔记本电脑上，我每秒处理大约 400 万个项目。

用新的解决方案重新定义第一个解决方案很简单，如下所示：

fun <T, K> Flow<T>.groupToList(getKey: (T) -> K): Flow<Pair<K, List<T>>> =
        groupBy(getKey).reducePerKey { toList() }

Answer 2

还没有，但你可以看看这个库https://github.com/akarnokd/kotlin-flow-extensions。