如何在PySpark DataFrame中拆分二元子数组组成的数组列为两列？

不用转换为RDD，直接用PySpark的DataFrame API就能高效实现需求，核心是利用transform函数（或SQL表达式形式的transform）对数组列的每个子数组提取对应元素：

方法1：适用于PySpark 3.1及以上版本 #

直接调用原生transform函数，遍历数组中的每个子数组，提取指定位置的元素：

from pyspark.sql import functions as F

# 假设你的DataFrame名为df，目标数组列为arr_col
df = df.withColumn("x_values", F.transform(F.col("arr_col"), lambda sub_arr: sub_arr[0])) \
       .withColumn("y_values", F.transform(F.col("arr_col"), lambda sub_arr: sub_arr[1]))

方法2：兼容PySpark 3.1以下版本 #

如果你的PySpark版本低于3.1，可通过expr函数使用SQL风格的transform表达式：

from pyspark.sql import functions as F

df = df.withColumn("x_values", F.expr("transform(arr_col, sub_arr -> sub_arr[0])")) \
       .withColumn("y_values", F.expr("transform(arr_col, sub_arr -> sub_arr[1])"))

原理说明 #

transform函数会逐个处理数组中的每个子数组，将所有子数组的第0个元素收集为新数组x_values，第1个元素收集为新数组y_values。该操作完全在DataFrame层面执行，PySpark的Catalyst优化器会对其进行查询优化，相比转换为RDD的map操作，避免了序列化/反序列化的开销，执行效率更高。

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