如何提升Spark中对结构体数组执行filter()的性能？

原方案对每行数组执行30次filter操作，计算量随键数量线性增长，导致性能瓶颈。核心优化思路是先对数组做一次分组聚合，将相同Key的Value合并到一个数组中，之后再基于分组结果构建目标结构体，全程仅需遍历数组一次。

具体实现逻辑 #

1. 数组分组：利用Spark的aggregate高阶函数，将原数组转换为Map<String, Array<String>>结构，Key对应原结构体的Key，Value对应该Key下所有Value的集合。
2. 构建目标结构体：遍历目标Schema的每个字段，根据字段类型从分组后的Map中提取对应值：
   • 若字段为Array<String>类型：直接取Map中对应的Value数组
   • 若字段为String类型：取Map中对应数组的第一个元素（数组为空或不存在时返回null）

优化后的Scala代码 #

import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._

def optimizeTransform(column: Column, targetSchema: StructType): Column = {
  // 1. 将原数组按Key分组，生成Key到Value数组的映射
  val groupedMap = aggregate(
    column,
    lit(Map.empty[String, Array[String]]).cast(MapType(StringType, ArrayType(StringType))),
    (acc, elem) => {
      val key = elem("Key").cast(StringType)
      val value = elem("Value").cast(StringType)
      // 已存在的Key追加Value，不存在则新建数组
      when(acc.contains(key),
        acc + (key -> array_union(acc(key), array(value)))
      ).otherwise(
        acc + (key -> array(value))
      )
    },
    identity // 单条数据的数组聚合无需跨分区合并
  )

  // 2. 遍历目标Schema，构建结构体字段
  struct(targetSchema.fields.map { field =>
    val fieldName = field.name
    field.dataType match {
      case ArrayType(StringType, _) =>
        groupedMap.get(fieldName).alias(fieldName)
      case StringType =>
        element_at(groupedMap.get(fieldName), 1).alias(fieldName)
    }
  }: _*)
}

性能提升说明 #

• 原方案时间复杂度为O(N*K)（N为数组元素数量，K为键的数量），优化后降至O(N + K)，当K=30时，计算量可减少约97%。
• 避免了重复遍历数组的冗余操作，简化Spark执行计划，降低资源消耗。

使用示例 #

假设你的DataFrame名为df，目标Schema已定义为targetSchema，调用方式如下：

val resultDF = df.withColumn("Collection", optimizeTransform(col("Collection"), targetSchema))

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