Spark中提取struct_key并按role过滤转换数组结构体的问题

场景回顾 #

输入Spark DataFrame的结构包含嵌套数组：

• 顶层数组a，每个元素是结构体：struct(struct_key: string, two: array<struct(name: string, role: string)>)
• 需要按role过滤后，生成以role为分组、每个组包含struct(struct_key, name)的结构体数组。

之前用转Map的方法丢失struct_key，核心问题是没有将外层的struct_key与内层two的元素绑定，下面提供两种高效的解决方式：

方法一：使用高阶函数（Spark 2.4+，无需展开数组） #

适合需要保留原DataFrame行结构的场景，全程在数组内操作，避免数据shuffle：

步骤&代码（Scala） #

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

// 假设你的输入DataFrame为df
val processedDf = df.withColumn(
  "role_groups",
  // 1. 遍历数组a，将每个元素的struct_key与two数组的元素绑定
  transform($"a", elem => 
    transform(elem.two, twoElem => 
      struct(
        elem.struct_key.alias("struct_key"),
        twoElem.name.alias("name"),
        twoElem.role.alias("role")
      )
    )
  )
  // 2. 扁平化所有嵌套数组，得到一维的绑定后元素数组
  .alias("flattened_pairs")
  // 3. 按role分组，收集对应的struct_key+name结构体
  .expr("""
    aggregate(
      flattened_pairs,
      cast(array() as array<struct<role:string, items:array<struct<struct_key:string, name:string>>>>),
      (acc, x) -> case
        when size(filter(acc, y -> y.role = x.role)) = 0 then
          array_union(acc, array(struct(x.role as role, array(struct(x.struct_key, x.name)) as items)))
        else
          transform(acc, y -> if(y.role = x.role, struct(y.role as role, array_union(y.items, array(struct(x.struct_key, x.name))) as items), y))
        end
    )
  """)
)

输出示例 #

role_groups列的结构为：

array<struct<role:string, items:array<struct<struct_key:string, name:string>>>>

对应数据样例：

[
  {"role": "admin", "items": [{"struct_key": "key1", "name": "alice"}, {"struct_key": "key2", "name": "charlie"}]},
  {"role": "user", "items": [{"struct_key": "key1", "name": "bob"}]}
]

方法二：使用Explode+GroupBy（直观易理解） #

适合需要单独处理每个role分组，或对Spark版本较低的场景：

步骤&代码（Scala） #

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

// 1. 展开顶层数组a，保留struct_key
val explodedA = df.select(explode($"a").alias("a_elem"))

// 2. 展开内层数组two，将struct_key与name/role放在同一行
val explodedTwo = explodedA.select(
  $"a_elem.struct_key",
  explode($"a_elem.two").alias("two_elem")
).select(
  $"struct_key",
  $"two_elem.name",
  $"two_elem.role"
)

// 3. （可选）过滤特定role，比如只保留admin
// val filtered = explodedTwo.filter($"role" === "admin")

// 4. 按role分组，收集struct_key+name的结构体数组
val groupedDf = explodedTwo.groupBy($"role")
  .agg(collect_list(struct($"struct_key", $"name")).alias("items"))

输出示例 #

groupedDf的结构：

root
 |-- role: string (nullable = true)
 |-- items: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- struct_key: string (nullable = true)
 |    |    |-- name: string (nullable = true)

对应数据样例：

+-----+-----------------------------------+
|role |items                              |
+-----+-----------------------------------+
|admin|[{key1, alice}, {key2, charlie}]   |
|user |[{key1, bob}]                      |
+-----+-----------------------------------+

关键注意点 #

• 无论哪种方法，核心都是将外层的struct_key与内层two的元素绑定，避免处理内层时丢失上下文信息。
• 如果需要为每个role单独生成列，可以基于方法一的结果，用transform+filter提取指定role的数组；或基于方法二的结果，用pivot将role转为列。

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