如何在PySpark中提取数组列元素并拆分存入新DataFrame

常见报错原因 #

PySpark数组下标从0开始计数，你需要提取的第4个元素对应下标为3，若之前调用getItem(4)会取到空值，导致逻辑不符合预期。

固定长度子数组拆分代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# 构造测试数据，实际使用时替换为自己的DataFrame即可
spark = SparkSession.builder.appName("array_split").getOrCreate()
test_data = [
    ([1,5,[],[2,3,22,42,3,-5]],),
    ([1,5,[],[-3,67,32,2,2.14,5]],),
    ([1,5,[],[32,1,3,34,6.7,90]],)
]
df = spark.createDataFrame(test_data, ["column_x"])

# 提取第4个元素（下标3）得到子数组
sub_array = col("column_x").getItem(3)

# 拆分生成新列
result_df = df.select(
    sub_array.getItem(0).alias("column1"),
    sub_array.getItem(1).alias("column2"),
    sub_array.getItem(2).alias("column3"),
    sub_array.getItem(3).alias("column4"),
    sub_array.getItem(4).alias("column5"),
    sub_array.getItem(5).alias("column6")
)

# 验证输出
result_df.show()

运行后输出结果与预期一致：

+-------+-------+-------+-------+-------+-------+
|column1|column2|column3|column4|column5|column6|
+-------+-------+-------+-------+-------+-------+
|      2|      3|     22|     42|      3|     -5|
|     -3|     67|     32|      2|   2.14|      5|
|     32|      1|      3|     34|    6.7|     90|
+-------+-------+-------+-------+-------+-------+

动态拆分方案（适配子数组长度不固定场景） #

如果子数组元素较多或长度不固定，可批量生成选择列，无需手动编写：

# 计算子数组最大长度，也可直接指定固定长度
max_sub_len = df.selectExpr("max(size(column_x[3]))").first()[0]
select_cols = [sub_array.getItem(i).alias(f"column{i+1}") for i in range(max_sub_len)]
result_df = df.select(*select_cols)

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