Spark ApplyInPandas原理与调优：Sentence Transformers句子编码实践

1. applyInPandas底层运行机制与模型副本问题 #

• 底层运行逻辑：applyInPandas会将Spark DataFrame按指定分组键拆分，每个分组的数据转换为Pandas DataFrame后，发送到Executor对应的Python Worker进程中执行你定义的处理函数。每个Python Worker对应一个Executor（JVM进程），同一个Executor内的多个Task会复用这个Python Worker进程。
• 模型副本情况：如果模型是在处理函数内部初始化（推荐方式），每个Executor的Python Worker进程会加载一份模型副本；如果在Driver端初始化模型再传递给函数，模型会被序列化后发送到每个Executor的Worker，本质也是每个Executor持有一份。不会每个Task都加载新模型，同Executor内的Task可以复用已加载的模型，避免重复开销。

2. 虚拟分组列数量选择与内存负载问题 #

• 分组数量建议：选数百个更合理。数千个分组会导致Task数量过多，增加调度和进程切换的开销，而且每个分组数据量过小，模型加载的固定开销占比会升高。数百个分组既能保证足够的并行度，又能让每个分组的数据量足够大，分摊模型加载成本。
• 内存负载说明：Executor不需要载入所有分组的数据，每个Task仅会将自己负责的单个分组数据载入内存处理，处理完成后释放该分组的内存。你需要确保单个分组的数据量（加上模型占用的1GB内存）不超过Executor给Python Worker分配的内存额度，避免OOM。比如模型占1GB，单个分组的数据建议控制在500MB以内，这样单任务内存总占用约1.5GB，留足冗余空间即可。

3. Executor数量、分区数及内存设置建议 #

• Executor数量：根据集群总CPU核数来定，一般每个Executor分配2-4核（避免单Executor核数过多导致资源竞争）。比如集群有40个可用核，可设置8-10个Executor（每个4核），预留少量核给Driver和集群其他进程。
• Executor内存：
  • 配置spark.executor.memory：建议设为3-4GB，其中一部分给JVM本身，另一部分分配给Python Worker。
  • 配置spark.executor.pyspark.memory：建议设为2-3GB，确保能容纳1GB的模型加上单个分组的数据（比如500MB），同时留足内存冗余。
• 分区数：设置为Executor总核数的2-3倍，比如8个Executor×4核=32核，分区数设64-96。虚拟分组列的数量可以和分区数保持一致，或者略多（比如100个），让Spark能灵活调度分组到不同Executor，避免数据倾斜。另外，要确保虚拟分组列的取值均匀分布，防止某几个分组数据量过大导致OOM或任务延迟。

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