Spark DataFrame中将结构体数组转换为字段数组的实现需求

要实现将Spark DataFrame中cust_detail数组内的同名字段合并为对应数组的转换，可以通过Spark内置聚合函数结合结构体构造完成。以下是具体实现：

输入输出结构说明 #

• 输入DataFrame结构：

{"emp_id":1,"emp_name":"John","cust_id":"c1","cust_detail":[{"name":"abc","acc_no":123,"mobile":0},{"name":"abc","acc_no":123,"mobile":111},{"name":"abc","acc_no":123,"mobile":222}]}

• 期望输出结构：

{"emp_id":1,"emp_name":"John","cust_id":"c1","cust_detail":[{"name":["abc"],"acc_no":[123],"mobile":[0,111,222]}]}

Python 实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, collect_set, collect_list, struct, array

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

# 构造输入数据
input_data = [
    (1, "John", "c1", [{"name":"abc","acc_no":123,"mobile":0}, {"name":"abc","acc_no":123,"mobile":111}, {"name":"abc","acc_no":123,"mobile":222}])
]

# 创建DataFrame
df = spark.createDataFrame(input_data, ["emp_id", "emp_name", "cust_id", "cust_detail"])

# 核心转换逻辑
processed_df = df.withColumn(
    "cust_detail",
    # 将合并后的结构体包装为数组，匹配输出格式
    array(
        struct(
            # 收集name字段值并去重（原数据重复，得到单元素数组）
            collect_set(col("cust_detail.name")).alias("name"),
            # 收集acc_no字段值并去重
            collect_set(col("cust_detail.acc_no")).alias("acc_no"),
            # 收集mobile字段所有值（保留全部原始数据）
            collect_list(col("cust_detail.mobile")).alias("mobile")
        )
    )
)

# 查看结果
processed_df.show(truncate=False)

# 输出JSON格式结果
print(processed_df.toJSON().collect()[0])

Scala 实现代码 #

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{col, collect_set, collect_list, struct, array}

object MergeArrayStructFields {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("MergeArrayStructFields").getOrCreate()
    import spark.implicits._

    // 构造输入数据
    val inputData = Seq(
      """{"emp_id":1,"emp_name":"John","cust_id":"c1","cust_detail":[{"name":"abc","acc_no":123,"mobile":0},{"name":"abc","acc_no":123,"mobile":111},{"name":"abc","acc_no":123,"mobile":222}]}"""
    )

    // 创建DataFrame
    val df = spark.read.json(inputData.toDS())

    // 核心转换逻辑
    val processedDF = df.withColumn(
      "cust_detail",
      array(
        struct(
          collect_set(col("cust_detail.name")).alias("name"),
          collect_set(col("cust_detail.acc_no")).alias("acc_no"),
          collect_list(col("cust_detail.mobile")).alias("mobile")
        )
      )
    )

    // 查看结果
    processedDF.show(false)

    // 输出JSON格式结果
    processedDF.toJSON.collect().foreach(println)
  }
}

关键逻辑说明 #

1. 字段聚合：
   • 用collect_set处理name和acc_no，原数组中这两个字段值重复，去重后得到单元素数组。
   • 用collect_list处理mobile，保留所有原始值不做去重。
2. 结构体重构：通过struct将聚合后的数组字段重新组合为新的结构体。
3. 数组包装：用array将新结构体包装为数组，与期望输出格式完全匹配。

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