如何用PySpark列值索引NumPy数组并添加为DataFrame新列？

直接用numpy数组索引Spark Column对象会报错，因为numpy无法识别Column类型。下面提供两种可行的解决方案：

方案一：广播变量+UDF #

适合复杂数组访问逻辑，通过广播变量避免数组在任务间重复传输：

from pyspark.sql.functions import udf, broadcast, col
from pyspark.sql.types import IntegerType
import numpy as np

# 初始化数据
array = np.array([[1, 2, 3], [4, 5, 6]])
df = spark.createDataFrame(
    [(0, 2), (1, 1), (1, 2)],
    ["x", "y"]
)

# 广播数组到所有节点
broadcast_array = spark.sparkContext.broadcast(array)

# 定义UDF提取对应值
@udf(IntegerType())
def get_array_value(x, y):
    return broadcast_array.value[x][y]

# 添加新列
df = df.withColumn("value", get_array_value(col("x"), col("y")))
df.show()

执行后输出：

+---+---+-----+
|  x|  y|value|
+---+---+-----+
|  0|  2|    3|
|  1|  1|    5|
|  1|  2|    6|
+---+---+-----+

方案二：Spark原生数组函数（推荐） #

利用Spark原生数组操作，无需Python UDF，性能更优：

from pyspark.sql.functions import lit, col
from pyspark.sql.types import ArrayType, IntegerType
import numpy as np

# 初始化数据
array = np.array([[1, 2, 3], [4, 5, 6]])
df = spark.createDataFrame(
    [(0, 2), (1, 1), (1, 2)],
    ["x", "y"]
)

# 将numpy数组转为Python列表，再转为Spark二维数组类型
array_list = array.tolist()
spark_array = lit(array_list).cast(ArrayType(ArrayType(IntegerType())))

# Spark数组索引从1开始，需对x、y加1
df = df.withColumn("value", spark_array[col("x") + 1][col("y") + 1])
df.show()

关键说明 #

• 原代码报错原因：array[col("x")]试图用Spark Column对象索引numpy数组，但numpy仅支持Python原生的整数、切片等索引类型，无法识别Column对象。
• 方案二优势：Spark原生函数运行在JVM层面，避免了Python UDF的序列化/反序列化开销，大数据量下性能更出色。

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