PySpark：基于另一DataFrame实现文本内容匿名化的需求

核心思路 #

先通过id关联原始文本数据和识别结果，再利用PySpark字符串操作，将指定偏移量、长度的文本段替换为等长星号。如果存在单条文本多个人名需要替换的场景，额外处理规则排序与批量替换逻辑。

基础场景（单条文本对应一个人名） #

假设每个id仅对应一条识别结果，直接使用内置字符串函数即可高效处理，无需UDF：

1. 创建示例测试数据

# 原始文本DataFrame
df_origin = spark.createDataFrame(
    [(1, "Lorem ipsum Jane dolor sit amet, consectetur adipiscing"),
     (2, "Ut enim ad minim veniam, Max nostrud exercitation")],
    ["id", "text"]
)

# 识别结果DataFrame
df_results = spark.createDataFrame(
    [(1, "Person", 34, 4, "Jane"),
     (2, "Person", 36, 3, "Max")],
    ["id", "category", "offsett", "length", "content"]
)

2. 执行匿名化处理

from pyspark.sql import functions as F

# 按id关联两个DataFrame
joined_df = df_origin.join(df_results, on="id", how="inner")

# 构造匿名化文本：前半段 + 等长星号 + 后半段
# 注意：PySpark的substring是1-based索引，若你的offsett是0-based，需加1转换
df_final = joined_df.withColumn(
    "anonymized_text",
    F.concat(
        # 截取偏移量之前的文本（0-based offsett对应1-based的1到offsett）
        F.substring(F.col("text"), 1, F.col("offsett")),
        # 生成对应长度的星号（Spark 2.4+支持repeat函数）
        F.repeat("*", F.col("length")),
        # 截取偏移量+长度之后的剩余文本
        F.substring(F.col("text"), F.col("offsett") + F.col("length") + 1, F.length(F.col("text")))
    )
).select("id", "anonymized_text").withColumnRenamed("anonymized_text", "text")

# 查看最终结果
df_final.show(truncate=False)

关键说明

• 确认偏移量offsett的索引类型：PySpark substring是1-based，若你的识别结果是0-based偏移量（多数NLP工具默认），代码中的转换逻辑已适配；如果是1-based，需调整substring的参数。
• 内置函数比UDF性能更高，适合单替换规则的场景。

进阶场景（单条文本对应多个人名） #

如果同一个id对应多条识别结果（单文本多个人名），需先聚合替换规则，再按偏移量从大到小替换（避免替换后偏移量错位）：

1. 聚合替换规则

# 按id分组，收集并按偏移量升序排序替换规则
grouped_results = df_results.groupBy("id").agg(
    F.sort_array(
        F.collect_list(F.struct(F.col("offsett"), F.col("length"))),
        asc=True
    ).alias("replace_rules")
)

2. 定义UDF处理批量替换

from pyspark.sql.types import StringType

def anonymize_multiple(text, rules):
    if not rules:
        return text
    # 从后往前替换，防止前面的替换影响后续偏移量
    sorted_rules = sorted(rules, key=lambda x: x["offsett"], reverse=True)
    for rule in sorted_rules:
        offset = rule["offsett"]
        length = rule["length"]
        text = text[:offset] + "*" * length + text[offset+length:]
    return text

# 注册UDF
anonymize_udf = F.udf(anonymize_multiple, StringType())

3. 执行批量匿名化

df_final_multi = df_origin.join(grouped_results, on="id", how="left").withColumn(
    "text",
    F.when(F.col("replace_rules").isNotNull(), anonymize_udf(F.col("text"), F.col("replace_rules"))).otherwise(F.col("text"))
).select("id", "text")

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