Получить индекс элемента в массиве, который является столбцом в кадре данных 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

Я использую версию 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]|
+---------------+-------------+

Answer 2

В версии 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, если указанное значение не может быть найдено в массиве.