PySpark中实现类似explode的数组列拆分转多列功能

我来帮你搞定这个问题！针对把DenseVector类型的列拆分成多个独立列的需求，这里有两种实用的解决方案，适配不同的场景：

方法一：固定向量长度的快速拆分 #

如果你已经明确知道DenseVector的长度（比如示例里的3），可以直接把向量转成数组后提取对应位置的元素：

from pyspark.sql import Row
from pyspark.mllib.linalg import DenseVector
from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType, DoubleType

# 初始化你的示例DataFrame
df = spark.createDataFrame([Row(a=1, intlist=DenseVector([1,2,3])), Row(a=2, intlist=DenseVector([4,5,6]))])

# 定义UDF，把DenseVector转换成Python列表（对应Spark的ArrayType）
vec_to_array = udf(lambda vec: vec.toArray().tolist(), ArrayType(DoubleType()))

# 将DenseVector列转为数组列
df_array = df.withColumn("intlist_arr", vec_to_array(col("intlist")))

# 提取数组中的每个元素作为独立列，并重命名
df_result = df_array.select(
    col("a"),
    col("intlist_arr").getItem(0).alias("_1"),
    col("intlist_arr").getItem(1).alias("_2"),
    col("intlist_arr").getItem(2).alias("_3")
).drop("intlist_arr")

# 查看结果
df_result.show()

运行后就能得到你想要的输出：

+---+---+---+---+
|  a| _1| _2| _3|
+---+---+---+---+
|  1|1.0|2.0|3.0|
|  2|4.0|5.0|6.0|
+---+---+---+---+

方法二：动态适配任意向量长度的通用方案 #

如果你的DenseVector长度不固定，或者不想硬编码索引，可以先获取向量的长度，再动态生成要提取的列：

from pyspark.sql import Row
from pyspark.mllib.linalg import DenseVector
from pyspark.sql.functions import udf, col
from pyspark.sql.types import ArrayType, DoubleType

# 初始化DataFrame
df = spark.createDataFrame([Row(a=1, intlist=DenseVector([1,2,3])), Row(a=2, intlist=DenseVector([4,5,6]))])

# 向量转数组的UDF
vec_to_array = udf(lambda vec: vec.toArray().tolist(), ArrayType(DoubleType()))
df_array = df.withColumn("intlist_arr", vec_to_array(col("intlist")))

# 获取DenseVector的长度（取第一行的向量长度即可）
vec_length = df.select(col("intlist")).first()[0].size

# 动态生成要选择的列：保留原a列，加上数组中每个位置的元素列
select_cols = [col("a")] + [col("intlist_arr").getItem(i).alias(f"_{i+1}") for i in range(vec_length)]

# 生成结果DataFrame
df_result = df_array.select(*select_cols).drop("intlist_arr")

df_result.show()

这个方法会自动根据向量的长度生成对应的列，不管是3维、5维还是其他长度都能适用。

小提示 #

• 注意DenseVector来自pyspark.mllib.linalg，如果是pyspark.ml.linalg里的DenseVector，用法是完全一样的，toArray()方法同样适用。
• 如果你使用的是Spark 3.0+，其实可以不用自定义UDF，直接用Spark内置的vector_to_array函数（性能更优），替换UDF部分的代码：

  from pyspark.sql.functions import vector_to_array
  df_array = df.withColumn("intlist_arr", vector_to_array(col("intlist")))
  

  这样代码更简洁，运行效率也更高。

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