时间戳

首页 > 解决方案 > Kotlin Flow 的 GroupBy 运算符

kotlin - Kotlin Flow 的 GroupBy 运算符

问题描述

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

标签： kotlinkotlin-flow

解决方案

从 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() }

推荐阅读