PySpark DataFrame如何拆分Array列提取指定元素生成新数组列

最优方案采用PySpark内置高阶函数transform直接处理数组列，全程无shuffle开销，性能远高于展开数组再聚合的实现方式。

核心实现代码 #

首先导入依赖函数：

from pyspark.sql import functions as F

核心转换逻辑：

df = df.withColumn(
    "newcol1",
    F.transform(
        F.col("col1"),
        lambda x: F.element_at(F.split(x, "_"), 3)
    )
)

逻辑说明 #

• transform会遍历col1数组的每个元素，在单条数据内完成处理直接生成新数组，不需要展开重聚合，无跨分区数据交互开销
• split(x, "_")将单个字符串按下划线拆分为数组
• element_at(拆分后数组, 3)直接取第三个位置的元素，若数组长度不足3会自动返回null，刚好匹配需求中不存在第三段则为null的要求，不需要额外做空值判断

完整测试示例 #

# 构造测试数据
test_data = [
    (["a1_b1_c1", "a2_b2_c2", "a3_b3_c3"],),
    (["aa1_bb1_cc1", "aa2_bb2_cc2", "aa3_bb3"],),
    (["aaa2_bbb2_ccc1", "aaa2_bbb2_cc2", "aaa3_bbb3"],)
]
df = spark.createDataFrame(test_data, schema=["col1"])

# 执行转换
df = df.withColumn(
    "newcol1",
    F.transform(F.col("col1"), lambda x: F.element_at(F.split(x, "_"), 3))
)

# 查看输出结果
df.select("newcol1").show(truncate=False)

运行输出和需求完全一致：

+---------------------+
|newcol1              |
+---------------------+
|[c1, c2, c3]         |
|[cc1, cc2, null]     |
|[ccc1, ccc2, null]   |
+---------------------+

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