使用udf过滤掉火花dataframe

Question

我有一个 pyspark dataframe 有两列， name和source 。 name列中的所有值都是不同的。 Source有多个字符串，用逗号 (,) 分隔。 我想过滤掉source列中的任何字符串包含整个name列中的任何值的所有行。

我正在使用以下 UDF：

def checkDependentKPI(df, name_list):
for row in df.collect():
    for src in row["source"].split(","):
        for name in name_list:
            if name in src:
                return row['name']
        return row['name']

我的最终目标是将所有此类行放在 dataframe 的末尾。我该怎么做？

样本 dataframe：

+--------------------+--------------------+
|                name|              source|
+--------------------+--------------------+
|dev.................|prod, sum, diff.....|
|prod................|dev, diff, avg......|
|stage...............|mean, mode..........|
|balance.............|median, mean........|
|target..............|avg, diff, sum......|
+--------------------+--------------------+

Answer 1

也许这个？

from pyspark.sql import functions as psf 

test_data = [('dev','prod,sum,diff')
  , ('prod','dev,diff,avg')
  , ('stage','mean,mode')
  , ('balance','median,mean')
  , ('target','avg,diff,sum')]
  
df = spark.createDataFrame(test_data, ['kpi_name','kpi_source_table'])

df = df.withColumn('kpi_source_table', psf.split('kpi_source_table', ','))
df_flat = df.agg(psf.collect_list('kpi_name').alias('flat_kpi'))

df = df.join(df_flat, how='cross')

df = df.withColumn('match', psf.array_intersect('kpi_source_table', 'flat_kpi'))
display(df.orderBy('match'))

Answer 2

您可以使用like()来利用 SQL like 表达式，而无需任何繁重的collect()操作和循环检查。 假设您已经有一个name列表：

from functools import reduce

df.filter(
    reduce(lambda x, y: x|y, [func.col('source').like(f"%{pattern}%") for pattern in name])
).show(20, False)