优化PySpark关键词统计代码：解决S3文件检索性能问题

我们的S3存储桶里有大量Parquet文件，目前用PySpark代码读取每个文件、提取列并查找指定关键词，最终返回包含关键词计数、对应列及文件名的DataFrame。但这段PySpark代码的检索性能远不如本地Python代码——本地代码虽然快，但每天都要配置AWS密钥，而且处理大文件时会触发内存错误。我们试过把like换成contains，但性能没改善。现在用的是Databricks环境，希望优化PySpark代码提升性能。

当前PySpark代码 #

import s3fs
fs = s3fs.S3FileSystem()
from pyspark.sql.functions import lower, col
keywords = ['%keyword1%','%keyword2%']
prefix = ''
deployment_id = ''
pull_id = ''
paths = fs.ls(prefix+'/'+deployment_id+'/'+pull_id)
result = []
errors = []
try:
    for path in paths:
        df = spark.read.parquet('s3://'+path)
        print(path)
        for keyword in keywords:
            for col in df.columns:
                filtered_df = df.filter(lower(df[col]).like(keyword))
                filtered_count = filtered_df.count()
                if filtered_count > 0 :
                    #print(col +' has '+ str(filtered_count) +' appearences')
                    result.append({'keyword': keyword, 'column': col, 'count': filtered_count,'table':path.split('/')[-1]})


except Exception as e:
    errors.append({'error_msg':e})


try:
    errors = spark.createDataFrame(errors)
except Exception as e:
    print('no errors')
try:
    result = spark.createDataFrame(result)
    result.display()
except Exception as e:
    print('problem with results. May be no results')

本地Python代码片段 #

def search_df(self,keyword,df,regex=False):
        start=time.time()
        
        if regex:
            mask = df.applymap(lambda x: re.search(keyword,x) is not None if isinstance(x,str) else False).to_numpy()
        else:
            mask = df.applymap(lambda x: keyword.lower() in x.lower() if isinstance(x,str) else False).to_numpy()

1. 批量读取文件，放弃单文件循环 #

原代码逐个读取S3文件，完全浪费了Spark的并行计算能力。直接用通配符读取整个目录的Parquet，同时通过input_file_name()获取文件名：

from pyspark.sql.functions import input_file_name, lower, col, sum as spark_sum, lit

# 直接读取目标目录下所有Parquet文件
df = spark.read.parquet(f"s3://{prefix}/{deployment_id}/{pull_id}/*")
# 提取文件名（去掉路径前缀）
df = df.withColumn("table", input_file_name())
df = df.withColumn("table", col("table").substr(col("table").lastIndexOf("/") + 1, col("table").length()))

2. 合并检索逻辑，减少重复计算 #

原代码嵌套循环关键词和列，每个组合都触发一次count()，会生成大量Spark Job，调度开销极大。改成一次性生成所有匹配标记，再聚合计数：

# 注意去掉关键词里的%，contains不需要通配符
keywords = ['keyword1', 'keyword2']
# 只保留字符串类型列，非字符串列无需匹配
string_cols = [c for c, dtype in df.dtypes if dtype == "string"]

# 生成所有匹配条件的列
match_columns = []
for keyword in keywords:
    for col_name in string_cols:
        # 标记该列是否匹配当前关键词
        match_col = f"match_{keyword}_{col_name}"
        match_columns.append(lower(col(col_name)).contains(keyword).cast("int").alias(match_col))

# 添加匹配列到DataFrame
df_with_matches = df.select("table", *match_columns)

# 聚合统计每个文件、关键词、列的匹配次数
aggregations = []
for keyword in keywords:
    for col_name in string_cols:
        match_col = f"match_{keyword}_{col_name}"
        aggregations.extend([
            spark_sum(col(match_col)).alias("count"),
            lit(keyword).alias("keyword"),
            lit(col_name).alias("column")
        ])

# 分组聚合并过滤掉无匹配的结果
result_df = df_with_matches.groupBy("table").agg(*aggregations)\
    .filter(col("count") > 0)\
    .select("keyword", "column", "count", "table")

result_df.display()

3. 移除s3fs依赖，用Databricks原生S3支持 #

Databricks已经集成了S3访问能力，不需要手动用s3fs列文件，直接用Spark路径通配符更高效，还能避免额外的客户端开销。

4. 优化数据处理流程 #

• 提前过滤非字符串列：只处理需要匹配的字符串类型列，减少数据量。
• 避免多次Action操作：原代码每个count()都是一次Action，改用聚合操作只触发一次计算。
• 合并小文件：如果S3里有大量小Parquet文件，可以先合并成大文件（比如用repartition或coalesce），减少IO次数。

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