请求修改PySpark数据掩码代码：仅掩码字段前5个字符

原代码对字段进行全哈希掩码，现在需要调整为仅掩码前5个字符，例如将邮箱customername@hotmail.com转换为******mername@hotmail.com。以下是修改后的实现：

修改后的代码 #

from pyspark.sql import SparkSession  
from pyspark.sql.functions import udf  
import hashlib  
  
class Mask:  
    def __init__(self, salt: str):  
        self.salt = salt  
      
    def sha512(self, value):  
        return hashlib.sha512(f'{value}{self.salt}'.encode()).hexdigest()  
  
    def shake_128(self, value):  
        return hashlib.shake_128(f'{value}{self.salt}'.encode()).hexdigest(32)  

    # 新增：仅掩码前5个字符的方法
    def mask_first_five(self, value):
        if not value:  # 处理空值情况
            return value
        # 若字符串长度小于等于5，全部替换为******；否则替换前5个字符为******，保留剩余部分
        if len(value) <= 5:
            return "******"
        return "******" + value[5:]
  
    def register(self, spark: SparkSession):  
        spark.udf.register('sha512', self.sha512)  
        spark.udf.register('shake128', self.shake_128)
        spark.udf.register('mask_first_five', self.mask_first_five)  # 注册新UDF


spark = SparkSession.builder.getOrCreate()  
m = Mask('123456789')  
m.register(spark) 

代码说明 #

• 新增mask_first_five方法：
  • 先判断输入值是否为空，为空则直接返回
  • 如果字符串长度≤5，直接返回6个*；否则将前5个字符替换为6个*，拼接剩余部分（匹配示例效果）
• 在register方法中注册新UDFmask_first_five，支持在Spark SQL中直接调用

使用示例 #

读取CSV文件并对邮箱字段仅掩码前5个字符：

spark.read \  
    .format('csv') \  
    .option('inferSchema', True) \  
    .option('header', True) \  
    .load(path) \  
    .selectExpr(['user_name', 'mask_first_five(email) as masked_email']) \  
    .write \  
    .mode('append') \  
    .saveAsTable('my_table')

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