Выберите строку по значению в наборе после collect_set с помощью pyspark

Question

Используя

from pyspark.sql import functions as f

и методы f.agg и f.collect_set Я создал столбец colSet в dataFrame какследует:

+-------+--------+
| index | colSet |
+-------+--------+
|      1|[11, 13]|
|      2|  [3, 6]|
|      3|  [3, 7]|
|      4|  [2, 7]|
|      5|  [2, 6]|
+-------+--------+

Теперь, как это возможно, используя python / и pyspark, чтобы выбрать только те строки, где, например, 3 является элементом массива в записи colSet (где вообще естьможет быть намного больше, чем только две записи!)?

Я пытался использовать функцию udf следующим образом:

isInSet = f.udf( lambda vcol, val: val in vcol, BooleanType())

вызывается через

dataFrame.where(isInSet(f.col('colSet'), 3))

Я такжепопытался удалить f.col из вызывающей программы и использовать его вместо определения isInSet, но ни одна из них не сработала, я получаю исключение:

AnalysisException: cannot resolve '3' given input columns: [index, colSet]

Любая помощь приветствуется при выборе строк с определенной записью.(или даже лучше подмножество !!!), учитывая строку с результатом collect_set.

Answer 1

С исходным UDF все в порядке, но чтобы использовать его, вам нужно передать значение 3 в виде литерала:

dataFrame.where(isInSet(f.col('colSet'), f.lit(3)))

Но, как указывает jxc в комментарии, использование array_contains, вероятно, лучшеВыбор:

dataFrame.where(f.array_contains(f.col('colSet'), 3))

Я не проводил никаких сравнительных тестов, но в целом использование пользовательских функций в PySpark медленнее, чем использование встроенных функций, из-за обратной связи между JVM и интерпретатором Python.

Answer 2

Я нашел решение сегодня (после провала в пятницу вечером) без использования udf-метода:

[3 in x[0] for x in list(dataFrame.select(['colSet']).collect())]

Надеюсь, это поможет кому-то еще в будущем.