Фильтр Pyspark Dataframe с udf на всю строку

Question

Есть ли способ выбрать всю строку в виде столбца для ввода в фильтр Pyspark udf?

У меня есть сложная функция фильтрации "my_filter", которую я хочу применить ко всему DataFrame:

my_filter_udf = udf(lambda r: my_filter(r), BooleanType())
new_df = df.filter(my_filter_udf(col("*"))

Но

col("*")

выдает ошибку, потому что это недопустимая операция.

Я знаю, что могу преобразовать фрейм данных в RDD, а затем использовать метод фильтра RDD, но я НЕ хочу конвертировать его в RDD и затем обратно в фрейм данных. Мой DataFrame имеет сложные вложенные типы, поэтому вывод схемы завершается неудачно, когда я снова пытаюсь преобразовать RDD в информационный кадр.

Answer 1

Вы должны писать все столбцы статически. Например:

from pyspark.sql import functions as F

# create sample df
df = sc.parallelize([
     (1, 'b'),
     (1, 'c'),

 ]).toDF(["id", "category"])

#simple filter function
@F.udf(returnType=BooleanType())
def my_filter(col1, col2):
    return (col1>0) & (col2=="b")

df.filter(my_filter('id', 'category')).show()

Результаты:

+---+--------+
| id|category|
+---+--------+
|  1|       b|
+---+--------+

Если у вас так много столбцов, и вы уверены, что порядок столбцов:

cols = df.columns
df.filter(my_filter(*cols)).show()

Дает тот же вывод.