PySpark Pandas中tolist()的替代方案及大数据性能优化问询

核心问题分析 #

你的代码性能瓶颈集中在将分布式ps.DataFrame强制转换为本地Pandas/DataFrame/列表，比如to_pandas()、tolist()、np.concatenate(...)，这会把集群上的全量数据拉到Driver节点内存，直接触发OOM和执行器丢失。下面针对关键代码块给出替代方案：

1. 替换np.concatenate(df["key_phrase_extract"].tolist()) #

原代码是为了收集所有关键词生成target_map，但全量拉取列表会占用巨量内存。改用Spark分布式操作提取唯一关键词：

# 替代原来的try-except块
# 分布式flatMap展开所有关键词，再去重，最后仅收集唯一值到Driver
targets = df["key_phrase_extract"].explode().distinct().tolist()
# 如果是空的会返回空列表，无需额外异常处理

这样只收集唯一关键词，而非全量重复数据，内存占用大幅降低。

2. 避免将ps.DataFrame转成本地Pandas处理 #

原代码中df = df.to_pandas()、revised_df = get_revised_keywords().toPandas()后用循环遍历生成字典，再用apply赋值，完全放弃了Spark的分布式计算能力。改用分布式UDF+Spark Join实现：

重构后的核心代码示例

import pyspark.sql.functions as F
from pyspark.sql.types import ArrayType, StringType
import pandas as pd

def get_product_keyword(df: ps.DataFrame, bu: constants.BU) -> ps.DataFrame:
    keywords = tools.get_keywords(bu, en=False)
    
    # --- 泰语关键词提取：用mapPartitions分布式处理每个分区 ---
    def extract_keywords_partition(partition):
        # 每个分区初始化一次模型，避免重复加载
        embedder = SentenceTransformer("mrp/simcse-model-m-bert-thai-cased")
        part_df = pd.DataFrame(partition)
        keyword_dict = thai_transformer.extractRelevantThaiwords(part_df, th_model=embedder, parse="dataframe")
        
        # 生成(product_code, keyword)行级数据
        rows = []
        for keyword, product_data in keyword_dict.items():
            product_list = list(product_data.values())[0][-1]
            for product_code in product_list:
                rows.append({"product_code": product_code, "keyword": keyword})
        return rows
    
    # 转换为Spark DataFrame进行分布式处理
    keyword_pair_df = df.rdd.mapPartitions(extract_keywords_partition).toDF()
    
    # --- 关键词修正：用Spark Join替代本地字典映射 ---
    revised_df = get_revised_keywords()  # 假设返回ps.DataFrame
    merged_df = keyword_pair_df.join(revised_df, keyword_pair_df.keyword == revised_df.key, how="left")
    merged_df = merged_df.withColumn("revised_keyword", F.coalesce(merged_df.value, merged_df.keyword))
    
    # --- 按product_code聚合去重关键词 ---
    result_df = merged_df.groupBy("product_code").agg(F.collect_set("revised_keyword").alias("keywords_extract"))
    
    # --- 关联回原DataFrame ---
    df = df.join(result_df, on="product_code", how="left")
    df = df.withColumn("keywords_extract", F.coalesce(df.keywords_extract, F.array()))
    
    # --- 中文关键词处理：用分布式UDF替代apply ---
    else:
        # 用Spark内置split替代apply
        df = df.withColumn("key_phrase_extract", F.split(df["key_phrase_extract"], ","))  # 分隔符按需调整
    
    # 生成target_map：仅收集唯一关键词
    targets = df["key_phrase_extract"].explode().distinct().tolist()
    tf = Transformer(lang=constants.Language.zh_hk, keywords=keywords)
    target_map = tf.get_transform_dict(bu, targets, execute_cache=False)
    
    # 广播变量：避免每个任务重复加载字典
    target_map_bc = spark.sparkContext.broadcast(target_map)
    block_list_bc = spark.sparkContext.broadcast(block_list_cn)
    
    @F.udf(returnType=ArrayType(StringType()))
    def process_keywords(words):
        filtered = []
        for word in words:
            if word in target_map_bc.value and len(target_map_bc.value[word]) > 1:
                filtered.append(target_map_bc.value[word])
        # 去重+过滤黑名单
        return list(set([w for w in filtered if w not in block_list_bc.value]))
    
    df = df.withColumn("keywords_extract", process_keywords(df["key_phrase_extract"]))
    
    return df

3. 关键优化点说明 #

• mapPartitions替代全量处理：每个分区初始化一次SentenceTransformer模型，避免重复加载，同时仅处理单分区数据，内存可控。
• 广播变量：把target_map和block_list_cn转成广播变量，分发到每个执行器，避免每个任务重复加载字典。
• Spark原生API替代本地循环：用join、groupBy、agg等分布式操作替代本地字典映射和apply，利用集群算力。
• 仅收集必要数据：用explode().distinct().tolist()替代全量tolist()，只收集唯一关键词，减少Driver内存占用。

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