Spark中将DataFrame的数组列拆分为多列的实现方法

实现步骤 #

核心思路是先将数组展开为单行单结构体，提取字段后通过pivot将行转列，最终聚合得到目标结果：

1. 使用explode函数拆分subjects数组列，把每个结构体单独拆成一行
2. 从结构体中提取subject和score字段
3. 以name分组，通过pivot将不同的subject值转为列，聚合获取对应的score

Python 代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode, col
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, ArrayType

# 初始化Spark会话
spark = SparkSession.builder.appName("struct_array_to_cols").getOrCreate()

# 构造输入数据与Schema
data = [
    ("Michael", [{"subject": "english", "score": 100}, {"subject": "math", "score": 80}]),
    ("John", [{"subject": "english", "score": 75}, {"subject": "math", "score": 66}])
]

schema = StructType([
    StructField("name", StringType()),
    StructField("subjects", ArrayType(StructType([
        StructField("subject", StringType()),
        StructField("score", IntegerType())
    ])))
])

df = spark.createDataFrame(data, schema)

# 1. 展开数组列
exploded_df = df.select("name", explode(col("subjects")).alias("subject_info"))
# 2. 提取结构体中的字段
extracted_df = exploded_df.select(
    "name",
    col("subject_info.subject").alias("subject"),
    col("subject_info.score").alias("score")
)
# 3. 行转列并聚合
result_df = extracted_df.groupBy("name").pivot("subject").agg({"score": "first"})

# 查看结果
result_df.show()

Scala 代码示例 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{explode, col}
import org.apache.spark.sql.types.{StructType, StructField, StringType, IntegerType, ArrayType}

object StructArrayToCols {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("struct_array_to_cols").getOrCreate()

    // 构造输入数据
    val data = Seq(
      ("Michael", Array(("english", 100), ("math", 80))),
      ("John", Array(("english", 75), ("math", 66)))
    ).map { case (name, subs) =>
      (name, subs.map { case (sub, score) => (sub, score) })
    }

    // 定义Schema
    val schema = StructType(Seq(
      StructField("name", StringType),
      StructField("subjects", ArrayType(StructType(Seq(
        StructField("subject", StringType),
        StructField("score", IntegerType)
      ))))
    ))

    val df = spark.createDataFrame(data).toDF(schema.fieldNames: _*)

    // 处理流程
    val explodedDF = df.select(col("name"), explode(col("subjects")).alias("subject_info"))
    val extractedDF = explodedDF.select(
      col("name"),
      col("subject_info.subject").alias("subject"),
      col("subject_info.score").alias("score")
    )
    val resultDF = extractedDF.groupBy("name").pivot("subject").agg(org.apache.spark.sql.functions.first("score"))

    // 输出结果
    resultDF.show()
  }
}

注意事项 #

• 如果同一name下存在重复的subject记录，聚合时可根据业务需求选择first、max、sum等函数
• 若subject的取值范围已知，在pivot时指定列名列表（如pivot("subject", Seq("english", "math"))）可提升执行性能

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