PySpark：含哈希顶层键的嵌套MapType子键过滤及UDF实现问询

场景说明 #

我们有一个PySpark DataFrame，其中顶层字段是哈希值作为键的MapType，每个哈希键对应的值又是一个嵌套的MapType。现在需要移除嵌套Map里指定的键（比如fname和lname），保留其余键值对。

输入示例（DataFrame结构） #

假设DataFrame的列名为data，数据结构如下：

{
  'h#9l00' : {
     'fname': 'fname1',
     'lname': 'lname1',
     'salary': 100,
     'city': 'xyz'
  },
  'o*5ftr': {
     'fname': 'fname2',
     'lname': 'lname2',
     'city': 'xyz'
  }
}

预期输出 #

处理后data列的结构：

{
  'h#9l00' : {
     'salary': 100,
     'city': 'xyz'
  },
  'o*5ftr': {
     'city': 'xyz'
  }
}

实现方案：使用PySpark UDF #

可以通过自定义UDF遍历顶层Map的每个键值对，对嵌套Map执行过滤操作。以下分两种场景给出实现：

1. 移除指定键列表 #

若已知需要移除的键（如fname、lname），可编写UDF过滤掉这些键：

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.sql.types import MapType, StringType, StructType, ObjectType

# 初始化SparkSession
spark = SparkSession.builder.appName("NestedMapFilter").getOrCreate()

# 定义过滤逻辑：移除嵌套Map中的指定键
def filter_nested_map_remove(top_map, keys_to_remove):
    result = {}
    for top_key, nested_map in top_map.items():
        filtered_nested = {k: v for k, v in nested_map.items() if k not in keys_to_remove}
        result[top_key] = filtered_nested
    return result

# 注册UDF，指定返回类型（适配嵌套Map的多值类型）
filter_remove_udf = udf(
    lambda x: filter_nested_map_remove(x, {'fname', 'lname'}),
    MapType(StringType(), MapType(StringType(), ObjectType()))
)

# 构造示例DataFrame
data = [
    ({
        'h#9l00': {'fname': 'fname1', 'lname': 'lname1', 'salary': 100, 'city': 'xyz'},
        'o*5ftr': {'fname': 'fname2', 'lname': 'lname2', 'city': 'xyz'}
    },)
]
df = spark.createDataFrame(
    data,
    schema=StructType().add("data", MapType(StringType(), MapType(StringType(), ObjectType())))
)

# 应用UDF得到结果
filtered_df = df.withColumn("filtered_data", filter_remove_udf(df["data"]))
filtered_df.select("filtered_data").show(truncate=False)

2. 保留指定键列表 #

若已知需要保留的键（如salary、city），可修改UDF逻辑为只保留这些键：

# 定义过滤逻辑：只保留嵌套Map中的指定键
def filter_nested_map_keep(top_map, keys_to_keep):
    result = {}
    for top_key, nested_map in top_map.items():
        filtered_nested = {k: v for k, v in nested_map.items() if k in keys_to_keep}
        result[top_key] = filtered_nested
    return result

# 注册UDF
filter_keep_udf = udf(
    lambda x: filter_nested_map_keep(x, {'salary', 'city'}),
    MapType(StringType(), MapType(StringType(), ObjectType()))
)

# 应用UDF得到结果
filtered_df_keep = df.withColumn("filtered_data", filter_keep_udf(df["data"]))
filtered_df_keep.select("filtered_data").show(truncate=False)

注意事项 #

• 类型匹配：UDF的返回类型需与原始DataFrame的Map类型严格匹配。如果嵌套Map的值类型多样（同时包含字符串、整数等），可用ObjectType()兼容；若类型固定，建议指定具体类型以保证性能。
• 性能优化：数据量极大时，UDF的性能可能不如原生Spark函数。如果嵌套Map的键是固定的，可考虑将嵌套Map展开为StructType，用drop或select操作后再重新组装成Map，但该方法仅适用于键已知且固定的场景。

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