如何将值为列表的字典转换为指定结构的PySpark DataFrame

问题背景 #

给定字典：

mydict = {
    'a': ["audi", "honda"],
    'b': ["toyota", "chevy"],
    'c': ["bmw", "ford"]
}

需要转换为以下结构的PySpark DataFrame，所有列均为字符串类型：

topic    id    brand
    a       s1    audi
    a       s2    honda
    b       s3    toyota
    b       s4    chevy
    c       s5    bmw
    c       s6    ford

列名定义为 col_names = ['topic', 'id', 'brand']。已有未包含id列的初步实现，需完善以生成目标DataFrame。

解决方案 #

方法1：完善原有循环逻辑（小数据场景） #

通过维护全局计数器生成连续ID，同时显式指定Schema确保列类型为字符串：

master_topic_list = []
master_id_list = []
master_brand_list = []
counter = 1  # 全局计数器，保证ID连续递增

for k, v in mydict.items():
    n = len(v)
    topic_list = [k] * n
    # 生成当前分组的ID列表
    id_list = [f's{counter + i}' for i in range(n)]
    counter += n
    
    master_topic_list.extend(topic_list)
    master_id_list.extend(id_list)
    master_brand_list.extend(v)

# 定义Schema确保所有列都是字符串类型
from pyspark.sql.types import StringType, StructType, StructField

schema = StructType([
    StructField('topic', StringType(), nullable=False),
    StructField('id', StringType(), nullable=False),
    StructField('brand', StringType(), nullable=False)
])

col_names = ['topic', 'id', 'brand']
df = spark.createDataFrame(zip(master_topic_list, master_id_list, master_brand_list), schema=schema)
df.show()

方法2：PySpark原生函数实现（大数据场景推荐） #

利用Spark分布式API处理，避免本地循环，适合大规模数据：

from pyspark.sql.functions import explode, concat_ws, lit, row_number
from pyspark.sql.window import Window
from pyspark.sql.types import StringType

# 第一步：将字典转为初始DataFrame
initial_df = spark.createDataFrame(mydict.items(), ['topic', 'brands'])

# 第二步：展开品牌列表
exploded_df = initial_df.select('topic', explode('brands').alias('brand'))

# 第三步：生成连续ID，按顺序拼接"s"前缀
window_spec = Window.orderBy('topic', 'brand')
final_df = exploded_df.withColumn('id', concat_ws('', lit('s'), row_number().over(window_spec)))

# 第四步：确保所有列类型为字符串并调整列顺序
final_df = final_df.select('topic', 'id', 'brand').withColumns({
    'topic': final_df['topic'].cast(StringType()),
    'id': final_df['id'].cast(StringType()),
    'brand': final_df['brand'].cast(StringType())
})

final_df.show()

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