Spark倒排索引构建代码优化问询：能否合并groupBy提升效率与简洁性

原代码存在两次shuffle开销，完全可以优化得更简洁高效，核心优化思路是减少shuffle数据量、避免不必要的多次分组。

最优性能优化方案 #

该方案通过先在单文档内完成词频统计，大幅降低shuffle阶段传输的数据量，仅触发一次分组shuffle，性能提升最明显，输出格式完全符合要求：

import org.apache.spark.sql.functions._
import org.apache.spark.storage.StorageLevel

pipelineClean(corpus)
  // 单文档内直接统计每个词的词频，避免重复数据进入shuffle
  .select(
    $"id".as("documentId"),
    transform(
      array_distinct($"tokens"),
      term => struct(
        term.as("term"),
        size(filter($"tokens", _ === term)).as("termFreq")
      )
    ).as("docTermFreq")
  )
  // 展开后每个词+文档仅对应一行，无重复数据
  .select(explode($"docTermFreq").as("termInfo"), $"documentId")
  .where($"termInfo.term" =!= "") // 过滤前置，减少无效计算
  // 单次分组收集倒排表，生成要求的Array[Long, Long]格式
  .groupBy($"termInfo.term".as("term"))
  .agg(collect_list(array($"documentId", $"termInfo.termFreq".cast("long"))).as("postingList"))
  .orderBy("term")
  .persist(StorageLevel.MEMORY_ONLY_SER)

简化可读性优化方案 #

如果不需要极致性能，仅想简化原有两次分组的写法、保持逻辑可读性，可以用下面的版本，相比原代码去掉了冗余的中间步骤：

import org.apache.spark.sql.functions._
import org.apache.spark.storage.StorageLevel

pipelineClean(corpus)
  .select($"id" as "documentId", explode($"tokens") as "term")
  .where($"term" =!= "") // 过滤前置，减少后续计算量
  .groupBy("term", "documentId").count
  .groupBy("term").agg(collect_list(array($"documentId", $"count".cast("long"))).as("postingList"))
  .orderBy("term")
  .persist(StorageLevel.MEMORY_ONLY_SER)

优化说明 #

• 原代码两次groupBy会触发两次shuffle，最优方案通过单文档预统计词频，将shuffle数据量降低数倍（同一文档内重复的词不会产生多条重复行参与shuffle），仅需一次分组即可完成聚合。
• 空term过滤前置，避免无效数据参与后续分组、统计逻辑，进一步降低计算开销。
• 对词频字段做了显式类型转换，保证输出的postingList格式完全符合要求的List[Array[Long, Long]]结构。

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