Spark DataFrame获取item列值在另一数组列中的索引位置的实现方法

解决方案 #

你想实现的动态传入item列作为查找值的需求，在Spark 2.4及以上版本可以直接通过array_position原生实现，第二个参数本身就支持传入列类型，只需要给计算结果列设置别名即可：

from pyspark.sql import functions as F

# 新增pos列存储计算结果
df = df.withColumn("pos", F.array_position(F.col("ls_rec_items"), F.col("item")))

# 查看结果
df.show()

如果偏好select写法，也可以按如下方式实现：

df.select(
    "user", 
    "item", 
    "ls_rec_items",
    F.array_position(F.col("ls_rec_items"), F.col("item")).alias("pos")
).show()

运行以上代码即可得到你期望的输出：

+-------+-------+--------------------+-----+
|   user|   item|        ls_rec_items|  pos|
+-------+-------+--------------------+-----+
|    321|      3|  [4, 3, 2, 6, 1, 5]|    2|
|    123|      2|  [5, 6, 3, 1, 2, 4]|    5|
|    123|      7|  [5, 6, 3, 1, 2, 4]|    0|
+-------+-------+--------------------+-----+

低版本兼容方案（Spark <2.4） #

如果你使用的Spark版本较低没有内置array_position函数，可以通过自定义UDF实现相同逻辑：

from pyspark.sql import functions as F

# 定义UDF，匹配Spark array_position的1开始索引规则，不存在返回0
get_arr_pos = F.udf(lambda arr, target: arr.index(target) + 1 if target in arr else 0, "int")

df = df.withColumn("pos", get_arr_pos(F.col("ls_rec_items"), F.col("item")))
df.show()

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