PySpark：实现两个函数按列表键过滤DataFrame字典列

现有如下PySpark DataFrame：

d1 = [({'the town': 1, 'County Council s': 2, 'email':5},2),
      ({'Mayor': 2, 'Indiana': 2}, 4),
      ({'Congress': 2, 'Justice': 2,'country': 2, 'veterans':1},6)
]
df1 = spark.createDataFrame(d1, ['dct', 'count'])
df1.show()

ignore_lst = ['County Council s', 'emal','Indiana']
filter_lst = ['Congress','town','Mayor', 'Indiana']

需要实现两个功能：

• 第一个函数：过滤dct列中不在ignore_lst内的键，生成新列
• 第二个函数：保留dct列中在filter_lst内的键，生成新列

可以通过PySpark的用户自定义函数（UDF）实现字典的键过滤逻辑，具体代码如下：

导入依赖模块 #

from pyspark.sql.functions import udf
from pyspark.sql.types import MapType, StringType, IntegerType

1. 过滤忽略列表内的键 #

定义UDF，保留字典中不在ignore_lst里的键值对：

def filter_ignore_keys(dct):
    if not dct:
        return {}
    return {k: v for k, v in dct.items() if k not in ignore_lst}

# 注册UDF，指定返回类型为字符串键、整数值的Map
filter_ignore_udf = udf(filter_ignore_keys, MapType(StringType(), IntegerType()))

2. 保留过滤列表内的键 #

定义UDF，只保留字典中存在于filter_lst里的键值对：

def keep_filter_keys(dct):
    if not dct:
        return {}
    return {k: v for k, v in dct.items() if k in filter_lst}

keep_filter_udf = udf(keep_filter_keys, MapType(StringType(), IntegerType()))

应用UDF生成新列 #

将两个UDF应用到原DataFrame，生成包含过滤后字典的新列：

df_result = df1.withColumn("filtered_dct_ignore", filter_ignore_udf(df1["dct"])) \
               .withColumn("filtered_dct_keep", keep_filter_udf(df1["dct"]))

df_result.show(truncate=False)

执行结果 #

+------------------------------------------------+-----+----------------------------------------+--------------------------+
|dct                                             |count|filtered_dct_ignore                     |filtered_dct_keep         |
+------------------------------------------------+-----+----------------------------------------+--------------------------+
|{the town -> 1, County Council s -> 2, email ->5}|2    |{the town -> 1, email -> 5}             |{}                        |
|{Mayor -> 2, Indiana -> 2}                      |4    |{Mayor -> 2}                            |{Mayor -> 2, Indiana -> 2}|
|{Congress -> 2, Justice -> 2, country ->2, veterans ->1}|6|{Congress ->2, Justice ->2, country ->2, veterans ->1}|{Congress ->2}|
+------------------------------------------------+-----+----------------------------------------+--------------------------+

补充说明 #

• 注意ignore_lst中的'emal'是拼写错误（原字典中为'email'），因此不会过滤'email'键，若需修正可直接调整列表内容。
• 如果需要匹配包含指定子串的键（比如用filter_lst中的'town'匹配字典里的'the town'），可修改UDF判断逻辑为：

def keep_filter_keys_contains(dct):
    if not dct:
        return {}
    return {k: v for k, v in dct.items() if any(sub in k for sub in filter_lst)}

内容的提问来源于stack exchange，提问作者Rory