获取作为 Spark 数据帧中一列的数组中项目的索引

Question

通过执行以下操作，我能够根据数组字段中是否存在特定值来过滤 Spark 数据帧（在 PySpark 中）：

from pyspark.sql.functions import array_contains
spark_df.filter(array_contains(spark_df.array_column_name, "value that I want")).show() 

有没有办法获取在数组中找到该项目的位置的索引？似乎应该存在，但我没有找到它。谢谢你。

Answer 1

在 spark 2.4+ 中，有以下array_position功能：

df = spark.createDataFrame([(["c", "b", "a"],), ([],)], ['data'])
df.show()
#+---------+
#|     data|
#+---------+
#|[c, b, a]|
#|       []|
#+---------+

from pyspark.sql.functions import array_position
df.select(df.data, array_position(df.data, "a").alias('a_pos')).show()
#+---------+-----+
#|     data|a_pos|
#+---------+-----+
#|[c, b, a]|    3|
#|       []|    0|
#+---------+-----+

文档中的注释：

1. 定位给定数组中给定值第一次出现的位置；

2. 该位置不是基于零的，而是基于 1 的索引。如果在数组中找不到给定的值，则返回 0。

Answer 2

我使用的是spark 2.3版本，所以我尝试使用udf。

df = spark.createDataFrame([(["c", "b", "a","e","f"],)], ['arraydata'])
+---------------+
|      arraydata|
+---------------+
|[c, b, a, e, f]|
+---------------+

user_func = udf (lambda x,y: [i for i, e in enumerate(x) if e==y ])

检查项目“b”的索引位置：

newdf = df.withColumn('item_position',user_func(df.arraydata,lit('b')))

>>> newdf.show();
+---------------+-------------+
|      arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]|          [1]|
+---------------+-------------+

检查项目“e”的索引位置：

newdf = df.withColumn('item_position',user_func(df.arraydata,lit('e')))

>>> newdf.show();
+---------------+-------------+
|      arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]|          [3]|
+---------------+-------------+