如何使用 pyspark 列值来索引 numpy 数组？

Question

我有一个 pyspark 数据框，其中两列代表数组的二维索引。我想将该数组的特定值作为新列添加到我的 df 中。

array = np.array([[1, 2, 3], [4, 5, 6]])

df = spark.createDataFrame(
    [
        (0, 2),
        (1, 1),
        (1, 2)
    ],
    ["x", "y"]  
)

我尝试将列传递给数组（这显然不起作用）。我收到错误： IndexError：只有整数、切片(':')、省略号('...')、numpy.newaxis ('None) 和整数或布尔数组是有效索引。

这是有道理的，因为我实际上尝试过通过一个专栏。

df.withColumn("value", array[col("x")][col("y")])

理想情况下，我的结果使用 df 中的索引来导航数组并提取值。

value = [3,5,6]

如果有任何帮助，我将不胜感激。我特别尝试利用 pyspark 功能来做到这一点。

Answer 1

我最终使用了矢量化 UDF 方法。

import pyspark.sql.functions as F
from pyspark.ml.feature import VectorAssembler
from pyspark.sql.types import IntegerType


vector =  VectorAssembler(inputCols =["x", "y"], 
                          outputCol= "v").transform(df)
def loc_img_value(x):
    return int(array[int(x[0])][int(x[1])])

get_img_value = F.udf(lambda x: loc_img_value(x), IntegerType())

vector = vector.withColumn("Value", get_img_value("v"))