如何在Apache Spark中单次扫描计算词频与词对计数

可行性结论 #

完全可以通过单次扫描完成三类计数，你的思路是合理的，下面给出落地实现和优化建议。

核心思路 #

利用Spark的mapPartitions维护分区内的遍历状态（记录前一个token的内容和类型），在遍历每个token时一次性生成三类统计所需的键值对，最后统一聚合。这样避免了多次扫描数据集，减少I/O开销。

具体代码实现 #

import org.apache.spark.rdd.RDD

// 假设已实现的类型判断函数
def isWord(token: String): Boolean = ???
def isNumber(token: String): Boolean = ???

def computeAllCounts(tokens: RDD[String]): (Map[String, Long], Map[(String, String), Long], Map[(String, String), Long]) = {
    val allStats = tokens.mapPartitions(iter => {
        var prevToken: Option[String] = None
        val output = scala.collection.mutable.ListBuffer[((Int, Any), Long)]()

        iter.foreach(curr => {
            // 1. 单个单词计数
            if (isWord(curr)) {
                output.append(((0, curr), 1L)) // 0标记为单词计数类型
            }

            // 2. 处理相邻对
            prevToken.foreach(prev => {
                // 单词-单词对
                if (isWord(prev) && isWord(curr)) {
                    output.append(((1, (prev, curr)), 1L)) // 1标记为单词对类型
                }
                // 数字-单词对（支持两种顺序，如需无序对可排序后存储）
                if ((isNumber(prev) && isWord(curr)) || (isWord(prev) && isNumber(curr))) {
                    // 无序对处理：val sortedPair = (prev, curr).sorted
                    output.append(((2, (prev, curr)), 1L)) // 2标记为数字单词对类型
                }
            })

            prevToken = Some(curr)
        })

        output.iterator
    }).reduceByKey(_ + _).collect()

    // 拆分三类统计结果
    val wordCounts = allStats.filter(_._1._1 == 0).map { case ((_, word), cnt) => (word.asInstanceOf[String], cnt) }.toMap
    val wordPairCounts = allStats.filter(_._1._1 == 1).map { case ((_, pair), cnt) => (pair.asInstanceOf[(String, String)], cnt) }.toMap
    val numWordPairCounts = allStats.filter(_._1._1 == 2).map { case ((_, pair), cnt) => (pair.asInstanceOf[(String, String)], cnt) }.toMap

    (wordCounts, wordPairCounts, numWordPairCounts)
}

思路合理性验证 #

单次扫描的方案完全可行：

• Spark的RDD按分区并行处理，mapPartitions可以在每个分区内维护遍历状态（比如前一个token），无需全局共享状态；
• 一次性生成所有统计项的键值对，仅需一次全局聚合，相比多步扫描减少了多次数据读写和网络传输的开销，性能更优。

优化方向 #

1. 无序对处理：如果数字-单词对无需区分顺序（如"123 apple"和"apple 123"算同一对），生成键时对元素排序即可：val sortedPair = if (prev < curr) (prev, curr) else (curr, prev)。
2. 结果存储：若数据集规模较大，避免用collect()将结果拉到Driver，可直接将三类计数分别写入外部存储（如HDFS、数据库）。
3. 类型安全：可以用密封特质或枚举类替代整数标记，提升代码可读性和安全性。

内容的提问来源于stack exchange，提问作者lowlypalace