Pyspark: сопоставить значения в одном столбце со списком в той же строке в другом столбце

Question

У меня есть фрейм данных, содержащий следующие 2 столбца, среди прочих: 1. ID 2. list_IDs

Я пытаюсь создать 3-й столбец, возвращающий логическое значение True или False, если идентификатор присутствует в столбце list_IDв той же строке

Я пытался использовать следующее:

df = sqlContext.createDataFrame([(1, [1, 2, 3,]), (2, [1, 3, 4])], ("ID", "list_IDs"))

df.withColumn("IDmatch", when(col("ID").isin(F.col("list_IDs")), True).otherwise(False)).show()

Это не работает. Однако, если бы я предоставил некоторый статический список для сравнения, он, конечно, работает.

df.withColumn("IDmatch", when(col("ID").isin([2, 3]), True).otherwise(False)).show()

Я могу использовать udf для возврата логического типа, и это также работает:

@udf(returnType=BooleanType())
def isinlist(x, y):
    return x in y

Однако я стараюсь избегать использования UDF в этом случае, если это возможно, иМне было интересно, можно ли использовать что-то родное .isin (), чтобы проверить, присутствует ли идентификатор в строке в списке значений в столбце list_ID для той же строки?

Answer 1

Метод 1:

Если вы используете Spark> = 2.4.0. Вы можете использовать встроенную функцию arrays_overlap. Эта функция принимает 2 массива и проверяет общие элементы среди них.

from pyspark.sql.functions import arrays_overlap, array

df.withColumn("IDmatch", arrays_overlap(df.list_IDs, array(df.ID))).show()

Вывод:

+---+---------+-------+
| ID| list_IDs|IDmatch|
+---+---------+-------+
|  1|[1, 2, 3]|   true|
|  2|[1, 3, 4]|  false|
+---+---------+-------+

Подробнее об этом можно прочитать здесь, https://spark.apache.org/docs/2.4.0/api/python/pyspark.sql.html#pyspark.sql.functions.arrays_overlap

Метод 2:

В качестве альтернативы, вы также можете использовать udf для получения того же результата

from pyspark.sql.functions import udf

element_check = udf(lambda elt_list, elt: elt in elt_list)

df.withColumn("IDmatch", element_check(df.list_IDs, df.ID)).show()

Answer 2

Вы должны определить пользовательский udf, который принимает 2 аргумента

val isInList = spark.udf
.register("isInList", (id: Int, ids: Seq[Int]) => ids.contains(id))

Вывод:

df.withColumn("IDmatch", isInList($"ID", $"list_IDs")).show()
+---+---------+-------+
| ID| list_IDs|IDmatch|
+---+---------+-------+
|  1|[1, 2, 3]|   true|
|  2|[1, 3, 4]|  false|
+---+---------+-------+

Я использую синтаксис Scala, но общее направление должно быть ясным.