PySpark中高效遍历DataFrame的优化方案咨询

我是Spark新手，想请教PySpark里更高效的DataFrame遍历方案。当前的操作逻辑是：筛选track_name为空的行，提取ID列转成列表；循环列表里的每个ID，查询对应的song_id，调用Spotify API获取缺失的track_name和artist_name，再更新原DataFrame的对应字段。

现有实现代码如下：

missing_track_name = df.filter(df['track_name'].isNull()).select(df['ID']).collect()

missing_list = [x[0] for x in missing_track_name]

for i in missing_list:

    track_id = df.filter(col('ID')==i).select(df.song_id).collect()

    url = 'https://api.spotify.com/v1/tracks/{0}'.format(track_id)
    request = requests.get(url, headers = header, params = {"limit" : 50})
    data = request.json()

    df = df.withColumn("track_name", when(col("ID") == i, data['name']).otherwise(col("track_name")))
    df = df.withColumn("artist_name", when(col("ID") == i, data['artists'][0]['name']).otherwise(col("artist_name")))

数据表样例如下（音乐来自朋友的Spotify账号）：

现有代码的问题 #

• 多次调用collect()把分布式数据拉到Driver端，大数据量下性能暴跌
• 循环里每次都过滤整个DataFrame获取track_id，重复计算完全没必要
• 每次withColumn都会生成新的DataFrame实例，频繁创建对象浪费资源
• 单个请求Spotify API，没有用批量接口，请求次数多、速度慢

优化后的实现方案 #

1. 一次性提取所有缺失数据的关联信息

先把需要补充的ID和对应的song_id一次性拉取到本地，避免多次查询：

# 提取所有track_name为空的行的ID和song_id，转成字典映射
missing_records = df.filter(df['track_name'].isNull()).select('ID', 'song_id').collect()
id_to_song = {row['ID']: row['song_id'] for row in missing_records}

2. 批量调用Spotify API获取补充数据

Spotify支持一次查询最多50个track_id，用ids参数批量请求，大幅减少API调用次数：

import requests

song_ids = list(id_to_song.values())
batch_size = 50
update_map = {}

# 按50个一组拆分请求
for idx in range(0, len(song_ids), batch_size):
    batch = song_ids[idx:idx+batch_size]
    url = f"https://api.spotify.com/v1/tracks?ids={','.join(batch)}"
    resp = requests.get(url, headers=header)
    tracks = resp.json().get('tracks', [])
    
    for track in tracks:
        # 通过song_id反向找到对应的原ID
        original_id = next(k for k, v in id_to_song.items() if v == track['id'])
        update_map[original_id] = {
            'track_name': track['name'],
            'artist_name': track['artists'][0]['name']
        }

3. 用Spark的Join操作批量更新DataFrame

把补充数据转成Spark DataFrame，通过左连接+coalesce函数批量替换空值，避免循环更新：

from pyspark.sql import Row
from pyspark.sql.functions import col, coalesce

# 把更新数据转成Spark DataFrame
update_rows = [Row(ID=k, new_track=v['track_name'], new_artist=v['artist_name']) for k, v in update_map.items()]
update_df = spark.createDataFrame(update_rows)

# 左连接原DataFrame，用coalesce优先取非空值
final_df = df.join(update_df, on='ID', how='left') \
    .withColumn('track_name', coalesce(col('track_name'), col('new_track'))) \
    .withColumn('artist_name', coalesce(col('artist_name'), col('new_artist'))) \
    .drop('new_track', 'new_artist')

优化核心点 #

• 减少collect()调用次数，只拉取必要的小量数据到本地
• 利用API批量能力，把N次请求压缩成N/50次，提升接口调用效率
• 用Spark分布式的Join操作替代循环更新，充分发挥集群计算优势
• 避免循环中重复过滤DataFrame，消除不必要的计算开销

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