PySpark交叉连接后自定义字符串匹配性能优化求助

问题概述 #

现有两个Spark DataFrame：

• df1：约40万条记录，含通过row_number()生成的索引列RN
• df2：约2.1万条记录

需执行交叉连接后，通过自定义StringMatch函数计算字符串相似度，过滤掉matchScore < 0.8的结果并写入ADLS。当前代码按RN分片处理，集群配置为1驱动节点+最多5个工作节点（总计48核，单工作节点8核32GB内存），但运行4小时无输出，需优化性能。

具体优化措施 #

1. 修复分片逻辑，覆盖全量数据 #

原分片起始从200000开始，且依赖错误的变量df_f.count()，导致仅处理部分数据。调整分片逻辑：

• 提前计算df1总条数，生成覆盖全部数据的分片范围，例如按每5万条分片（可根据实际情况调整）：

  df1_count = df1.count()
  split_range_val = [(i, min(i + 50000 - 1, df1_count - 1)) for i in range(0, df1_count, 50000)]
  

• 避免循环内重复调用count()，减少重复计算开销。

2. 调整 repartition 时机与策略 #

• 延后 repartition 操作：交叉连接后数据量极大（单10万分片会生成21亿条数据），此时执行repartition(96)会引发海量shuffle。应在过滤低匹配分数据后再做repartition，过滤后数据量会大幅降低，shuffle开销也会减少。
• 取消强制单分区输出：repartition(1)会将所有数据集中到单个节点，成为严重性能瓶颈。若需减少输出文件数量，改用coalesce(8)（与工作节点核数匹配），避免全量shuffle；或直接让Spark自动管理分区数。

3. 优化自定义字符串匹配UDF性能 #

自定义UDF是核心性能瓶颈，可从以下两点优化：

• 改用Pandas向量化UDF：替代普通UDF，利用Pandas批量处理提升计算效率，示例：

  from pyspark.sql.functions import pandas_udf
  import pandas as pd
  
  @pandas_udf("double")
  def string_match_vectorized(col_df1: pd.Series, col_df2: pd.Series, param: pd.Series) -> pd.Series:
      # 替换为你的字符串匹配逻辑，示例用fuzzywuzzy批量计算
      from fuzzywuzzy import fuzz
      return pd.Series([fuzz.ratio(a, b)/100 for a, b in zip(col_df1, col_df2)])
  

• 优先使用Spark内置函数：若业务允许，改用Spark内置的levenshtein等字符串相似度函数，避免UDF的序列化/反序列化开销。

4. 优化广播与内存配置 #

• 确认广播生效：执行df_cross.explain()查看物理计划，确保出现BroadcastHashJoin（说明broadcast(df2)生效）。若df2数据量超过默认广播阈值（10MB），可调整参数：

  spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "50mb")
  

• 调整Executor内存：针对32GB内存的工作节点，设置合理的内存分配：

  spark.conf.set("spark.executor.memory", "24g")
  spark.conf.set("spark.executor.memoryOverhead", "8g")
  

  确保有足够内存存储广播变量和执行计算，避免OOM或频繁GC。

5. 减少循环中的重复IO与计算 #

• 提前缓存数据：在循环前缓存df1和df2，避免每次循环重新读取源数据：

  df1.cache()
  df2.cache()
  

• 合并输出操作：不要每次分片都写入ADLS，将所有分片结果收集后统一写入，减少IO开销：

  from pyspark.sql import DataFrame
  from functools import reduce
  
  result_dfs = []
  for li_val in split_range_val:
      df_f = df1.filter((col('RN') >= li_val[0]) & (col('RN') <= li_val[1]))
      df_cross = df_f.crossJoin(broadcast(df2))
      df_cross = df_cross.withColumn('matchScore', string_match_vectorized(col('col_df1_1'), col('col_df2'), lit(2)))
      df_cross = df_cross.filter(df_cross['matchScore'] > 0.80)
      result_dfs.append(df_cross)
  
  # 合并所有结果并写入
  final_df = reduce(DataFrame.unionAll, result_dfs)
  final_df.coalesce(8).write.format('com.databricks.spark.csv')\
          .mode("append").option('header','true').csv(dest_adls+"/curatedFinal")
  

6. 监控与调试 #

• 查看Databricks的Spark UI：重点关注Stage执行时间、Shuffle读写量、任务数据倾斜情况，定位具体瓶颈阶段。
• 先做小数据测试：用1000条df1数据验证代码逻辑和优化效果，确认可行后再跑全量数据。

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