在PySpark中展开嵌套数组并提取第三个cn数组的方法

实现逻辑 #

1. 确认嵌套结构：先通过printSchema()查看数据表的Schema，明确cn数组所在的嵌套层级（比如是在某个数组字段的元素结构体中）。
2. 定位目标数组：根据层级关系，利用PySpark的数组索引语法或element_at函数，定位到第三个cn数组（注意：PySpark中col("arr")[index]索引从0开始，element_at函数索引从1开始）。
3. 处理边界情况：如果父数组长度不足3，通过size函数判断后返回空值，避免索引越界报错。

示例代码 #

1. 环境初始化与数据加载 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, element_at, size, when

# 创建SparkSession
spark = SparkSession.builder.appName("ExtractThirdCn").getOrCreate()

# 读取数据（以JSON格式为例，根据你的数据源调整）
df = spark.read.json("your_data_path.json")

# 查看Schema确认结构
df.printSchema()

2. 提取第三个cn数组（两种方式） #

假设你的数据Schema对应示意图结构如下：

|-- items: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- cn: array (nullable = true)
 |    |    |    |-- element: string (containsNull = true)

方式一：使用数组下标索引（从0开始，第三个元素对应索引2）

# 提取第三个cn数组，处理数组长度不足的情况
df_result = df.withColumn(
    "third_cn_array",
    when(size(col("items")) >= 3, col("items")[2]["cn"]).otherwise(None)
)

方式二：使用element_at函数（从1开始，第三个元素对应参数3）

# 直接提取items数组中第三个元素的cn字段
df_result = df.withColumn(
    "third_cn_array",
    when(size(col("items")) >= 3, element_at(col("items"), 3)["cn"]).otherwise(None)
)

3. 查看结果 #

# 展示提取后的数组内容
df_result.select("third_cn_array").show(truncate=False)

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