在整行上使用udf过滤Pyspark数据帧

Question

有没有办法选择整行作为一列输入到Pyspark过滤器udf中？

我有一个复杂的过滤函数“my_filter”，我想应用于整个数据帧：

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

但是

col("*")

引发错误，因为这不是有效的操作。

我知道我可以将数据帧转换为RDD，然后使用RDD的过滤方法，但我不想将其转换为RDD，然后再转换回数据帧。我的数据帧具有复杂的嵌套类型，因此当我尝试再次将 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()

产生相同的输出。